#!/usr/bin/env python3
"""
scripts/hf_mcp_server.py — Local Model Context Protocol (MCP) Server for Hugging Face Hub.

Provides direct tool integration for:
- HF User and PRO subscription inspection
- Spaces status, hardware management, restarts, and secrets
- Model / Dataset discovery, file uploads, and model card management
- Local ZeroGPU hub interaction

Runs as a standard stdio JSON-RPC MCP server.
"""

import sys
import os
import json
import warnings
import traceback

warnings.filterwarnings("ignore")
from typing import Any, Dict, List

try:
    from huggingface_hub import HfApi, SpaceHardware, SpaceStage
except ImportError:
    HfApi = None


def get_hf_api() -> HfApi:
    token = os.environ.get("HF_TOKEN")
    if not token:
        raise ValueError("HF_TOKEN environment variable is not set.")
    return HfApi(token=token)


# ─── Tool Implementations ───────────────────────────────────────────────────────

def tool_hf_whoami(args: Dict[str, Any]) -> str:
    api = get_hf_api()
    user_info = api.whoami()
    summary = {
        "username": user_info.get("name"),
        "fullname": user_info.get("fullname"),
        "email": user_info.get("email"),
        "is_pro": user_info.get("isPro", False),
        "can_pay": user_info.get("canPay", False),
        "orgs": [org.get("name") for org in user_info.get("orgs", [])],
        "auth_type": user_info.get("auth", {}).get("type"),
    }
    return json.dumps(summary, indent=2)


def tool_hf_space_info(args: Dict[str, Any]) -> str:
    api = get_hf_api()
    space_id = args.get("space_id", "abalanescu/flow")
    info = api.get_space_runtime(repo_id=space_id)
    out = {
        "space_id": space_id,
        "stage": str(info.stage),
        "hardware": str(info.hardware),
        "requested_hardware": str(info.requested_hardware),
        "gc_timeout": info.gc_timeout,
        "raw": {
            "current_hardware": str(info.hardware),
            "stage": str(info.stage),
        }
    }
    return json.dumps(out, indent=2)


def tool_hf_space_restart(args: Dict[str, Any]) -> str:
    api = get_hf_api()
    space_id = args.get("space_id", "abalanescu/flow")
    factory_reboot = bool(args.get("factory_reboot", False))
    res = api.restart_space(repo_id=space_id, factory_reboot=factory_reboot)
    return json.dumps({"status": "restarting", "space_id": space_id, "factory_reboot": factory_reboot, "response": str(res)}, indent=2)


def tool_hf_list_user_repos(args: Dict[str, Any]) -> str:
    api = get_hf_api()
    user_info = api.whoami()
    username = user_info.get("name")
    repo_type = args.get("repo_type", "all")  # model, dataset, space, all
    
    results = {}
    if repo_type in ("all", "model"):
        models = list(api.list_models(author=username, limit=30))
        results["models"] = [{"id": m.id, "downloads": m.downloads, "likes": m.likes, "private": m.private} for m in models]
    if repo_type in ("all", "space"):
        spaces = list(api.list_spaces(author=username, limit=30))
        results["spaces"] = [{"id": s.id, "likes": s.likes, "private": s.private} for s in spaces]
    if repo_type in ("all", "dataset"):
        datasets = list(api.list_datasets(author=username, limit=30))
        results["datasets"] = [{"id": d.id, "downloads": d.downloads, "likes": d.likes, "private": d.private} for d in datasets]
        
    return json.dumps(results, indent=2)


def tool_hf_search_models(args: Dict[str, Any]) -> str:
    api = get_hf_api()
    query = args.get("query", "")
    limit = int(args.get("limit", 10))
    filter_tag = args.get("filter")
    models = list(api.list_models(search=query, filter=filter_tag, limit=limit, sort="downloads", direction=-1))
    out = [
        {
            "id": m.id,
            "downloads": m.downloads,
            "likes": m.likes,
            "pipeline_tag": getattr(m, "pipeline_tag", None),
            "private": m.private,
        }
        for m in models
    ]
    return json.dumps(out, indent=2)


def tool_hf_upload_file(args: Dict[str, Any]) -> str:
    api = get_hf_api()
    path_or_fileobj = args.get("local_path")
    path_in_repo = args.get("path_in_repo")
    repo_id = args.get("repo_id")
    repo_type = args.get("repo_type", "space")  # model, dataset, space
    commit_message = args.get("commit_message", f"Upload {path_in_repo}")

    if not path_or_fileobj or not path_in_repo or not repo_id:
        return json.dumps({"error": "Missing required arguments: local_path, path_in_repo, repo_id"})

    if not os.path.exists(path_or_fileobj):
        return json.dumps({"error": f"Local file not found: {path_or_fileobj}"})

    url = api.upload_file(
        path_or_fileobj=path_or_fileobj,
        path_in_repo=path_in_repo,
        repo_id=repo_id,
        repo_type=repo_type,
        commit_message=commit_message,
    )
    return json.dumps({"status": "uploaded", "repo_id": repo_id, "url": str(url)}, indent=2)


# ─── MCP Tools Catalog ──────────────────────────────────────────────────────────

TOOLS_METADATA = [
    {
        "name": "hf_whoami",
        "description": "Inspect authenticated Hugging Face account profile, PRO subscription status, and organizations.",
        "inputSchema": {
            "type": "object",
            "properties": {},
        },
    },
    {
        "name": "hf_space_info",
        "description": "Get current stage, hardware tier (ZeroGPU/A10G/CPU), and runtime state for an HF Space.",
        "inputSchema": {
            "type": "object",
            "properties": {
                "space_id": {"type": "string", "description": "Space repo ID, e.g. abalanescu/flow"}
            },
        },
    },
    {
        "name": "hf_space_restart",
        "description": "Restart or factory-reboot a Hugging Face Space.",
        "inputSchema": {
            "type": "object",
            "properties": {
                "space_id": {"type": "string", "description": "Space repo ID, e.g. abalanescu/flow"},
                "factory_reboot": {"type": "boolean", "description": "Whether to perform a clean factory rebuild"}
            },
        },
    },
    {
        "name": "hf_list_user_repos",
        "description": "List all models, spaces, and datasets owned by the authenticated HF user.",
        "inputSchema": {
            "type": "object",
            "properties": {
                "repo_type": {"type": "string", "enum": ["all", "model", "space", "dataset"], "description": "Type of repos to list"}
            },
        },
    },
    {
        "name": "hf_search_models",
        "description": "Search public and private models on Hugging Face Hub by keyword or pipeline tag.",
        "inputSchema": {
            "type": "object",
            "properties": {
                "query": {"type": "string", "description": "Search keyword (e.g. qwen, gemma, flux)"},
                "limit": {"type": "integer", "description": "Max models to return (default: 10)"},
                "filter": {"type": "string", "description": "Filter by pipeline tag, e.g. text-generation"}
            },
            "required": ["query"],
        },
    },
    {
        "name": "hf_upload_file",
        "description": "Upload a file from local filesystem to a Hugging Face model, space, or dataset repo.",
        "inputSchema": {
            "type": "object",
            "properties": {
                "local_path": {"type": "string", "description": "Absolute path to local file"},
                "path_in_repo": {"type": "string", "description": "Target filename/path in repo"},
                "repo_id": {"type": "string", "description": "Target repository ID (e.g. abalanescu/flow)"},
                "repo_type": {"type": "string", "enum": ["space", "model", "dataset"], "description": "Repository type"},
                "commit_message": {"type": "string", "description": "Git commit message"}
            },
            "required": ["local_path", "path_in_repo", "repo_id"],
        },
    },
]

TOOL_DISPATCH = {
    "hf_whoami": tool_hf_whoami,
    "hf_space_info": tool_hf_space_info,
    "hf_space_restart": tool_hf_space_restart,
    "hf_list_user_repos": tool_hf_list_user_repos,
    "hf_search_models": tool_hf_search_models,
    "hf_upload_file": tool_hf_upload_file,
}


def handle_request(req: Dict[str, Any]) -> Dict[str, Any]:
    req_id = req.get("id")
    method = req.get("method")
    params = req.get("params", {})

    if method == "initialize":
        return {
            "jsonrpc": "2.0",
            "id": req_id,
            "result": {
                "protocolVersion": "2024-11-05",
                "capabilities": {
                    "tools": {}
                },
                "serverInfo": {
                    "name": "hf-local-mcp-server",
                    "version": "1.0.0"
                }
            }
        }

    elif method == "notifications/initialized":
        return None

    elif method == "tools/list":
        return {
            "jsonrpc": "2.0",
            "id": req_id,
            "result": {
                "tools": TOOLS_METADATA
            }
        }

    elif method == "tools/call":
        tool_name = params.get("name")
        arguments = params.get("arguments", {})
        if tool_name not in TOOL_DISPATCH:
            return {
                "jsonrpc": "2.0",
                "id": req_id,
                "error": {
                    "code": -32601,
                    "message": f"Unknown tool: {tool_name}"
                }
            }

        try:
            fn = TOOL_DISPATCH[tool_name]
            result_text = fn(arguments)
            return {
                "jsonrpc": "2.0",
                "id": req_id,
                "result": {
                    "content": [
                        {
                            "type": "text",
                            "text": result_text
                        }
                    ]
                }
            }
        except Exception as e:
            return {
                "jsonrpc": "2.0",
                "id": req_id,
                "error": {
                    "code": -32000,
                    "message": f"Tool execution failed: {str(e)}",
                    "data": traceback.format_exc()
                }
            }

    else:
        return {
            "jsonrpc": "2.0",
            "id": req_id,
            "error": {
                "code": -32601,
                "message": f"Method not found: {method}"
            }
        }


def main():
    for line in sys.stdin:
        line = line.strip()
        if not line:
            continue
        try:
            req = json.loads(line)
            res = handle_request(req)
            if res is not None:
                sys.stdout.write(json.dumps(res) + "\n")
                sys.stdout.flush()
        except Exception as e:
            err_res = {
                "jsonrpc": "2.0",
                "id": None,
                "error": {
                    "code": -32700,
                    "message": f"Parse error: {str(e)}"
                }
            }
            sys.stdout.write(json.dumps(err_res) + "\n")
            sys.stdout.flush()


if __name__ == "__main__":
    main()