"""Claude Code-style MCP server with a private workspace administration dashboard.

The browser dashboard is available at / and asks for MCP_AUTH_TOKEN.
The MCP endpoint is /gradio_api/mcp/ and continues to accept the token as
Authorization: Bearer ..., X-MCP-Token: ... or a query parameter.
"""

import asyncio
import hmac
import json
import os
import re
import time

import gradio as gr
import uvicorn
from fastapi import FastAPI, HTTPException, Query, Request
from fastapi.responses import HTMLResponse, JSONResponse, PlainTextResponse, Response, StreamingResponse
from starlette.concurrency import run_in_threadpool

import workspaces
import runtime_client
from tool_dispatch import adapter
from admin_ui import DASHBOARD_PAGE, LOCKED_PAGE, LOGIN_PAGE
from tools import ALL_TOOLS

TOKEN = os.environ.get("MCP_AUTH_TOKEN", "")
TOKEN_OK = len(TOKEN) >= 16
PORT = int(os.environ.get("PORT", 7860))
SESSION_COOKIE = "ccmcp_admin"
SESSION_TTL_SECONDS = 12 * 60 * 60
WORKSPACE_ID_RE = re.compile(r"^ws-(?:[0-9a-f]{6}|[0-9a-f]{32})$")
JOB_ID_RE = re.compile(r"^bg-?(?:[0-9a-f]{12}|[0-9]+)$")

with gr.Blocks(title="Claude Code MCP") as demo:
    # Keep MCP tool schemas registered with Gradio. queue=False avoids serializing
    # independent chats through the same Gradio event queue.
    for fn in ALL_TOOLS:
        gr.api(adapter(fn), queue=False)

app = FastAPI(title="Claude Code MCP")


def _session_signature(expires: int) -> str:
    message = f"workspace-admin:{expires}".encode("ascii")
    return hmac.new(TOKEN.encode("utf-8"), message, digestmod="sha256").hexdigest()


def _make_session(expires: int) -> str:
    return f"{expires}.{_session_signature(expires)}"


def _mark_session_cookie_partitioned(response: Response) -> None:
    """Partition the admin cookie so it works in HF's cross-site Space iframe."""
    cookie_prefix = f"{SESSION_COOKIE}=".encode("ascii")
    response.raw_headers = [
        (name, value + b"; Partitioned")
        if name.lower() == b"set-cookie" and value.startswith(cookie_prefix)
        else (name, value)
        for name, value in response.raw_headers
    ]


def _valid_admin_session(request: Request) -> bool:
    if not TOKEN_OK:
        return False
    raw = request.cookies.get(SESSION_COOKIE, "")
    try:
        expires_text, supplied = raw.split(".", 1)
        expires = int(expires_text)
    except (ValueError, TypeError):
        return False
    if expires <= int(time.time()):
        return False
    expected = _session_signature(expires)
    return hmac.compare_digest(supplied, expected)


def _admin_headers(response: Response) -> Response:
    response.headers["Cache-Control"] = "no-store"
    response.headers["X-Content-Type-Options"] = "nosniff"
    response.headers["Referrer-Policy"] = "no-referrer"
    # Hugging Face serves Space apps inside an iframe on huggingface.co. Do not
    # send X-Frame-Options (DENY/SAMEORIGIN would block that embed); use CSP to
    # permit only the HF Space page and same-origin frames.
    response.headers["Content-Security-Policy"] = (
        "default-src 'none'; style-src 'unsafe-inline'; script-src 'unsafe-inline'; "
        "connect-src 'self'; img-src 'self' data:; form-action 'self'; "
        "frame-ancestors 'self' https://huggingface.co; base-uri 'none'"
    )
    return response


@app.middleware("http")
async def require_token_or_admin_session(request: Request, call_next):
    path = request.url.path

    # Readiness is intentionally public so Hugging Face can check the Space.
    # The page itself only displays the login screen until a signed session exists.
    if path in ("/", "/healthz"):
        response = await call_next(request)
        return _admin_headers(response) if path == "/" else response

    # Login is the sole public administrative action. Every admin API, including
    # logout, requires the short-lived HttpOnly browser session.
    if path.startswith("/admin/"):
        if path == "/admin/login" and request.method == "POST":
            response = await call_next(request)
            return _admin_headers(response)
        if not TOKEN_OK:
            return JSONResponse({"error": "locked: configure MCP_AUTH_TOKEN"}, status_code=503)
        if not _valid_admin_session(request):
            return _admin_headers(JSONResponse({"error": "admin login required"}, status_code=401))
        if request.method in {"POST", "PUT", "PATCH", "DELETE"}:
            origin = request.headers.get("origin", "")
            scheme = request.headers.get("x-forwarded-proto", request.url.scheme).split(",", 1)[0].strip()
            expected = f"{scheme if scheme in {'http', 'https'} else request.url.scheme}://{request.url.netloc}"
            if request.headers.get("x-admin-action") != "1" or (origin and origin != expected):
                return _admin_headers(JSONResponse({"error": "same-origin administrative action required"}, status_code=403))
        response = await call_next(request)
        return _admin_headers(response)

    # MCP and Gradio API routes retain their shared-token authentication.
    if not TOKEN_OK:
        return JSONResponse({"error": "locked: set the MCP_AUTH_TOKEN secret (16+ chars)"}, status_code=503)
    auth = request.headers.get("authorization", "")
    supplied = (
        request.headers.get("x-mcp-token")
        or (auth[7:] if auth.lower().startswith("bearer ") else "")
        or request.query_params.get("token", "")
    )
    if not hmac.compare_digest(supplied.encode("utf-8"), TOKEN.encode("utf-8")):
        return JSONResponse({"error": "unauthorized"}, status_code=401)
    return await call_next(request)


@app.get("/", response_class=HTMLResponse)
def admin_home(request: Request):
    if not TOKEN_OK:
        return HTMLResponse(LOCKED_PAGE, status_code=503)
    page = DASHBOARD_PAGE if _valid_admin_session(request) else LOGIN_PAGE
    return HTMLResponse(page)


@app.post("/admin/login")
async def admin_login(request: Request):
    if not TOKEN_OK:
        return _admin_headers(JSONResponse({"error": "MCP_AUTH_TOKEN is not configured"}, status_code=503))
    body = await request.body()
    if len(body) > 4096:
        return _admin_headers(JSONResponse({"error": "request too large"}, status_code=413))
    try:
        payload = json.loads(body or b"{}")
    except (ValueError, TypeError):
        return _admin_headers(JSONResponse({"error": "invalid JSON body"}, status_code=400))
    supplied = payload.get("token", "") if isinstance(payload, dict) else ""
    if not isinstance(supplied, str) or not hmac.compare_digest(supplied.encode("utf-8"), TOKEN.encode("utf-8")):
        return _admin_headers(JSONResponse({"error": "Token inválido"}, status_code=401))

    expires = int(time.time()) + SESSION_TTL_SECONDS
    response = JSONResponse({"ok": True, "expires_at": expires})
    forwarded_proto = request.headers.get("x-forwarded-proto", "").split(",", 1)[0].strip().lower()
    secure = request.url.scheme == "https" or forwarded_proto == "https"
    response.set_cookie(
        SESSION_COOKIE,
        _make_session(expires),
        max_age=SESSION_TTL_SECONDS,
        httponly=True,
        secure=secure,
        # SameSite=None is required in the cross-site huggingface.co -> hf.space
        # iframe. Partitioned keeps the session scoped to that top-level site.
        samesite="none" if secure else "lax",
        path="/",
    )
    if secure:
        _mark_session_cookie_partitioned(response)
    return _admin_headers(response)


@app.post("/admin/logout")
def admin_logout(request: Request):
    response = JSONResponse({"ok": True})
    forwarded_proto = request.headers.get("x-forwarded-proto", "").split(",", 1)[0].strip().lower()
    secure = request.url.scheme == "https" or forwarded_proto == "https"
    response.delete_cookie(
        SESSION_COOKIE,
        httponly=True,
        secure=secure,
        samesite="none" if secure else "lax",
        path="/",
    )
    if secure:
        _mark_session_cookie_partitioned(response)
    return _admin_headers(response)


def _load_job_rows(ws: workspaces.Workspace) -> list[dict]:
    row = next((row for row in admin_overview()["workspaces"] if row["id"] == ws.id), None)
    return row["jobs"] if row else []


@app.get("/admin/api/overview")
def admin_overview():
    return runtime_client.get_client().request("overview")


def _cached_overview() -> dict:
    return admin_overview()


def _read_log_tail(ws: workspaces.Workspace, job_id: str, tail: int) -> dict:
    try:
        return runtime_client.get_client().request("logs", workspace=ws.id, job_id=job_id, tail=tail)
    except RuntimeError as exc:
        raise HTTPException(status_code=404, detail=str(exc)) from exc


def _live_workspace(workspace_id: str, metrics: dict, log_signatures: dict) -> dict:
    row = next((row for row in admin_overview()["workspaces"] if row["id"] == workspace_id), None)
    if row is None:
        raise HTTPException(status_code=404, detail="Workspace not found")
    jobs = row["jobs"][:20]
    for job in jobs:
        signature = (job.get("log_bytes", 0), job.get("base_offset", 0), job["status"])
        if log_signatures.get(job["id"]) != signature:
            job["log"] = _read_log_tail(workspaces.Workspace(workspace_id), job["id"], 16 * 1024)
            log_signatures[job["id"]] = signature
    return {"id": row["id"], "status": row["status"], "resources": row["resources"],
            "running_jobs": row["running_jobs"], "queued_jobs": row["queued_jobs"],
            "jobs": jobs, "history_total": row.get("history_total", len(row["jobs"])),
            "maintenance_error": row.get("maintenance_error"), "file_operations": row.get("file_operations", [])}


@app.get("/admin/api/metrics")
def admin_metrics():
    return runtime_client.get_client().request("metrics")


@app.get("/admin/api/workspaces/{workspace_id}/live")
def admin_workspace_live(workspace_id: str):
    if not WORKSPACE_ID_RE.fullmatch(workspace_id):
        raise HTTPException(status_code=404, detail="Workspace not found")
    return _live_workspace(workspace_id, admin_metrics(), {})


@app.get("/admin/api/live")
async def admin_live(request: Request, workspace: str = Query(default="")):
    if workspace and not WORKSPACE_ID_RE.fullmatch(workspace):
        raise HTTPException(status_code=400, detail="Invalid workspace")

    async def events():
        log_signatures = {}
        last_inventory_at = -10.
        yield ": connected\n\n"
        while not await request.is_disconnected():
            if not _valid_admin_session(request):
                yield 'event: session-expired\ndata: {}\n\n'
                return

            def collect():
                nonlocal last_inventory_at
                metrics = admin_metrics()
                payload = {"metrics": metrics}
                now = time.monotonic()
                if now - last_inventory_at >= 5:
                    payload["overview"] = _cached_overview()
                    last_inventory_at = now
                if workspace:
                    try:
                        payload["workspace"] = _live_workspace(workspace, metrics, log_signatures)
                    except HTTPException:
                        payload["workspace"] = {"id": workspace, "jobs": [], "missing": True}
                return payload

            payload = await run_in_threadpool(collect)
            yield "event: snapshot\ndata: " + json.dumps(payload, ensure_ascii=False, separators=(",", ":")) + "\n\n"
            await asyncio.sleep(1)

    return StreamingResponse(events(), media_type="text/event-stream",
                             headers={"Cache-Control": "no-store", "X-Accel-Buffering": "no"})


@app.get("/admin/api/workspaces/{workspace_id}/jobs/{job_id}/logs")
def admin_job_logs(workspace_id: str, job_id: str, tail: int = Query(default=64 * 1024, ge=1024, le=128 * 1024)):
    if not WORKSPACE_ID_RE.fullmatch(workspace_id) or not JOB_ID_RE.fullmatch(job_id):
        raise HTTPException(status_code=404, detail="Workspace or job not found")
    return _read_log_tail(workspaces.Workspace(workspace_id), job_id, int(tail))


def _admin_action(action, **parameters):
    try:
        return runtime_client.get_client().request(action, **parameters)
    except RuntimeError as exc:
        raise HTTPException(status_code=409, detail=str(exc)) from exc


@app.post("/admin/api/workspaces/{workspace_id}/jobs/{job_id}/cancel")
def admin_cancel_job(workspace_id: str, job_id: str):
    return _admin_action("cancel", workspace=workspace_id, job_id=job_id)


@app.post("/admin/api/workspaces/{workspace_id}/stop")
async def admin_stop_workspace(workspace_id: str, request: Request):
    body = await request.json()
    return await run_in_threadpool(_admin_action, "stop", workspace=workspace_id, recursive=bool(body.get("include_workflows", False)))


@app.post("/admin/api/workspaces/{workspace_id}/hibernate")
async def admin_hibernate_workspace(workspace_id: str, request: Request):
    body = await request.json()
    return await run_in_threadpool(_admin_action, "maintenance", workspace=workspace_id, operation="hibernate", manual=bool(body.get("manual", False)))


@app.post("/admin/api/workspaces/{workspace_id}/resume")
def admin_resume_workspace(workspace_id: str):
    return _admin_action("maintenance", workspace=workspace_id, operation="restore")


@app.post("/admin/api/executor/restart")
def admin_restart_executor():
    from tool_dispatch import stop_file_workers
    stop_file_workers()
    runtime_client.get_client().close()
    runtime_client.get_client()
    return {"status": "restarted", "message": "Command executor restarted; interrupted commands are not automatically rerun."}


@app.get("/healthz")
def healthz():
    state = "ready" if TOKEN_OK else "LOCKED: add the MCP_AUTH_TOKEN secret in Settings > Variables and secrets"
    storage = (
        "persistent when a Bucket is mounted at /data"
        if str(workspaces.ROOT).startswith("/data/")
        else "ephemeral; configure a Bucket at /data for persistence"
    )
    return PlainTextResponse(
        f"Claude Code MCP server: {state}\n"
        f"Runtime version: 2.1.0\n"
        f"Admin dashboard: / (MCP_AUTH_TOKEN login required)\n"
        f"MCP endpoint: /gradio_api/mcp/ (MCP_AUTH_TOKEN required)\n"
        f"workspace storage: {storage}\n"
        f"worker pool: 1 active workspace family, 1.5 vCPU / 12 GB RAM; storage budget: 50 GB global\n"
    )


# ssr_mode=False is required on Spaces: HF sets GRADIO_SSR_MODE=True, and mounting then starts
# a Node SSR server on port 7860, so uvicorn below dies with "address already in use".
app = gr.mount_gradio_app(app, demo, path="/", mcp_server=True, ssr_mode=False)


if __name__ == "__main__":
    uvicorn.run(app, host="0.0.0.0", port=PORT)