import gradio as gr
from ultralytics import YOLO
from PIL import Image
import os
from huggingface_hub import hf_hub_download
import numpy as np

# --- 1. SETUP & MODEL LOADING (UNCHANGED) ---
MODEL_REPO = "youkii-xr/hieroglyphic-detection"
MODEL_FILENAME = "best.pt"

print(f"Server Status: Public MCP Endpoint Active")
print(f"Security: Model weights are protected (private repo)")

try:
    model_path = hf_hub_download(
        repo_id=MODEL_REPO,
        filename=MODEL_FILENAME,
        token=os.environ.get("HF_TOKEN")
    )
    print(f"System: Model loaded successfully from private storage.")
    model = YOLO(model_path)
except Exception as e:
    print(f"CRITICAL ERROR: Could not load model. Check HF_TOKEN in Settings. {e}")
    model = None

# --- 2. DETECTION LOGIC (UNCHANGED) ---
def detect_hieroglyphs(image: Image.Image, conf_threshold: float = 0.25):
    """
    Analyzes an image to find Egyptian hieroglyphs.
    """
    if image is None:
        return None, {"error": "No image provided"}
    
    if model is None:
        return None, {"error": "Server Error: Model not loaded."}

    try:
        # Run Inference
        results = model.predict(
            source=image,
            conf=conf_threshold,
            iou=0.45,
            imgsz=640,
            verbose=False,
            device='cpu',
            max_det=300
        )
        
        # 1. Generate Visual Output (RGB Image)
        annotated_array = results[0].plot() 
        annotated_image = Image.fromarray(annotated_array[..., ::-1])
        
        # 2. Generate Data Output (JSON)
        detections = []
        gardiner_counts = {}
        
        for box in results[0].boxes:
            if box.cls.numel() > 0:
                cls_id = int(box.cls[0])
                if 0 <= cls_id < len(model.names):
                    code = model.names[cls_id]
                    conf = float(box.conf[0])
                    
                    if code not in gardiner_counts: gardiner_counts[code] = 0
                    gardiner_counts[code] += 1
                    
                    detections.append({
                        "code": code,
                        "confidence": round(conf, 2),
                        "box": [round(x, 1) for x in box.xyxy[0].cpu().numpy().tolist()]
                    })

        summary = {
            "status": "success",
            "total_found": len(detections),
            "unique_symbols": list(gardiner_counts.keys()),
            "counts": gardiner_counts
        }
        
        return annotated_image, summary

    except Exception as e:
        print(f"Inference Error: {e}")
        return None, {"error": str(e)}

# --- 3. UI/UX CONFIGURATION ---

# A. Custom CSS for Animations and Fonts
# Imports 'Cinzel' font for that ancient feel and defines a pulsing gold button
custom_css = """
@import url('https://fonts.googleapis.com/css2?family=Cinzel:wght@400;700&display=swap');

body, .gradio-container {
    background-color: #fdf6e3; /* Light Papyrus */
}

h1, h2, h3 {
    font-family: 'Cinzel', serif !important;
    color: #8b4513 !important; /* SaddleBrown */
}

/* The Magic Button Animation */
@keyframes goldenPulse {
    0% { box-shadow: 0 0 0 0 rgba(212, 175, 55, 0.7); transform: scale(1); }
    50% { box-shadow: 0 0 0 10px rgba(212, 175, 55, 0); transform: scale(1.02); }
    100% { box-shadow: 0 0 0 0 rgba(212, 175, 55, 0); transform: scale(1); }
}

#magic-btn {
    background: linear-gradient(135deg, #b8860b 0%, #d4af37 100%); /* Gold Gradient */
    border: 1px solid #8b4513;
    color: white;
    font-family: 'Cinzel', serif;
    font-weight: bold;
    font-size: 1.2em;
    animation: goldenPulse 2s infinite;
    transition: all 0.3s ease;
}

#magic-btn:hover {
    animation: none;
    transform: translateY(-2px);
    box-shadow: 0 5px 15px rgba(139, 69, 19, 0.4);
}

.json-output {
    background-color: #fff8dc; /* Cornsilk */
    border: 1px solid #d4af37;
}
"""

# B. Custom Theme (Sand, Gold, Lapis)
theme = gr.themes.Soft(
    primary_hue="amber",
    secondary_hue="slate",
    neutral_hue="stone",
).set(
    body_background_fill="#fdf6e3",
    block_background_fill="#ffffff",
    block_border_color="#d4af37", # Gold borders
    button_primary_background_fill="#d4af37",
    button_primary_text_color="white",
)

# C. Claude Configuration JSON Generator
claude_config_content = """
{
  "mcpServers": {
    "hieroglyph-detector": {
      "command": "uv",
      "args": [
        "python",
        "client.py" 
      ],
      "env": {
        "GRADIO_SERVER_URL": "YOUR_SPACE_URL_HERE"
      }
    }
  }
}
"""

# --- 4. BUILD THE APP WITH BLOCKS ---
with gr.Blocks(theme=theme, css=custom_css, title="Horus Vision") as demo:
    
    # Header Area
    with gr.Row():
        with gr.Column(scale=1):
            gr.Markdown("""
            # 👁️ Horus Vision
            ### AI Hieroglyphic Decoder & Classifier
            """)
        with gr.Column(scale=3):
            gr.Markdown("""
            > *"The eye sees all."* Upload an image of Egyptian text. 
            > This AI identifies Gardiner codes using the YOLO architecture.
            """)

    # Main Tabs
    with gr.Tabs():
        
        # TAB 1: The Detector
        with gr.TabItem("🔍 Decoder"):
            with gr.Row():
                # Left Column: Inputs
                with gr.Column():
                    img_input = gr.Image(type="pil", label="Upload Papyrus/Image", sources=["upload", "clipboard"])
                    conf_slider = gr.Slider(minimum=0.1, maximum=1.0, value=0.25, step=0.05, label="Confidence Threshold")
                    
                    # The Animated Button
                    analyze_btn = gr.Button("🔮 Decipher Symbols", elem_id="magic-btn", variant="primary")
                    
                # Right Column: Outputs
                with gr.Column():
                    img_output = gr.Image(label="Annotated Result", interactive=False)
                    json_output = gr.JSON(label="Glyph Data", elem_classes="json-output")
            
            # Event Listener
            analyze_btn.click(
                fn=detect_hieroglyphs,
                inputs=[img_input, conf_slider],
                outputs=[img_output, json_output]
            )

        # TAB 2: Claude Desktop Config
        with gr.TabItem("🤖 Connect to Claude"):
            gr.Markdown("### MCP Server Configuration")
            gr.Markdown("Copy the JSON below into your Claude Desktop config file to use this model directly within Claude.")
            
            gr.Code(
                value=claude_config_content, 
                language="json", 
                label="claude_desktop_config.json",
                interactive=False
            )

    # Footer
    gr.Markdown("---")
    gr.Markdown(f"*Powered by YOLOv8 | Model: {MODEL_REPO} | Private Weights Active*")

# --- 5. LAUNCH ---
if __name__ == "__main__":
    demo.launch(
        mcp_server=True, 
        ssr_mode=False, 
        allowed_paths=["/tmp"]
    )