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 ---
# Replace with your actual private repo ID
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:
    # 🔒 SECURE DOWNLOAD: 
    # This uses the 'HF_TOKEN' Secret from Space Settings to authenticate.
    # Users of the Space CANNOT see this token or the downloaded file.
    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 ---
def detect_hieroglyphs(image: Image.Image, conf_threshold: float = 0.25):
    """
    Analyzes an image to find Egyptian hieroglyphs.
    
    Args:
        image: The image to analyze.
        conf_threshold: Confidence level (0.1 to 1.0). Default is 0.25.
    
    Returns:
        A tuple containing the annotated image and a JSON summary of findings.
    """
    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. INTERFACE ---
demo = gr.Interface(
    fn=detect_hieroglyphs,
    inputs=[
        gr.Image(type="pil", label="Upload Image"),
        gr.Number(value=0.25, label="Confidence")
    ],
    outputs=[
        gr.Image(label="Annotated Result"),
        gr.JSON(label="Detection Data")
    ],
    title="Egyptian Hieroglyph MCP Server",
    description="Public MCP Endpoint for Hieroglyph Detection. (Model Weights are Private)"
)

if __name__ == "__main__":
    # Settings to ensure Public access works without 403 errors:
    # ssr_mode=False: Disables Server-Side Rendering (helps with API proxies)
    # allowed_paths: Grants permission to read temp files uploaded by MCP
    demo.launch(
        mcp_server=True, 
        ssr_mode=False, 
        allowed_paths=["/tmp"]
    )