"""Streamable HTTP transport for MCP (with elicitation support).

Exposes ``POST /mcp`` for client→server requests. When a tool call drives
the skill into ``ctx.collect`` / ``ctx.ask``, the server pushes
``elicitation/create`` requests back to the client interleaved on the SSE
response, and routes the client's responses (delivered via a follow-up
``POST /mcp`` on the same ``Mcp-Session-Id``) to the in-flight future.

Content negotiation:
  * ``Accept: application/json`` (or no Accept) — one-shot JSON response.
    Skills that call ``ctx.collect``/``ctx.ask`` over this path will time
    out, because there is no channel to deliver the request to the client.
  * ``Accept`` includes ``text/event-stream`` — SSE response. Elicitation
    works end-to-end.

Mountable into an existing FastAPI app via :func:`mount_http`, or
serveable standalone via :func:`build_http_app`.
"""
from __future__ import annotations

import asyncio
import json
import logging
import os
import secrets
from typing import Any, AsyncIterator

from fastapi import APIRouter, FastAPI, Header, HTTPException, Request, Response
from fastapi.responses import JSONResponse, StreamingResponse

from ..agent import A2AAgent
from ..auth import NoAuth
from ..context import LocalRunContext, RunContext
from .server import (
    MCPServer,
    _INVALID_REQUEST,
    _PARSE_ERROR,
    _error,
)

log = logging.getLogger(__name__)


# Caller identity cache: token-hash → (user_id, expires_at). Avoids
# hammering CP's /v1/me on every tools/call. 60s TTL — short enough to
# pick up password rotations / revocations within a minute.
_ME_CACHE: dict[str, tuple[int, float]] = {}
_ME_CACHE_TTL_S = 60.0


def _extract_bearer(authorization: str | None) -> str | None:
    if not authorization or not authorization.lower().startswith("bearer "):
        return None
    return authorization.split(None, 1)[1].strip() or None


def _check_key(authorization: str | None) -> None:
    """Legacy single-shared-key gate. Honored when A2A_API_KEY is set
    (used for self-hosted deployments without a control plane)."""
    api_key = os.environ.get("A2A_API_KEY")
    if api_key is None:
        return
    token = _extract_bearer(authorization)
    if token is None:
        raise HTTPException(401, "missing bearer token")
    if token != api_key:
        raise HTTPException(401, "invalid bearer token")


async def _verify_cp_jwt(token: str, cp_url: str) -> int | None:
    """Return the caller's user_id by hitting ``GET {cp_url}/v1/me`` with
    the bearer token. Cached briefly. Returns None on any failure so the
    caller decides between "anonymous OK" (public agent) or 401 (private)."""
    import hashlib

    import httpx

    h = hashlib.sha256(token.encode()).hexdigest()
    now = asyncio.get_event_loop().time()
    cached = _ME_CACHE.get(h)
    if cached and cached[1] > now:
        return cached[0]
    try:
        async with httpx.AsyncClient(timeout=5.0) as c:
            r = await c.get(
                f"{cp_url.rstrip('/')}/v1/me",
                headers={"authorization": f"bearer {token}"},
            )
    except httpx.HTTPError:
        return None
    if r.status_code != 200:
        return None
    body = r.json() or {}
    uid = body.get("id")
    if not isinstance(uid, int):
        return None
    _ME_CACHE[h] = (uid, now + _ME_CACHE_TTL_S)
    return uid


async def _enforce_agent_auth(authorization: str | None) -> int | None:
    """Apply platform-level auth (called in addition to ``_check_key``).

    Returns the verified caller's user_id, or None for an unauthenticated
    request to a public agent. Raises 401/403 for unauthenticated calls
    to a private agent or for owner mismatches.

    Env vars read on every call (cheap, lets the operator hot-swap a
    deploy without restarting the pod). Agents that don't have these set
    fall back to the legacy ``A2A_API_KEY`` behavior — fully open if
    that's also unset.
    """
    cp_url = os.environ.get("A2A_CP_URL")
    public_env = os.environ.get("A2A_AGENT_PUBLIC", "").strip().lower()
    owner_env = os.environ.get("A2A_AGENT_OWNER_ID", "").strip()
    is_public = public_env in {"", "true", "1", "yes"}

    token = _extract_bearer(authorization)
    if token is None:
        if not is_public:
            raise HTTPException(401, "agent is private; missing bearer token")
        return None
    if not cp_url:
        # No control plane configured for verification — allow with no
        # identity (preserves self-hosted deployments).
        return None
    uid = await _verify_cp_jwt(token, cp_url)
    if uid is None:
        if not is_public:
            raise HTTPException(401, "invalid bearer token")
        return None
    if not is_public and owner_env:
        try:
            owner_id = int(owner_env)
        except ValueError:
            raise HTTPException(500, "invalid A2A_AGENT_OWNER_ID env") from None
        if uid != owner_id:
            raise HTTPException(403, "not the owner of this private agent")
    return uid


def _wants_sse(accept: str | None) -> bool:
    if not accept:
        return False
    return "text/event-stream" in accept.lower()


def _is_jsonrpc_response(body: dict[str, Any]) -> bool:
    """A JSON-RPC message is a response (vs request/notification) when it
    has an id and a result/error but no method."""
    if "method" in body:
        return False
    if body.get("id") is None:
        return False
    return "result" in body or "error" in body


class ElicitError(Exception):
    """Raised when the client declines/cancels an elicitation request or
    returns a JSON-RPC error in response to one."""


class _Session:
    """Per-client MCP session. Tracks server→client requests in flight so
    the client's response POSTs can resolve the right pending future."""

    def __init__(self, agent: A2AAgent) -> None:
        self.id = secrets.token_urlsafe(16)
        self.agent = agent
        self._pending: dict[str, asyncio.Future[Any]] = {}
        self._next_id = 0

    def _new_request_id(self) -> str:
        self._next_id += 1
        return f"srv-{self.id[:8]}-{self._next_id}"

    async def send_request(
        self,
        queue: "asyncio.Queue[dict[str, Any] | None]",
        method: str,
        params: dict[str, Any],
    ) -> Any:
        """Push a JSON-RPC request onto ``queue`` and await the response
        the client delivers via :meth:`deliver_response`."""
        rid = self._new_request_id()
        loop = asyncio.get_running_loop()
        fut: asyncio.Future[Any] = loop.create_future()
        self._pending[rid] = fut
        await queue.put(
            {"jsonrpc": "2.0", "id": rid, "method": method, "params": params}
        )
        try:
            return await fut
        finally:
            self._pending.pop(rid, None)

    def deliver_response(self, message: dict[str, Any]) -> bool:
        """Route a JSON-RPC response from the client to the matching
        pending server→client future. Returns True on hit."""
        rid = message.get("id")
        if not isinstance(rid, str):
            return False
        fut = self._pending.get(rid)
        if fut is None or fut.done():
            return False
        if "error" in message:
            err = message["error"] or {}
            fut.set_exception(
                ElicitError(err.get("message") or "client returned error")
            )
        else:
            fut.set_result(message.get("result"))
        return True


def _ask_schema() -> dict[str, Any]:
    """JSON Schema used to elicit a plain-text answer for ``ctx.ask``."""
    return {
        "type": "object",
        "required": ["answer"],
        "properties": {
            "answer": {
                "type": "string",
                "description": "Your answer",
            }
        },
    }


def _build_elicit_context(
    agent: A2AAgent, skill_name: str, session: _Session,
    queue: "asyncio.Queue[dict[str, Any] | None]",
    progress_token: Any = None,
) -> LocalRunContext[Any]:
    """Build a per-call context whose ``on_event`` translates skill
    events into MCP JSON-RPC traffic on the SSE stream:

    * ``input_request`` / ``question`` → ``elicitation/create`` request
      (server→client RPC; resolved when the client responds).
    * ``progress`` / ``text_delta`` / any other event → MCP
      ``notifications/progress`` notification. Lets long-running skills
      reset the client's tool-call idle timer; without these, clients
      like Claude Code abort the call after ~60s even though the HTTP
      stream stays open.

    ``progress_token`` is what the MCP client passed in
    ``params._meta.progressToken``; we echo it on every progress
    notification so the client can correlate with the originating
    tools/call. If the client didn't supply one we synthesize a token
    (the upstream tools/call request id) — most clients accept this and
    use it as a keep-alive signal.
    """
    auth_cls = type(agent).auth_model
    try:
        auth = auth_cls()
    except Exception:  # noqa: BLE001
        auth = NoAuth()
    ctx = LocalRunContext(auth=auth, task_id=f"mcp-{skill_name}")

    async def on_event(ev: Any) -> None:
        kind = getattr(ev, "kind", None)
        payload = getattr(ev, "payload", {}) or {}
        if kind == "input_request":
            request_id = payload.get("request_id")
            if not request_id:
                return
            title = payload.get("title") or ""
            reason = payload.get("reason") or ""
            message = f"{title}\n\n{reason}".strip() if reason else title
            schema = payload.get("schema") or {"type": "object"}
            try:
                result = await session.send_request(
                    queue,
                    "elicitation/create",
                    {"message": message, "requestedSchema": schema},
                )
            except ElicitError as exc:
                log.warning("elicitation/create failed: %s", exc)
                return
            if not isinstance(result, dict):
                return
            if result.get("action") == "accept":
                content = result.get("content")
                if isinstance(content, dict):
                    RunContext.submit_input(request_id, content)
            # decline / cancel: leave the registry future unresolved;
            # skill's own asyncio.wait_for timeout will fire.
            return
        if kind == "question":
            question_id = payload.get("question_id")
            if not question_id:
                return
            prompt = payload.get("prompt") or ""
            try:
                result = await session.send_request(
                    queue,
                    "elicitation/create",
                    {"message": prompt, "requestedSchema": _ask_schema()},
                )
            except ElicitError as exc:
                log.warning("elicitation/create failed: %s", exc)
                return
            if not isinstance(result, dict):
                return
            if result.get("action") == "accept":
                content = result.get("content") or {}
                answer = content.get("answer", "")
                if isinstance(answer, str):
                    RunContext.answer(question_id, answer)
            return
        # Any other event kind → forward as notifications/progress so
        # the MCP client gets liveness signals and doesn't abort the
        # tool call on its idle timer. Fire-and-forget — we don't wait
        # for a response from the client.
        message: str | None = None
        if isinstance(payload, dict):
            for key in ("message", "text", "summary"):
                value = payload.get(key)
                if isinstance(value, str) and value:
                    message = value
                    break
        progress_params: dict[str, Any] = {
            "progressToken": progress_token,
            "kind": kind,
        }
        if message is not None:
            progress_params["message"] = message
        await queue.put({
            "jsonrpc": "2.0",
            "method": "notifications/progress",
            "params": progress_params,
        })

    ctx._on_event = on_event  # noqa: SLF001
    return ctx


def mount_http(app: FastAPI, agent: A2AAgent, *, prefix: str = "/mcp") -> None:
    """Mount the MCP endpoint onto an existing FastAPI app."""
    # Sessions live for the lifetime of this app instance. The header
    # value is opaque to clients — they just echo whatever we returned
    # on the response.
    sessions: dict[str, _Session] = {}

    router = APIRouter()

    def _get_or_create_session(session_id: str | None) -> _Session:
        if session_id and session_id in sessions:
            return sessions[session_id]
        s = _Session(agent)
        sessions[s.id] = s
        return s

    @router.post(prefix)
    async def mcp_endpoint(
        request: Request,
        authorization: str | None = Header(default=None),
        accept: str | None = Header(default=None),
        mcp_session_id: str | None = Header(default=None, alias="Mcp-Session-Id"),
    ) -> Any:
        _check_key(authorization)
        caller_user_id = await _enforce_agent_auth(authorization)
        try:
            body = await request.json()
        except Exception as exc:  # noqa: BLE001
            return _error(None, _PARSE_ERROR, f"parse error: {exc}")
        if not isinstance(body, dict):
            return _error(None, _INVALID_REQUEST, "request must be a JSON object")

        # JSON-RPC response from the client (e.g. an elicitation result)?
        if _is_jsonrpc_response(body):
            session = sessions.get(mcp_session_id or "")
            if session is None:
                raise HTTPException(404, "unknown Mcp-Session-Id")
            if not session.deliver_response(body):
                raise HTTPException(404, "no pending server-initiated request with that id")
            return Response(status_code=202)

        # Client request / notification.
        session = _get_or_create_session(mcp_session_id)
        method = body.get("method")

        # When the caller is authenticated (has a verified CP JWT), inject
        # _meta.cp_jwt + cp_url + bucket onto tools/call params so skills
        # see the same caller context they'd get via the orchestrator.
        # Direct callers (Authorization header) and gateway-forwarded
        # callers (_meta on the body) end up in the same code path.
        if method == "tools/call" and caller_user_id is not None:
            token = _extract_bearer(authorization)
            if token:
                params = body.get("params") or {}
                if isinstance(params, dict):
                    meta = dict(params.get("_meta") or {})
                    meta.setdefault("cp_jwt", token)
                    meta.setdefault(
                        "cp_url",
                        os.environ.get("A2A_CP_URL", ""),
                    )
                    meta.setdefault("bucket", f"user-{caller_user_id}-files")
                    params["_meta"] = meta
                    body["params"] = params

        # Tools/call gets the elicit-capable path when the client accepts
        # SSE. Other methods (initialize, tools/list, ping, notifications)
        # never elicit; keep them as one-shot JSON.
        if method == "tools/call" and _wants_sse(accept):
            return _stream_tools_call(agent, session, body)

        server = MCPServer(agent)
        response = await server.handle(body)
        if response is None:
            return Response(
                status_code=202, headers={"Mcp-Session-Id": session.id}
            )
        return JSONResponse(response, headers={"Mcp-Session-Id": session.id})

    @router.get(prefix)
    async def mcp_get_unsupported() -> Any:
        # Streamable HTTP allows GET for standalone server-initiated
        # streams. We attach them inline to POST responses instead.
        raise HTTPException(405, "GET not supported; use POST /mcp")

    app.include_router(router)


def _stream_tools_call(
    agent: A2AAgent, session: _Session, body: dict[str, Any]
) -> StreamingResponse:
    """SSE response for ``tools/call``. Interleaves any
    ``elicitation/create`` requests the skill emits via
    ``ctx.collect`` / ``ctx.ask`` before yielding the final tool result."""
    queue: "asyncio.Queue[dict[str, Any] | None]" = asyncio.Queue()
    msg_id = body.get("id")
    params = body.get("params") or {}
    skill_name = params.get("name") if isinstance(params, dict) else None
    # Prefer the progressToken the MCP client supplied (per spec). Fall
    # back to the tools/call id so clients that don't opt in still get a
    # consistent token they can map to the call.
    progress_token: Any = msg_id
    if isinstance(params, dict):
        meta = params.get("_meta")
        if isinstance(meta, dict) and "progressToken" in meta:
            progress_token = meta.get("progressToken")

    # Build a context that turns local events into elicit-create requests
    # plus progress notifications.
    def build_ctx(name: str) -> LocalRunContext[Any]:
        return _build_elicit_context(
            agent, name, session, queue, progress_token=progress_token,
        )

    server = MCPServer(agent, context_builder=build_ctx)

    async def run_dispatch() -> dict[str, Any] | None:
        try:
            return await server.handle(body)
        except Exception as exc:  # noqa: BLE001
            log.exception("mcp tools/call dispatch failed")
            return {
                "jsonrpc": "2.0",
                "id": msg_id,
                "error": {
                    "code": -32603,
                    "message": f"{type(exc).__name__}: {exc}",
                },
            }

    async def gen() -> AsyncIterator[bytes]:
        dispatch_task = asyncio.create_task(run_dispatch())
        try:
            while True:
                queue_task = asyncio.create_task(queue.get())
                done, _ = await asyncio.wait(
                    {queue_task, dispatch_task},
                    return_when=asyncio.FIRST_COMPLETED,
                )
                if queue_task in done:
                    msg = queue_task.result()
                    if msg is not None:
                        yield (b"data: " + json.dumps(msg).encode() + b"\n\n")
                else:
                    queue_task.cancel()
                if dispatch_task in done:
                    final = dispatch_task.result()
                    # Drain any elicit requests still sitting in the queue.
                    while not queue.empty():
                        leftover = queue.get_nowait()
                        if leftover is not None:
                            yield (
                                b"data: "
                                + json.dumps(leftover).encode()
                                + b"\n\n"
                            )
                    if final is not None:
                        yield (b"data: " + json.dumps(final).encode() + b"\n\n")
                    break
        finally:
            if not dispatch_task.done():
                dispatch_task.cancel()
        _ = skill_name  # silence unused-var warning when debug logging is off

    return StreamingResponse(
        gen(),
        media_type="text/event-stream",
        headers={
            "Cache-Control": "no-cache",
            "X-Accel-Buffering": "no",
            "Mcp-Session-Id": session.id,
        },
    )


def build_http_app(agent: A2AAgent) -> FastAPI:
    """Build a minimal FastAPI app exposing only the MCP endpoint."""
    app = FastAPI(title=f"{type(agent).name} (MCP)", version=type(agent).version)
    mount_http(app, agent)
    return app
