import gradio as gr
import json
import struct
import os
import tempfile
from huggingface_hub import hf_hub_download, list_repo_files

def inspect_gguf(repo_id: str, filename: str = "model.gguf") -> str:
    """Inspect a GGUF model file on HuggingFace — get metadata, quant type, and size.
    
    Use this tool when a user asks about a GGUF model's properties, quantization level,
    or whether a model will fit on their device. Reads GGUF metadata from the Hub.
    
    Args:
        repo_id: HuggingFace repo ID (e.g., "dispatchAI/SmolLM2-135M-Instruct-mobile")
        filename: GGUF filename in the repo (default: "model.gguf")
    
    Returns:
        JSON string with GGUF metadata, quant type, and size info
    """
    try:
        # List files to find GGUF
        files = list_repo_files(repo_id)
        gguf_files = [f for f in files if f.endswith(".gguf")]
        
        if not gguf_files:
            return json.dumps({"error": f"No GGUF files found in {repo_id}", "all_files": files[:20]})
        
        if filename not in gguf_files:
            filename = gguf_files[0]
        
        # Download just the header (first 1MB) to read metadata
        # We use hf_hub_download with a partial download
        gguf_path = hf_hub_download(repo_id, filename)
        file_size = os.path.getsize(gguf_path)
        
        # Read GGUF header
        with open(gguf_path, "rb") as f:
            magic = f.read(4)
            if magic != b"GGUF":
                return json.dumps({"error": "Not a valid GGUF file", "magic": magic.hex()})
            
            version = struct.unpack("<I", f.read(4))[0]
            tensor_count = struct.unpack("<Q", f.read(8))[0]
            metadata_count = struct.unpack("<Q", f.read(8))[0]
            
            metadata = {}
            for _ in range(min(metadata_count, 100)):  # Limit to prevent huge reads
                try:
                    key_len = struct.unpack("<Q", f.read(8))[0]
                    key = f.read(key_len).decode("utf-8", errors="replace")
                    value_type = struct.unpack("<I", f.read(4))[0]
                    
                    if value_type == 8:  # STRING
                        val_len = struct.unpack("<Q", f.read(8))[0]
                        val = f.read(val_len).decode("utf-8", errors="replace")
                        metadata[key] = val
                    elif value_type == 4:  # UINT32
                        metadata[key] = struct.unpack("<I", f.read(4))[0]
                    elif value_type == 5:  # INT32
                        metadata[key] = struct.unpack("<i", f.read(4))[0]
                    elif value_type == 10:  # UINT64
                        metadata[key] = struct.unpack("<Q", f.read(8))[0]
                    elif value_type == 6:  # FLOAT32
                        metadata[key] = struct.unpack("<f", f.read(4))[0]
                    elif value_type == 7:  # BOOL
                        metadata[key] = struct.unpack("<?", f.read(1))[0]
                    elif value_type == 2:  # UINT8
                        metadata[key] = struct.unpack("<B", f.read(1))[0]
                    else:
                        metadata[key] = f"<type_{value_type}>"
                        break
                except Exception:
                    break
            
            # Extract key info
            arch = metadata.get("general.architecture", "unknown")
            name = metadata.get("general.name", "unknown")
            quant_version = metadata.get("general.quantization_version", "unknown")
            
            # Determine quant type from tensor info
            # Look for quantization in metadata
            quant_type = "unknown"
            for k, v in metadata.items():
                if "quant" in k.lower() and isinstance(v, (int, float)):
                    quant_type = str(v)
            
            # Clean up
            os.remove(gguf_path)
        
        size_mb = file_size / (1024 * 1024)
        
        return json.dumps({
            "repo_id": repo_id,
            "filename": filename,
            "file_size_mb": round(size_mb, 1),
            "file_size_gb": round(size_mb / 1024, 3),
            "gguf_version": version,
            "tensor_count": tensor_count,
            "architecture": arch,
            "model_name": name,
            "quantization_version": quant_version,
            "metadata": {k: v for k, v in metadata.items() if not isinstance(v, (bytes,))},
            "all_gguf_files": gguf_files,
            "fits_in_2gb_ram": size_mb < 2048,
            "fits_in_4gb_ram": size_mb < 4096,
            "fits_in_8gb_ram": size_mb < 8192,
            "huggingface_url": f"https://huggingface.co/{repo_id}",
        }, indent=2)
    
    except Exception as e:
        return json.dumps({"error": str(e), "repo_id": repo_id})

def list_dispatchai_gguf() -> str:
    """List all GGUF models available in the dispatchAI organization.
    
    Returns:
        JSON string with all dispatchAI models that have GGUF files
    """
    from huggingface_hub import HfApi
    api = HfApi()
    models = list(api.list_models(author="dispatchAI"))
    
    result = []
    for m in models:
        try:
            files = list_repo_files(m.id)
            gguf_files = [f for f in files if f.endswith(".gguf")]
            if gguf_files:
                result.append({
                    "repo_id": m.id,
                    "gguf_files": gguf_files,
                    "url": f"https://huggingface.co/{m.id}",
                })
        except:
            pass
    
    return json.dumps({"total": len(result), "models": result}, indent=2)

with gr.Blocks(title="dispatchAI GGUF Inspector MCP") as demo:
    gr.Markdown("## 🔍 dispatchAI GGUF Inspector (MCP Tool)")
    
    with gr.Row():
        repo = gr.Textbox(label="Repo ID", placeholder="dispatchAI/SmolLM2-135M-Instruct-mobile", scale=3)
        fname = gr.Textbox(label="Filename", value="model.gguf", scale=2)
    
    btn = gr.Button("Inspect GGUF", variant="primary")
    out = gr.Textbox(label="GGUF Metadata (JSON)", lines=20)
    btn.click(fn=inspect_gguf, inputs=[repo, fname], outputs=out)
    
    list_btn = gr.Button("List All dispatchAI GGUF Models")
    list_out = gr.Textbox(label="All Models (JSON)", lines=15)
    list_btn.click(fn=list_dispatchai_gguf, outputs=list_out)

demo.launch(mcp_server=True)