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"

# Fix for the Ultralytics config warning in containers
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. DETECTION LOGIC ---
def detect_hieroglyphs(image: Image.Image, conf_threshold: float = 0.25):
    if image is None: return None, None
    if model is None: return None, {"error": "Model failed to load"}

    try:
        results = model.predict(source=image, conf=conf_threshold, iou=0.45, imgsz=640, verbose=False, device='cpu', max_det=300)
        
        # Visual
        annotated_array = results[0].plot() 
        annotated_image = Image.fromarray(annotated_array[..., ::-1])
        
        # 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 = {
            "status": "success",
            "total_found": len(detections),
            "counts": gardiner_counts
        }
        return annotated_image, summary
    except Exception as e:
        return None, {"error": str(e)}

# --- 3. UI/UX: THE EGYPTIAN NIGHT STYLE ---

custom_css = """
@import url('https://fonts.googleapis.com/css2?family=Cinzel:wght@400;700&display=swap');

:root {
    --body-background-fill: #050510 !important;
    --background-fill-primary: #0a0f1e !important;
    --border-color-primary: #d4af37 !important; /* Gold */
    --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; 
}

/* --- THE GLASS CARDS --- */
.card {
    background: rgba(10, 15, 30, 0.6) !important;
    border: 1px solid rgba(212, 175, 55, 0.3) !important;
    border-radius: 20px;
    padding: 24px;
    box-shadow: 0 8px 32px rgba(0, 0, 0, 0.4);
    backdrop-filter: blur(10px);
}

.card-title {
    font-size: 18px;
    font-weight: 700;
    color: #d4af37; 
    text-transform: uppercase;
    letter-spacing: 2px;
    margin-bottom: 20px;
    display: flex;
    align-items: center;
    gap: 10px;
    border-bottom: 1px solid rgba(212, 175, 55, 0.1);
    padding-bottom: 10px;
}

/* --- THE TERMINAL WINDOW (Quick Start) --- */
.terminal-window {
    background: #0f0f15;
    border: 1px solid #333;
    border-radius: 12px;
    overflow: hidden;
    font-family: 'Courier New', monospace;
    box-shadow: 0 10px 30px rgba(0,0,0,0.8);
}
.terminal-header {
    background: #1a1a20;
    padding: 10px 15px;
    display: flex;
    gap: 8px;
    border-bottom: 1px solid #333;
}
.dot { width: 12px; height: 12px; border-radius: 50%; }
.red { background: #ff5f56; } .yellow { background: #ffbd2e; } .green { background: #27c93f; }

.terminal-body {
    padding: 20px;
    color: #a9b1d6;
    font-size: 13px;
    line-height: 1.6;
}
.json-key { color: #7aa2f7; } 
.json-string { color: #e0af68; } 

/* --- 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;
    transition: all 0.3s ease;
    box-shadow: var(--gold-glow);
}
button.primary-btn:hover {
    transform: translateY(-2px);
    box-shadow: 0 0 25px rgba(212, 175, 55, 0.6);
}

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

# HTML for the Header (Floating Style)
header_html = """
<div style="display: flex; align-items: center; justify-content: space-between; margin-bottom: 30px; padding: 20px;">
    <div style="display: flex; align-items: center; gap: 20px;">
        <svg width="60" height="60" 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: 36px; color: #d4af37; text-shadow: 0 0 10px rgba(212,175,55,0.5);">HORUS VISION</h1>
            <p style="margin: 0; color: #a5b4fc; font-size: 14px; letter-spacing: 2px;">HIEROGLYPHIC INTELLIGENCE SYSTEM</p>
        </div>
    </div>
    <div style="border: 1px solid #d4af37; padding: 5px 15px; border-radius: 20px; color: #d4af37; font-size: 12px;">
        ● SYSTEM ONLINE
    </div>
</div>
"""

terminal_html = """
<div class="card">
    <div class="card-title">⚡ CLAUDE DESKTOP CONFIG</div>
    <div class="terminal-window">
        <div class="terminal-header">
            <div class="dot red"></div><div class="dot yellow"></div><div class="dot green"></div>
            <div style="margin-left: auto; color: #555; font-size: 10px;">claude_config.json</div>
        </div>
        <div class="terminal-body">
            {<br>
            &nbsp;&nbsp;<span class="json-key">"mcpServers"</span>: {<br>
            &nbsp;&nbsp;&nbsp;&nbsp;<span class="json-key">"horus-vision"</span>: {<br>
            &nbsp;&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;<span class="json-key">"command"</span>: <span class="json-string">"uv"</span>,<br>
            &nbsp;&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;<span class="json-key">"args"</span>: [<span class="json-string">"python"</span>, <span class="json-string">"client.py"</span>]<br>
            &nbsp;&nbsp;&nbsp;&nbsp;}<br>
            &nbsp;&nbsp;}<br>
            }
        </div>
    </div>
</div>
"""

# --- 4. APP ASSEMBLY ---
# FIX: Removed 'css' argument here. We inject it via gr.HTML below.
with gr.Blocks(title="Horus Vision") as demo:
    
    # FIX: Inject CSS directly into the HTML. This works on ALL Gradio versions.
    gr.HTML(f"<style>{custom_css}</style>")
    
    # 1. Header Injection
    gr.HTML(header_html)
    
    # 2. Top Row: Info & Config
    with gr.Row():
        with gr.Column(scale=1):
            gr.HTML("""
            <div class="card" style="height: 100%;">
                <div class="card-title">📜 MISSION BRIEF</div>
                <p style="color: #ccc; line-height: 1.6;">
                    This system utilizes the YOLOv8 architecture to detect and classify Gardiner codes.
                    <br><br>
                    <span style="color: #d4af37;">> Model:</span> YOLOv8 Custom<br>
                    <span style="color: #d4af37;">> Security:</span> MCP Standard
                </p>
            </div>
            """)
            
        with gr.Column(scale=1):
            gr.HTML(terminal_html)

    # 3. Main Detector Area
    gr.HTML("<br>") 
    
    with gr.Row():
        with gr.Column():
            gr.HTML('<div class="card"> <div class="card-title">👁️ INPUT SOURCE</div>')
            img_input = gr.Image(type="pil", label="Upload Papyrus", sources=["upload", "clipboard", "webcam"], elem_classes="gradio-image")
            conf_slider = gr.Slider(0.1, 1.0, 0.25, step=0.05, label="Detection Confidence")
            analyze_btn = gr.Button("🔮 DECIPHER SYMBOLS", elem_classes="primary-btn")
            gr.HTML('</div>') 

        with gr.Column():
            gr.HTML('<div class="card"> <div class="card-title">💎 DECODED ARTIFACT</div>')
            img_output = gr.Image(label="Annotated Result", interactive=False, elem_classes="gradio-image")
            json_output = gr.JSON(label="Glyph Data", elem_classes="gradio-json")
            gr.HTML('</div>') 

    # 4. Wiring
    analyze_btn.click(
        fn=detect_hieroglyphs,
        inputs=[img_input, conf_slider],
        outputs=[img_output, json_output]
    )

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