from __future__ import annotations

import asyncio
import base64
import hashlib
import io
import json
import os
import time
from functools import partial
from typing import Any, Literal
from urllib.parse import quote

import httpx
from huggingface_hub import HfApi, hf_hub_download
from mcp.server import MCPServer
from mcp.server.auth.middleware.auth_context import get_access_token
from mcp.server.auth.provider import AccessToken, TokenVerifier
from mcp.server.auth.settings import AuthSettings
from mcp.types import ToolAnnotations
from pydantic import AnyHttpUrl

HF_ORIGIN = "https://huggingface.co"
OPENAPI_MD = f"{HF_ORIGIN}/.well-known/openapi.md"
ALL_SCOPES = [
    "openid",
    "profile",
    "email",
    "read-billing",
    "read-memberships",
    "read-repos",
    "gated-repos",
    "contribute-repos",
    "write-repos",
    "manage-repos",
    "read-collections",
    "write-collections",
    "inference-api",
    "read-endpoints",
    "write-endpoints",
    "jobs",
    "webhooks",
    "write-discussions",
]

SPACE_HOST = os.getenv("SPACE_HOST")
RESOURCE_URL = (
    f"https://{SPACE_HOST}/mcp" if SPACE_HOST else "http://127.0.0.1:7860/mcp"
)
PORT = int(os.getenv("PORT", "7860"))
MAX_RESPONSE_BYTES = int(os.getenv("MAX_RESPONSE_BYTES", str(2 * 1024 * 1024)))

READ = ToolAnnotations(readOnlyHint=True, destructiveHint=False, idempotentHint=True, openWorldHint=True)
WRITE = ToolAnnotations(readOnlyHint=False, destructiveHint=False, idempotentHint=False, openWorldHint=True)
DESTRUCTIVE = ToolAnnotations(readOnlyHint=False, destructiveHint=True, idempotentHint=False, openWorldHint=True)
IDEMPOTENT_WRITE = ToolAnnotations(readOnlyHint=False, destructiveHint=False, idempotentHint=True, openWorldHint=True)


def _jsonable(value: Any) -> Any:
    # Some huggingface_hub result objects (notably CommitInfo and RepoUrl)
    # subclass str while also carrying structured fields in __dict__. Preserve
    # those fields before falling back to scalar-string serialization.
    if value is None or type(value) in (str, int, float, bool):
        return value
    if isinstance(value, bytes):
        return {"base64": base64.b64encode(value).decode("ascii"), "bytes": len(value)}
    if isinstance(value, dict):
        return {str(k): _jsonable(v) for k, v in value.items()}
    if isinstance(value, (list, tuple, set)):
        return [_jsonable(v) for v in value]
    if hasattr(value, "model_dump"):
        return _jsonable(value.model_dump())
    if hasattr(value, "__dict__"):
        return _jsonable(vars(value))
    if isinstance(value, str):
        return str(value)
    return str(value)


def _commit_summary(value: Any) -> dict[str, Any]:
    return {
        "commit_url": getattr(value, "commit_url", str(value)),
        "oid": getattr(value, "oid", None),
        "pr_url": getattr(value, "pr_url", None),
        "pr_revision": getattr(value, "pr_revision", None),
        "pr_num": getattr(value, "pr_num", None),
    }


def _current_token() -> str:
    access = get_access_token()
    if access is None or not access.token:
        raise PermissionError("Hugging Face OAuth token is missing from this MCP request")
    return access.token


def _validate_path(path: str) -> str:
    if not path.startswith("/"):
        raise ValueError("path must start with /")
    if "://" in path or path.startswith("//"):
        raise ValueError("absolute/external URLs are forbidden; this server only calls huggingface.co")
    return path


async def _hf_http(
    method: str,
    path: str,
    *,
    params: dict[str, Any] | None = None,
    json_body: Any = None,
    form_body: dict[str, Any] | None = None,
) -> dict[str, Any]:
    path = _validate_path(path)
    token = _current_token()
    headers = {
        "Authorization": f"Bearer {token}",
        "User-Agent": "patdev-hf-full-mcp/1.0",
        "Accept": "application/json, text/plain, */*",
    }
    async with httpx.AsyncClient(timeout=90.0, follow_redirects=False) as client:
        response = await client.request(
            method.upper(),
            HF_ORIGIN + path,
            headers=headers,
            params=params,
            json=json_body,
            data=form_body,
        )
    content = response.content[:MAX_RESPONSE_BYTES]
    ctype = response.headers.get("content-type", "")
    result: dict[str, Any] = {
        "status": response.status_code,
        "content_type": ctype,
        "truncated": len(response.content) > len(content),
        "bytes": len(response.content),
    }
    if "application/json" in ctype:
        try:
            result["body"] = response.json()
        except Exception:
            result["body_text"] = content.decode("utf-8", errors="replace")
    elif ctype.startswith("text/") or "markdown" in ctype or "xml" in ctype:
        result["body_text"] = content.decode("utf-8", errors="replace")
    else:
        result["body_base64"] = base64.b64encode(content).decode("ascii")
    if response.status_code >= 400:
        result["error"] = True
    return result


class HFTokenVerifier(TokenVerifier):
    """Validate opaque HF OAuth tokens without storing them server-side."""

    def __init__(self) -> None:
        self._valid: dict[str, tuple[float, str]] = {}
        self._lock = asyncio.Lock()

    async def verify_token(self, token: str) -> AccessToken | None:
        digest = hashlib.sha256(token.encode()).hexdigest()
        now = time.monotonic()
        async with self._lock:
            cached = self._valid.get(digest)
            if cached and cached[0] > now:
                return AccessToken(token=token, client_id=cached[1], scopes=ALL_SCOPES)
        try:
            async with httpx.AsyncClient(timeout=15.0) as client:
                r = await client.get(
                    f"{HF_ORIGIN}/api/whoami-v2",
                    headers={"Authorization": f"Bearer {token}"},
                )
            if r.status_code != 200:
                return None
            data = r.json()
            subject = data.get("name") or data.get("fullname") or "hf-user"
        except Exception:
            return None
        async with self._lock:
            self._valid[digest] = (now + 60.0, str(subject))
        # HF access tokens are opaque and HF does not expose a public scope-introspection
        # endpoint. Upstream Hub APIs still enforce the actual token permissions.
        return AccessToken(token=token, client_id=f"hf:{subject}", scopes=ALL_SCOPES)


mcp = MCPServer(
    "Hugging Face Full Access",
    title="Hugging Face Full MCP",
    description=(
        "Full Hugging Face Hub MCP: repositories, Spaces, discussions/PRs, Jobs, "
        "Inference Endpoints, collections, webhooks, inference and raw Hub API access."
    ),
    instructions=(
        "Use convenience tools when available. For unsupported/new HF operations, search "
        "the official OpenAPI with hf_openapi_search, then call hf_get/hf_write/hf_delete. "
        "Never send the bearer token to any host other than huggingface.co. Treat repo deletion, "
        "Space secrets, hardware changes, endpoint changes, webhooks and merges as consequential."
    ),
    token_verifier=HFTokenVerifier(),
    auth=AuthSettings(
        issuer_url=AnyHttpUrl(HF_ORIGIN),
        resource_server_url=AnyHttpUrl(RESOURCE_URL),
        required_scopes=ALL_SCOPES,
        # HF OAuth access tokens are opaque and HF exposes no public audience/scope
        # introspection endpoint. HFTokenVerifier validates each token with whoami-v2,
        # and every upstream Hub API enforces the actual grant.
        validate_token_resource=False,
    ),
)


async def _thread(fn: Any, /, *args: Any, **kwargs: Any) -> Any:
    return await asyncio.to_thread(partial(fn, *args, **kwargs))


def _api() -> HfApi:
    return HfApi(token=_current_token())


@mcp.tool(annotations=READ)
async def hf_whoami() -> dict[str, Any]:
    """Return the authenticated Hugging Face identity and effective account information."""
    return await _hf_http("GET", "/api/whoami-v2")


@mcp.tool(annotations=READ)
async def hf_openapi_search(query: str, max_matches: int = 12, context_lines: int = 5) -> dict[str, Any]:
    """Search Hugging Face's live official OpenAPI Markdown reference for endpoints/parameters."""
    if not query.strip():
        raise ValueError("query cannot be empty")
    max_matches = max(1, min(max_matches, 30))
    context_lines = max(0, min(context_lines, 12))
    async with httpx.AsyncClient(timeout=30.0) as client:
        r = await client.get(OPENAPI_MD)
        r.raise_for_status()
        lines = r.text.splitlines()
    q = query.casefold()
    matches: list[dict[str, Any]] = []
    used: set[int] = set()
    for i, line in enumerate(lines):
        if q not in line.casefold():
            continue
        start = max(0, i - context_lines)
        end = min(len(lines), i + context_lines + 1)
        if any(j in used for j in range(start, end)):
            continue
        used.update(range(start, end))
        matches.append({
            "line": i + 1,
            "excerpt": "\n".join(f"{j+1}: {lines[j]}" for j in range(start, end)),
        })
        if len(matches) >= max_matches:
            break
    return {"query": query, "matches": matches, "reference": OPENAPI_MD}


@mcp.tool(annotations=READ)
async def hf_get(path: str, params_json: str = "{}") -> dict[str, Any]:
    """GET any Hugging Face Hub API path. The hostname is fixed to huggingface.co."""
    params = json.loads(params_json or "{}")
    if not isinstance(params, dict):
        raise ValueError("params_json must decode to an object")
    return await _hf_http("GET", path, params=params)


@mcp.tool(annotations=WRITE)
async def hf_write(
    method: Literal["POST", "PUT", "PATCH"],
    path: str,
    params_json: str = "{}",
    body_json: str | None = None,
    form_json: str | None = None,
) -> dict[str, Any]:
    """POST/PUT/PATCH any Hub API path. Search hf_openapi_search first when uncertain."""
    params = json.loads(params_json or "{}")
    body = json.loads(body_json) if body_json is not None else None
    form = json.loads(form_json) if form_json is not None else None
    if not isinstance(params, dict) or (form is not None and not isinstance(form, dict)):
        raise ValueError("params_json/form_json must decode to objects")
    if body is not None and form is not None:
        raise ValueError("use either body_json or form_json, not both")
    return await _hf_http(method, path, params=params, json_body=body, form_body=form)


@mcp.tool(annotations=DESTRUCTIVE)
async def hf_delete(path: str, params_json: str = "{}", body_json: str | None = None) -> dict[str, Any]:
    """DELETE any Hub API path. This can permanently remove resources; verify the target first."""
    params = json.loads(params_json or "{}")
    body = json.loads(body_json) if body_json is not None else None
    if not isinstance(params, dict):
        raise ValueError("params_json must decode to an object")
    return await _hf_http("DELETE", path, params=params, json_body=body)


@mcp.tool(annotations=READ)
async def repo_list_files(
    repo_id: str,
    repo_type: Literal["model", "dataset", "space"] = "model",
    revision: str = "main",
) -> list[str]:
    """List files in a model, dataset or Space repository."""
    result = await _thread(_api().list_repo_files, repo_id=repo_id, repo_type=repo_type, revision=revision)
    return list(result)


@mcp.tool(annotations=READ)
async def repo_read_text(
    repo_id: str,
    path_in_repo: str,
    repo_type: Literal["model", "dataset", "space"] = "model",
    revision: str = "main",
    max_bytes: int = 500_000,
) -> dict[str, Any]:
    """Read a UTF-8 text file from a Hub repo without writing it to server storage."""
    local_path = await _thread(
        hf_hub_download,
        repo_id=repo_id,
        filename=path_in_repo,
        repo_type=repo_type,
        revision=revision,
        token=_current_token(),
    )
    with open(local_path, "rb") as f:
        raw = f.read(max_bytes + 1)
    truncated = len(raw) > max_bytes
    raw = raw[:max_bytes]
    return {
        "repo_id": repo_id,
        "repo_type": repo_type,
        "revision": revision,
        "path": path_in_repo,
        "body_text": raw.decode("utf-8", errors="replace"),
        "bytes_read": len(raw),
        "truncated": truncated,
    }


@mcp.tool(annotations=WRITE)
async def repo_create(
    repo_id: str,
    repo_type: Literal["model", "dataset", "space"] = "model",
    private: bool = True,
    space_sdk: Literal["gradio", "streamlit", "docker", "static"] | None = None,
) -> dict[str, Any]:
    """Create a model, dataset or Space repository."""
    kwargs: dict[str, Any] = {"repo_id": repo_id, "repo_type": repo_type, "private": private, "exist_ok": False}
    if repo_type == "space" and space_sdk:
        kwargs["space_sdk"] = space_sdk
    out = await _thread(_api().create_repo, **kwargs)
    return {"repo_id": repo_id, "repo_type": repo_type, "url": str(out)}


@mcp.tool(annotations=DESTRUCTIVE)
async def repo_delete(repo_id: str, repo_type: Literal["model", "dataset", "space"] = "model") -> dict[str, Any]:
    """Permanently delete a model, dataset or Space repository."""
    out = await _thread(_api().delete_repo, repo_id=repo_id, repo_type=repo_type)
    return {"deleted": True, "repo_id": repo_id, "repo_type": repo_type, "result": _jsonable(out)}


@mcp.tool(annotations=IDEMPOTENT_WRITE)
async def repo_upload_text(
    repo_id: str,
    path_in_repo: str,
    content: str,
    repo_type: Literal["model", "dataset", "space"] = "model",
    revision: str = "main",
    commit_message: str = "Update via HF Full MCP",
    create_pr: bool = False,
) -> dict[str, Any]:
    """Upload/overwrite a UTF-8 text file. Set create_pr=true to propose instead of directly committing."""
    data = io.BytesIO(content.encode("utf-8"))
    out = await _thread(
        _api().upload_file,
        path_or_fileobj=data,
        path_in_repo=path_in_repo,
        repo_id=repo_id,
        repo_type=repo_type,
        revision=revision,
        commit_message=commit_message,
        create_pr=create_pr,
    )
    return _commit_summary(out)


@mcp.tool(annotations=IDEMPOTENT_WRITE)
async def repo_upload_base64(
    repo_id: str,
    path_in_repo: str,
    content_base64: str,
    repo_type: Literal["model", "dataset", "space"] = "model",
    revision: str = "main",
    commit_message: str = "Upload via HF Full MCP",
    create_pr: bool = False,
) -> dict[str, Any]:
    """Upload/overwrite a binary file supplied as base64."""
    raw = base64.b64decode(content_base64, validate=True)
    if len(raw) > 8 * 1024 * 1024:
        raise ValueError("inline upload limited to 8 MiB; use Jobs/Hub tooling for larger artifacts")
    out = await _thread(
        _api().upload_file,
        path_or_fileobj=io.BytesIO(raw),
        path_in_repo=path_in_repo,
        repo_id=repo_id,
        repo_type=repo_type,
        revision=revision,
        commit_message=commit_message,
        create_pr=create_pr,
    )
    return _commit_summary(out)


@mcp.tool(annotations=DESTRUCTIVE)
async def repo_delete_path(
    repo_id: str,
    path_in_repo: str,
    repo_type: Literal["model", "dataset", "space"] = "model",
    revision: str = "main",
    commit_message: str = "Delete via HF Full MCP",
    create_pr: bool = False,
    is_folder: bool = False,
) -> dict[str, Any]:
    """Delete a file or folder from a Hub repository, optionally through a PR."""
    method = _api().delete_folder if is_folder else _api().delete_file
    out = await _thread(
        method,
        path_in_repo=path_in_repo,
        repo_id=repo_id,
        repo_type=repo_type,
        revision=revision,
        commit_message=commit_message,
        create_pr=create_pr,
    )
    return _commit_summary(out)


@mcp.tool(annotations=WRITE)
async def discussion_create(
    repo_id: str,
    title: str,
    description: str | None = None,
    repo_type: Literal["model", "dataset", "space"] = "model",
    pull_request: bool = False,
) -> dict[str, Any]:
    """Create a discussion or draft Pull Request."""
    out = await _thread(
        _api().create_discussion,
        repo_id=repo_id,
        title=title,
        description=description,
        repo_type=repo_type,
        pull_request=pull_request,
    )
    return _jsonable(out)


@mcp.tool(annotations=WRITE)
async def discussion_comment(
    repo_id: str,
    discussion_num: int,
    comment: str,
    repo_type: Literal["model", "dataset", "space"] = "model",
) -> dict[str, Any]:
    """Post a comment on a Discussion or Pull Request."""
    out = await _thread(
        _api().comment_discussion,
        repo_id=repo_id,
        discussion_num=discussion_num,
        comment=comment,
        repo_type=repo_type,
    )
    return _jsonable(out)


@mcp.tool(annotations=DESTRUCTIVE)
async def pull_request_merge(
    repo_id: str,
    discussion_num: int,
    repo_type: Literal["model", "dataset", "space"] = "model",
    comment: str | None = None,
) -> dict[str, Any]:
    """Merge a Pull Request on the Hub."""
    out = await _thread(
        _api().merge_pull_request,
        repo_id=repo_id,
        discussion_num=discussion_num,
        repo_type=repo_type,
        comment=comment,
    )
    return {"merged": True, "repo_id": repo_id, "discussion_num": discussion_num, "result": _jsonable(out)}


@mcp.tool(annotations=READ)
async def space_runtime(repo_id: str) -> dict[str, Any]:
    """Return the current runtime/build/hardware state of a Space."""
    return _jsonable(await _thread(_api().get_space_runtime, repo_id=repo_id))


@mcp.tool(annotations=WRITE)
async def space_restart(repo_id: str) -> dict[str, Any]:
    """Restart a Space while preserving its repository and configured storage."""
    return _jsonable(await _thread(_api().restart_space, repo_id=repo_id))


@mcp.tool(annotations=WRITE)
async def space_pause(repo_id: str) -> dict[str, Any]:
    """Pause a Space."""
    return _jsonable(await _thread(_api().pause_space, repo_id=repo_id))


@mcp.tool(annotations=WRITE)
async def space_request_hardware(repo_id: str, hardware: str, sleep_time: int | None = None) -> dict[str, Any]:
    """Request/change Space hardware. This can affect billing; verify the hardware tier first."""
    kwargs: dict[str, Any] = {"repo_id": repo_id, "hardware": hardware}
    if sleep_time is not None:
        kwargs["sleep_time"] = sleep_time
    return _jsonable(await _thread(_api().request_space_hardware, **kwargs))


@mcp.tool(annotations=IDEMPOTENT_WRITE)
async def space_set_secret(repo_id: str, key: str, value: str, description: str | None = None) -> dict[str, Any]:
    """Add/update a Space secret. Secret values are write-only and are never returned by HF."""
    await _thread(_api().add_space_secret, repo_id=repo_id, key=key, value=value, description=description)
    return {"updated": True, "repo_id": repo_id, "key": key}


@mcp.tool(annotations=DESTRUCTIVE)
async def space_delete_secret(repo_id: str, key: str) -> dict[str, Any]:
    """Delete a Space secret."""
    await _thread(_api().delete_space_secret, repo_id=repo_id, key=key)
    return {"deleted": True, "repo_id": repo_id, "key": key}


@mcp.tool(annotations=IDEMPOTENT_WRITE)
async def space_set_variable(repo_id: str, key: str, value: str, description: str | None = None) -> dict[str, Any]:
    """Add/update a public Space environment variable."""
    await _thread(_api().add_space_variable, repo_id=repo_id, key=key, value=value, description=description)
    return {"updated": True, "repo_id": repo_id, "key": key, "value": value}


@mcp.tool(annotations=DESTRUCTIVE)
async def space_delete_variable(repo_id: str, key: str) -> dict[str, Any]:
    """Delete a Space environment variable."""
    await _thread(_api().delete_space_variable, repo_id=repo_id, key=key)
    return {"deleted": True, "repo_id": repo_id, "key": key}


if __name__ == "__main__":
    mcp.run(
        "streamable-http",
        host="0.0.0.0",
        port=PORT,
        streamable_http_path="/mcp",
        json_response=True,
        stateless_http=True,
    )