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 ---
MODEL_REPO = "youkii-xr/hieroglyphic-detection"
MODEL_FILENAME = "best.pt"
os.environ["YOLO_CONFIG_DIR"] = "/tmp/Ultralytics"

try:
    print("System: Downloading model weights...")
    model_path = hf_hub_download(
        repo_id=MODEL_REPO,
        filename=MODEL_FILENAME,
        token=os.environ.get("HF_TOKEN")
    )
    model = YOLO(model_path)
    print("System: Model loaded successfully.")
except Exception as e:
    print(f"Error: {e}")
    model = None

# --- 2. LOGIC: DETECTION & REPORTING ---
def generate_human_report(detections, counts):
    """Generates a natural language summary of the findings."""
    if not detections:
        return "The Eye of Horus sees no intelligible glyphs in this image."
    
    total = len(detections)
    # Sort counts by frequency
    sorted_counts = sorted(counts.items(), key=lambda item: item[1], reverse=True)
    
    report = f"🔎 ANALYSIS COMPLETE\n"
    report += f"-----------------------------------\n"
    report += f"Total Glyphs Detected: {total}\n\n"
    report += "📝 SYMBOL INVENTORY:\n"
    
    for code, count in sorted_counts:
        report += f"• Gardiner Code '{code}': {count} instance(s)\n"
        
    report += f"\n-----------------------------------\n"
    report += f"CONFIDENCE ASSESSMENT: High\n"
    report += f"TRANSLATION STATUS: Ready for context analysis."
    return report

def detect_hieroglyphs(image: Image.Image, conf_threshold: float = 0.25):
    if image is None: 
        return None, None, "Please provide an image input."
    if model is None: 
        return None, {"error": "Model failed"}, "System Error: Model not loaded."

    try:
        results = model.predict(source=image, conf=conf_threshold, iou=0.45, imgsz=640, verbose=False, device='cpu', max_det=300)
        
        # 1. Visual
        annotated_array = results[0].plot() 
        annotated_image = Image.fromarray(annotated_array[..., ::-1])
        
        # 2. Data
        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)})

        summary_json = {
            "status": "success",
            "total_found": len(detections),
            "counts": gardiner_counts
        }
        
        # 3. Text Report
        text_report = generate_human_report(detections, gardiner_counts)
        
        return annotated_image, summary_json, text_report

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

# --- 3. UI STYLING (EGYPTIAN NIGHT) ---
custom_css = """
@import url('https://fonts.googleapis.com/css2?family=Cinzel:wght@400;700&display=swap');
@import url('https://fonts.googleapis.com/css2?family=Courier+Prime&display=swap');

:root {
    --body-background-fill: #050510 !important;
    --background-fill-primary: #0a0f1e !important;
    --border-color-primary: #d4af37 !important; 
    --text-body: #e0e7ff !important;
    --gold-glow: 0 0 15px rgba(212, 175, 55, 0.3);
}

body, .gradio-container {
    background: radial-gradient(circle at 50% 0%, #1a1f35 0%, #050510 100%) !important;
    font-family: 'Cinzel', serif !important; 
}

/* CARDS */
.card {
    background: rgba(10, 15, 30, 0.7) !important;
    border: 1px solid rgba(212, 175, 55, 0.3) !important;
    border-radius: 15px;
    padding: 20px;
    box-shadow: 0 8px 32px rgba(0, 0, 0, 0.4);
    backdrop-filter: blur(10px);
    margin-bottom: 20px;
}
.card-title {
    font-size: 16px;
    font-weight: 700;
    color: #d4af37; 
    text-transform: uppercase;
    letter-spacing: 2px;
    border-bottom: 1px solid rgba(212, 175, 55, 0.2);
    padding-bottom: 8px;
    margin-bottom: 15px;
}

/* BUTTONS */
button.primary-btn {
    background: linear-gradient(135deg, #b8860b 0%, #d4af37 100%) !important;
    border: 1px solid #ffd700 !important;
    color: #000 !important;
    font-weight: bold !important;
    font-family: 'Cinzel', serif !important;
    text-transform: uppercase;
    letter-spacing: 1px;
    box-shadow: var(--gold-glow);
}

/* REPORT TEXTBOX */
.report-box textarea {
    background-color: #1a1a25 !important;
    border: 1px solid #d4af37 !important;
    font-family: 'Courier Prime', monospace !important;
    color: #d4af37 !important;
    font-size: 14px !important;
}

/* JSON CONFIG BOX */
.config-box {
    background-color: #111 !important;
    border: 1px solid #444 !important;
    font-family: 'Courier Prime', monospace !important;
    padding: 15px;
    border-radius: 8px;
    color: #0f0;
    font-size: 12px;
    overflow-x: auto;
}

/* CLEANUP */
.gradio-image, .gradio-json { background: transparent !important; border: none !important; }
"""

# HTML COMPONENTS
header_html = """
<div style="display: flex; align-items: center; gap: 20px; padding: 20px 0;">
    <svg width="50" height="50" viewBox="0 0 100 100" fill="none">
         <path d="M10,50 Q50,10 90,50 Q50,90 10,50" stroke="#d4af37" stroke-width="3" fill="none"/>
         <circle cx="50" cy="50" r="15" fill="#d4af37"/>
         <path d="M50,65 L50,90 L30,90" stroke="#d4af37" stroke-width="3" fill="none"/>
    </svg>
    <div>
        <h1 style="margin: 0; font-size: 32px; color: #d4af37;">HORUS VISION</h1>
        <p style="margin: 0; color: #a5b4fc; font-size: 12px; letter-spacing: 2px;">HIEROGLYPHIC TRANSLATION & PRESERVATION</p>
    </div>
</div>
"""

mission_html = """
<div class="card">
    <div class="card-title">📜 THE MISSION</div>
    <p style="color: #ccc; font-size: 14px; line-height: 1.6;">
        <b>To preserve the past is to save the future.</b><br>
        This tool is part of a larger initiative to digitize and translate Ancient Egyptian inscriptions. 
        Using YOLOv8, we identify Gardiner codes instantly, bridging the gap between stone artifacts and modern understanding.
    </p>
</div>
"""

claude_guide_html = """
<div class="card" style="border-color: #ff6b6b;">
    <div class="card-title" style="color: #ff6b6b;">🤖 CLAUDE DESKTOP INTEGRATION GUIDE</div>
    
    <div style="background: rgba(255, 107, 107, 0.1); border-left: 4px solid #ff6b6b; padding: 10px; margin-bottom: 15px;">
        <strong style="color: #ff6b6b;">⚠️ IMPORTANT:</strong> 
        Do <b>NOT</b> drag and drop images directly into Claude Desktop. The local server cannot see them. 
        You must place images in the dedicated folder created below.
    </div>

    <div style="color: #ddd; font-size: 14px; line-height: 1.8;">
        <b>Step 1:</b> Create a folder on your computer: <code style="background:#333; padding:2px;">C:\\Claude_Work</code><br>
        <b>Step 2:</b> Place your hieroglyph images inside that folder.<br>
        <b>Step 3:</b> Open your Claude Desktop Config file (usually <code>%APPDATA%\\Claude\\claude_desktop_config.json</code>).<br>
        <b>Step 4:</b> Paste the configuration below inside the <code>"mcpServers"</code> block.<br>
        <b>Step 5:</b> Restart Claude Desktop.
    </div>
    <br>
    <div class="config-box">
{
  "mcpServers": {
    "gradio": {
      "command": "npx",
      "args": [
        "mcp-remote",
        "https://youkii-xr-hieroglyph-mcp-server.hf.space/gradio_api/mcp/",
        "--transport",
        "streamable-http"
      ]
    },
    "upload_helper": {
      "command": "C:\\\\Python313\\\\python.exe",
      "args": [
        "-m",
        "gradio",
        "upload-mcp",
        "https://youkii-xr-hieroglyph-mcp-server.hf.space/",
        "C:\\\\Claude_Work"
      ]
    }
  }
}
    </div>
    <p style="color: #888; font-size: 12px; margin-top: 10px;">
        *Note: Check that 'C:\\Python313\\python.exe' matches your actual Python installation path.
    </p>
</div>
"""

# --- 4. MAIN APPLICATION ---
with gr.Blocks(title="Horus Vision") as demo:
    gr.HTML(f"<style>{custom_css}</style>")
    gr.HTML(header_html)
    gr.HTML(mission_html)
    
    # --- WORKSPACE AREA ---
    with gr.Row():
        
        # LEFT COLUMN: INPUTS (Now with Tabs)
        with gr.Column(scale=1):
            gr.HTML('<div class="card"><div class="card-title">SOURCE ARTIFACT</div>')
            
            with gr.Tabs():
                # TAB 1: UPLOAD
                with gr.TabItem("📜 Upload Scroll"):
                    img_upload = gr.Image(
                        type="pil", 
                        sources=["upload", "clipboard"], 
                        label="Upload Image",
                        height=300
                    )
                    slider_upload = gr.Slider(0.1, 1.0, 0.25, label="Confidence")
                    btn_upload = gr.Button("🔮 DECIPHER SCROLL", elem_classes="primary-btn")
                
                # TAB 2: CAMERA
                with gr.TabItem("👁️ Divine Sight"):
                    img_cam = gr.Image(
                        type="pil", 
                        sources=["webcam"], 
                        label="Camera Input",
                        height=300
                    )
                    slider_cam = gr.Slider(0.1, 1.0, 0.25, label="Confidence")
                    btn_cam = gr.Button("🔮 DECIPHER VISION", elem_classes="primary-btn")
            
            gr.HTML('</div>')

        # RIGHT COLUMN: OUTPUTS
        with gr.Column(scale=1):
            gr.HTML('<div class="card"><div class="card-title">INTERPRETATION</div>')
            
            # The Annotated Image
            out_image = gr.Image(label="Annotated Artifact", interactive=False)
            
            # The Human Readable Summary
            out_report = gr.Textbox(
                label="Scribe's Report", 
                lines=6, 
                elem_classes="report-box",
                placeholder="Analysis results will appear here..."
            )
            
            # The Raw JSON (Collapsed by default ideally, but standard here)
            with gr.Accordion("Raw Glyph Data (JSON)", open=False):
                out_json = gr.JSON(label="JSON Data")
            
            gr.HTML('</div>')

    # --- CLAUDE GUIDE FOOTER ---
    gr.HTML(claude_guide_html)

    # --- WIRING ---
    # We bind both buttons to the same function, targeting the same outputs
    btn_upload.click(
        fn=detect_hieroglyphs,
        inputs=[img_upload, slider_upload],
        outputs=[out_image, out_json, out_report]
    )
    
    btn_cam.click(
        fn=detect_hieroglyphs,
        inputs=[img_cam, slider_cam],
        outputs=[out_image, out_json, out_report]
    )

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