import requests
import base64
from io import BytesIO
from PIL import Image
import json
import matplotlib.pyplot as plt
import numpy as np

# This is an example client that demonstrates how to use the MCP server
# You would replace this URL with the actual URL of your deployed Huggingface Space
MCP_SERVER_URL = "https://your-username-nomic-vision-mcp.hf.space/mcp"

def embed_image_from_url(image_url):
    """
    Generate embeddings for an image using the MCP server's embed_image tool
    
    Args:
        image_url: URL of the image to embed
        
    Returns:
        The embedding vector and its dimension
    """
    # Prepare the MCP request
    mcp_request = {
        "jsonrpc": "2.0",
        "method": "callTool",
        "params": {
            "name": "embed_image",
            "arguments": {
                "image_url": image_url
            }
        },
        "id": 1
    }
    
    # Send the request to the MCP server
    response = requests.post(MCP_SERVER_URL, json=mcp_request)
    
    # Parse the response
    result = response.json()
    
    if "error" in result:
        print(f"Error: {result['error']['message']}")
        return None
    
    # Extract the embedding from the response
    content = result["result"]["content"][0]["text"]
    embedding_data = json.loads(content)
    
    return embedding_data

def embed_image_from_file(image_path):
    """
    Generate embeddings for an image using the MCP server's embed_image tool
    
    Args:
        image_path: Path to the image file
        
    Returns:
        The embedding vector and its dimension
    """
    # Load the image
    with open(image_path, "rb") as f:
        image_data = f.read()
    
    # Encode the image as base64
    image_base64 = base64.b64encode(image_data).decode("utf-8")
    
    # Prepare the MCP request
    mcp_request = {
        "jsonrpc": "2.0",
        "method": "callTool",
        "params": {
            "name": "embed_image",
            "arguments": {
                "image_data": image_base64
            }
        },
        "id": 1
    }
    
    # Send the request to the MCP server
    response = requests.post(MCP_SERVER_URL, json=mcp_request)
    
    # Parse the response
    result = response.json()
    
    if "error" in result:
        print(f"Error: {result['error']['message']}")
        return None
    
    # Extract the embedding from the response
    content = result["result"]["content"][0]["text"]
    embedding_data = json.loads(content)
    
    return embedding_data

def visualize_embedding(embedding):
    """
    Visualize the embedding vector
    
    Args:
        embedding: The embedding vector
    """
    # Convert the embedding to a numpy array
    embedding_array = np.array(embedding)
    
    # Plot the embedding
    plt.figure(figsize=(10, 5))
    plt.plot(embedding_array)
    plt.title("Embedding Vector")
    plt.xlabel("Dimension")
    plt.ylabel("Value")
    plt.grid(True)
    plt.show()
    
    # Plot the histogram of the embedding
    plt.figure(figsize=(10, 5))
    plt.hist(embedding_array, bins=50)
    plt.title("Embedding Histogram")
    plt.xlabel("Value")
    plt.ylabel("Frequency")
    plt.grid(True)
    plt.show()

if __name__ == "__main__":
    # Example usage with an image URL
    image_url = "https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/transformers/model_doc/bert-architects.png"
    print(f"Generating embedding for image: {image_url}")
    
    embedding_data = embed_image_from_url(image_url)
    
    if embedding_data:
        print(f"Embedding dimension: {embedding_data['dimension']}")
        print(f"First 10 values of embedding: {embedding_data['embedding'][:10]}...")
        
        # Visualize the embedding
        visualize_embedding(embedding_data['embedding'])
    
    # Example usage with a local image file
    # Uncomment the following lines to use a local image file
    # image_path = "path/to/your/image.jpg"
    # print(f"Generating embedding for image: {image_path}")
    # embedding_data = embed_image_from_file(image_path)
    # 
    # if embedding_data:
    #     print(f"Embedding dimension: {embedding_data['dimension']}")
    #     print(f"First 10 values of embedding: {embedding_data['embedding'][:10]}...")
    #     
    #     # Visualize the embedding
    #     visualize_embedding(embedding_data['embedding'])