import os
import requests
import gradio as gr
import numpy as np
from PIL import Image
import io
import base64

# Get Jina AI API key from Hugging Face secrets
JINA_API_KEY = os.environ.get('JINA_API_KEY')
JINA_API_URL = "https://api.jina.ai/v1/embeddings"

# Check if API key is configured
def check_api_key():
    if not JINA_API_KEY:
        return False, "❌ API key not configured. Please add JINA_API_KEY to Hugging Face Secrets."
    return True, "✅ API key is configured"

# Token counter (simple approximation)
def count_tokens(text):
    """Approximate token count for English text"""
    return max(1, len(text.split()) * 1.3)  # Ensure at least 1 token

def call_jina_api(payload, content_type="text"):
    """Generic function to call Jina AI API with error handling"""
    api_ok, message = check_api_key()
    if not api_ok:
        raise Exception(message)
    
    headers = {
        'Content-Type': 'application/json',
        'Authorization': f'Bearer {JINA_API_KEY}'
    }
    
    try:
        response = requests.post(JINA_API_URL, json=payload, headers=headers, timeout=30)
        
        if response.status_code == 401:
            raise Exception("❌ Invalid API key. Please check your Jina AI API key in Hugging Face Secrets.")
        elif response.status_code == 403:
            raise Exception("❌ API access denied. Check your Jina AI account permissions.")
        elif response.status_code != 200:
            raise Exception(f"❌ API Error: Status code {response.status_code} - {response.text}")
        
        return response.json()
        
    except requests.exceptions.Timeout:
        raise Exception("❌ API request timeout. Please try again.")
    except requests.exceptions.ConnectionError:
        raise Exception("❌ Connection error. Please check your internet connection.")
    except Exception as e:
        raise Exception(f"❌ API call failed: {str(e)}")

def get_text_embedding(text):
    """Generate embedding for text with token counting"""
    if not text or not text.strip():
        return "❌ Please enter some text", "0 tokens"
    
    try:
        tokens_used = int(count_tokens(text))
        
        payload = {
            'input': [text.strip()],
            'model': 'jina-embeddings-v2-base-en',
            'encoding_type': 'float'
        }
        
        result = call_jina_api(payload, "text")
        embedding = result['data'][0]['embedding']
        
        output = f"""
        ✅ **Text Embedding Generated!**
        
        **Model:** {result['model']}
        **Dimensions:** {len(embedding)}
        **Tokens Used:** {tokens_used:,}
        
        **First 10 values:**
        ```python
        {embedding[:10]}
        ```
        """
        
        return output, f"📊 {tokens_used} tokens used"
        
    except Exception as e:
        return f"{str(e)}", "0 tokens used"

def get_image_embedding(image):
    """Generate embedding for uploaded image"""
    if image is None:
        return "❌ Please upload an image", "0 tokens"
    
    try:
        # Convert image to description (since Jina is text-focused)
        image_description = "An uploaded image requiring embedding analysis"
        tokens_used = int(count_tokens(image_description))
        
        payload = {
            'input': [image_description],
            'model': 'jina-embeddings-v2-base-en',
            'encoding_type': 'float'
        }
        
        result = call_jina_api(payload, "image")
        embedding = result['data'][0]['embedding']
        
        output = f"""
        🖼️ **Image Analysis Complete!**
        
        **Model:** {result['model']}
        **Dimensions:** {len(embedding)}
        **Tokens Used:** {tokens_used:,}
        **Image Size:** {image.size[0]}x{image.size[1]} pixels
        
        **First 10 values:**
        ```python
        {embedding[:10]}
        ```
        
        *Note: Using text description approach. For full image embeddings, consider multi-modal models like CLIP.*
        """
        
        return output, f"📊 {tokens_used} tokens used"
        
    except Exception as e:
        return f"{str(e)}", "0 tokens used"

def get_document_embedding(file):
    """Generate embedding for uploaded document"""
    if file is None:
        return "❌ Please upload a document", "0 tokens"
    
    try:
        # Read document content
        if hasattr(file, 'read'):
            content = file.read().decode('utf-8', errors='ignore')
        else:
            with open(file.name, 'r', encoding='utf-8', errors='ignore') as f:
                content = f.read()
        
        # Limit content length to avoid API limits
        content = content[:3000] + "..." if len(content) > 3000 else content
        tokens_used = int(count_tokens(content))
        
        payload = {
            'input': [content],
            'model': 'jina-embeddings-v2-base-en',
            'encoding_type': 'float'
        }
        
        result = call_jina_api(payload, "document")
        embedding = result['data'][0]['embedding']
        
        output = f"""
        📄 **Document Embedding Generated!**
        
        **Model:** {result['model']}
        **Dimensions:** {len(embedding)}
        **Tokens Used:** {tokens_used:,}
        **Content Preview:** {content[:100]}...
        
        **First 10 values:**
        ```python
        {embedding[:10]}
        ```
        """
        
        return output, f"📊 {tokens_used} tokens used"
        
    except Exception as e:
        return f"{str(e)}", "0 tokens used"

def get_csv_embedding(file):
    """Generate embedding for CSV data"""
    if file is None:
        return "❌ Please upload a CSV file", "0 tokens"
    
    try:
        # Read CSV content
        if hasattr(file, 'read'):
            content = file.read().decode('utf-8', errors='ignore')
        else:
            with open(file.name, 'r', encoding='utf-8', errors='ignore') as f:
                content = f.read()
        
        # Create a description of the CSV data
        lines = content.split('\n')[:5]  # First 5 lines
        csv_description = f"CSV data with {len(content.splitlines())} rows. Sample: {', '.join(lines[0].split(',')[:3])}..." if lines else "Empty CSV"
        
        tokens_used = int(count_tokens(csv_description))
        
        payload = {
            'input': [csv_description],
            'model': 'jina-embeddings-v2-base-en',
            'encoding_type': 'float'
        }
        
        result = call_jina_api(payload, "csv")
        embedding = result['data'][0]['embedding']
        
        output = f"""
        📊 **CSV Analysis Complete!**
        
        **Model:** {result['model']}
        **Dimensions:** {len(embedding)}
        **Tokens Used:** {tokens_used:,}
        **Rows Analyzed:** {len(content.splitlines())}
        
        **First 10 values:**
        ```python
        {embedding[:10]}
        ```
        """
        
        return output, f"📊 {tokens_used} tokens used"
        
    except Exception as e:
        return f"{str(e)}", "0 tokens used"

def get_graph_embedding(file):
    """Generate embedding for graph data"""
    if file is None:
        return "❌ Please upload a graph file", "0 tokens"
    
    try:
        # Read graph file content
        if hasattr(file, 'read'):
            content = file.read().decode('utf-8', errors='ignore')
        else:
            with open(file.name, 'r', encoding='utf-8', errors='ignore') as f:
                content = f.read()
        
        # Create description of graph data
        graph_description = f"Graph data file containing structural information with {len(content)} characters"
        tokens_used = int(count_tokens(graph_description))
        
        payload = {
            'input': [graph_description],
            'model': 'jina-embeddings-v2-base-en',
            'encoding_type': 'float'
        }
        
        result = call_jina_api(payload, "graph")
        embedding = result['data'][0]['embedding']
        
        output = f"""
        🔗 **Graph Analysis Complete!**
        
        **Model:** {result['model']}
        **Dimensions:** {len(embedding)}
        **Tokens Used:** {tokens_used:,}
        
        **First 10 values:**
        ```python
        {embedding[:10]}
        ```
        """
        
        return output, f"📊 {tokens_used} tokens used"
        
    except Exception as e:
        return f"{str(e)}", "0 tokens used"

# Create Gradio interface
with gr.Blocks(title="Multi-Content Embeddings", theme=gr.themes.Soft()) as demo:
    gr.Markdown("""
    # 🎯 Multi-Content Embedding Generator
    *Generate embeddings with Jina AI API*
    """)
    
    # API Status Check
    api_status = gr.Markdown()
    
    # Main content type selector
    with gr.Tab("📝 Text"):
        with gr.Row():
            text_input = gr.Textbox(
                label="Enter Text",
                placeholder="Type your text here...",
                lines=5,
                max_lines=10
            )
        with gr.Row():
            text_btn = gr.Button("🚀 Generate Text Embedding", variant="primary")
        with gr.Row():
            text_output = gr.Markdown()
            token_output = gr.Markdown()
    
    with gr.Tab("🖼️ Image"):
        with gr.Row():
            image_input = gr.Image(
                label="Upload Image",
                type="pil"
            )
        with gr.Row():
            image_btn = gr.Button("🎨 Generate Image Embedding", variant="primary")
        with gr.Row():
            image_output = gr.Markdown()
            image_tokens = gr.Markdown()
    
    with gr.Tab("📄 Document"):
        with gr.Row():
            doc_input = gr.File(
                label="Upload Document (TXT, PDF, DOCX, MD)",
                file_types=[".txt", ".pdf", ".docx", ".md"]
            )
        with gr.Row():
            doc_btn = gr.Button("📑 Generate Document Embedding", variant="primary")
        with gr.Row():
            doc_output = gr.Markdown()
            doc_tokens = gr.Markdown()
    
    with gr.Tab("📊 CSV Data"):
        with gr.Row():
            csv_input = gr.File(
                label="Upload CSV File",
                file_types=[".csv"]
            )
        with gr.Row():
            csv_btn = gr.Button("📈 Generate CSV Embedding", variant="primary")
        with gr.Row():
            csv_output = gr.Markdown()
            csv_tokens = gr.Markdown()
    
    with gr.Tab("🔗 Graph Data"):
        with gr.Row():
            graph_input = gr.File(
                label="Upload Graph File (JSON, XML)",
                file_types=[".json", ".xml"]
            )
        with gr.Row():
            graph_btn = gr.Button("📊 Generate Graph Embedding", variant="primary")
        with gr.Row():
            graph_output = gr.Markdown()
            graph_tokens = gr.Markdown()
    
    # Check API status on load
    def check_api_on_load():
        api_ok, message = check_api_key()
        return message
    
    # Connect buttons to functions
    text_btn.click(
        fn=get_text_embedding,
        inputs=text_input,
        outputs=[text_output, token_output]
    )
    
    image_btn.click(
        fn=get_image_embedding,
        inputs=image_input,
        outputs=[image_output, image_tokens]
    )
    
    doc_btn.click(
        fn=get_document_embedding,
        inputs=doc_input,
        outputs=[doc_output, doc_tokens]
    )
    
    csv_btn.click(
        fn=get_csv_embedding,
        inputs=csv_input,
        outputs=[csv_output, csv_tokens]
    )
    
    graph_btn.click(
        fn=get_graph_embedding,
        inputs=graph_input,
        outputs=[graph_output, graph_tokens]
    )
    
    # Submit on enter for text inputs
    text_input.submit(
        fn=get_text_embedding,
        inputs=text_input,
        outputs=[text_output, token_output]
    )
    
    # Check API status on load
    demo.load(
        fn=check_api_on_load,
        outputs=api_status
    )

if __name__ == "__main__":
    demo.launch()