from __future__ import annotations

import os
from typing import Any

from grantflow.mcp.tools import DEFAULT_BASE_URL, GrantFlowToolClient


def create_fastmcp_server() -> Any:
    try:
        from mcp.server.fastmcp import FastMCP
    except ImportError as exc:
        raise RuntimeError(
            "Install the MCP extra to run the production MCP transport: pip install 'grantflow[mcp]'"
        ) from exc

    mcp = FastMCP("GrantFlow", stateless_http=True, json_response=True)

    def client() -> GrantFlowToolClient:
        return GrantFlowToolClient(
            base_url=os.getenv("GRANTFLOW_BASE_URL", DEFAULT_BASE_URL),
            api_key=os.getenv("GRANTFLOW_API_KEY") or None,
            bearer_token=os.getenv("GRANTFLOW_BEARER_TOKEN") or None,
        )

    @mcp.tool()
    def grantflow_onboard_agent(arguments: dict[str, Any]) -> dict[str, Any]:
        """Request API-key or OAuth client-credentials onboarding metadata."""
        with client() as grantflow:
            return grantflow.call_tool("grantflow_onboard_agent", arguments)

    @mcp.tool()
    def grantflow_register_agent(arguments: dict[str, Any]) -> dict[str, Any]:
        """Register a sandbox-scoped agent identity."""
        with client() as grantflow:
            return grantflow.call_tool("grantflow_register_agent", arguments)

    @mcp.tool()
    def grantflow_create_session(arguments: dict[str, Any]) -> dict[str, Any]:
        """Create a short-lived GrantFlow runtime session credential."""
        with client() as grantflow:
            return grantflow.call_tool("grantflow_create_session", arguments)

    @mcp.tool()
    def grantflow_introspect_agent(arguments: dict[str, Any]) -> dict[str, Any]:
        """Validate X-API-Key or Authorization Bearer credentials."""
        with client() as grantflow:
            return grantflow.call_tool("grantflow_introspect_agent", arguments)

    @mcp.tool()
    def grantflow_exchange_oauth_token(arguments: dict[str, Any]) -> dict[str, Any]:
        """Exchange self-serve OAuth client credentials for a Bearer access token."""
        with client() as grantflow:
            return grantflow.call_tool("grantflow_exchange_oauth_token", arguments)

    @mcp.tool()
    def grantflow_rotate_credential(arguments: dict[str, Any]) -> dict[str, Any]:
        """Rotate a self-serve API key or OAuth client secret."""
        with client() as grantflow:
            return grantflow.call_tool("grantflow_rotate_credential", arguments)

    @mcp.tool()
    def grantflow_revoke_credential(arguments: dict[str, Any]) -> dict[str, Any]:
        """Revoke a self-serve API key, Bearer token, or OAuth client secret."""
        with client() as grantflow:
            return grantflow.call_tool("grantflow_revoke_credential", arguments)

    @mcp.tool()
    def grantflow_run_preflight(arguments: dict[str, Any]) -> dict[str, Any]:
        """Run donor/readiness preflight before generation."""
        with client() as grantflow:
            return grantflow.call_tool("grantflow_run_preflight", arguments)

    @mcp.tool()
    def grantflow_start_generation(arguments: dict[str, Any]) -> dict[str, Any]:
        """Start a GrantFlow generation job."""
        with client() as grantflow:
            return grantflow.call_tool("grantflow_start_generation", arguments)

    @mcp.tool()
    def grantflow_get_status(arguments: dict[str, Any]) -> dict[str, Any]:
        """Read generation status by job_id."""
        with client() as grantflow:
            return grantflow.call_tool("grantflow_get_status", arguments)

    @mcp.tool()
    def grantflow_get_quality(arguments: dict[str, Any]) -> dict[str, Any]:
        """Read quality and grounding signals by job_id."""
        with client() as grantflow:
            return grantflow.call_tool("grantflow_get_quality", arguments)

    @mcp.tool()
    def grantflow_get_events(arguments: dict[str, Any]) -> dict[str, Any]:
        """Read job lifecycle events by job_id."""
        with client() as grantflow:
            return grantflow.call_tool("grantflow_get_events", arguments)

    @mcp.tool()
    def grantflow_run_sandbox_happy_path(arguments: dict[str, Any]) -> dict[str, Any]:
        """Run onboarding, generation, status, quality, and events as one smoke workflow."""
        with client() as grantflow:
            return grantflow.call_tool("grantflow_run_sandbox_happy_path", arguments)

    @mcp.tool()
    def grantflow_ingest_text(arguments: dict[str, Any]) -> dict[str, Any]:
        """Ingest pre-extracted text into the RAG corpus for a given donor and tenant.

        Use this when you have text content ready (e.g. extracted from a document or
        synthesised from research). No file upload required — pass the text as a string.
        Required: text, donor_id. Optional: title, tenant_id, doc_family, chunk_size, overlap, metadata.
        """
        with client() as grantflow:
            return grantflow.call_tool("grantflow_ingest_text", arguments)

    @mcp.tool()
    def grantflow_hitl_approve(arguments: dict[str, Any]) -> dict[str, Any]:
        """Approve or reject a pending Human-in-the-Loop (HITL) checkpoint.

        Required: checkpoint_id (str), approved (bool).
        Optional: feedback (str), request_id (str for idempotency).
        A job with status=pending_hitl cannot proceed until its checkpoint is decided.
        """
        with client() as grantflow:
            return grantflow.call_tool("grantflow_hitl_approve", arguments)

    @mcp.tool()
    def grantflow_hitl_list_pending(arguments: dict[str, Any]) -> dict[str, Any]:
        """List all pending HITL checkpoints waiting for a decision.

        Optional: tenant_id, donor_id, limit (default 50).
        Use this to discover which checkpoints need approval before jobs can continue.
        """
        with client() as grantflow:
            return grantflow.call_tool("grantflow_hitl_list_pending", arguments)

    @mcp.tool()
    def grantflow_get_export_payload(arguments: dict[str, Any]) -> dict[str, Any]:
        """Read the structured export payload for a completed job.

        Required: job_id.
        Returns the full export-ready JSON including toc_draft, logframe, quality signals,
        export_readiness status, and export_contract. Use this before deciding whether
        to trigger a document export or request further review.
        """
        with client() as grantflow:
            return grantflow.call_tool("grantflow_get_export_payload", arguments)

    return mcp


def main() -> None:
    transport = os.getenv("GRANTFLOW_MCP_TRANSPORT", "streamable-http")
    base_url = os.getenv("GRANTFLOW_BASE_URL", DEFAULT_BASE_URL)
    try:
        server = create_fastmcp_server()
    except RuntimeError as exc:
        print(
            f"\n[GrantFlow MCP] Cannot start streamable-http transport.\n"
            f"  Reason: {exc}\n\n"
            f"  Option A — install the MCP extra and retry:\n"
            f"    pip install 'grantflow[mcp]'\n"
            f"    GRANTFLOW_BASE_URL={base_url} python -m grantflow.mcp.fastmcp_server\n\n"
            f"  Option B — use the zero-dependency stdio server instead:\n"
            f"    GRANTFLOW_BASE_URL={base_url} python -m grantflow.mcp.server\n",
            flush=True,
        )
        raise SystemExit(1) from exc
    print(
        f"[GrantFlow MCP] Starting transport={transport} base_url={base_url}",
        flush=True,
    )
    server.run(transport=transport)


if __name__ == "__main__":
    main()
