"""MCP server with ACT + ECT + HTTP-signature enforcement middleware.

The server exposes two tools via FastMCP streamable-HTTP:

    search(query: str)      — returns a list of fake hits (str[])
    summarize(text: str)    — returns a short synthetic summary (str)

Every POST to /mcp goes through ``ActEctAuthMiddleware`` which:

    1. parses ``Authorization: Bearer <act-mandate>`` and verifies the
       Phase 1 mandate (ACT §8.1);
    2. parses ``Wimse-ECT: <ect-compact>`` and verifies the ECT
       (draft-nennemann-wimse-ect-01 §7);
    3. verifies the RFC 9421 HTTP-signature over
       @method/@target-uri/content-digest/wimse-ect with
       wimse-aud=mcp-server (per draft-ietf-wimse-http-signature-03);
    4. cross-checks exec_act ∈ mandate.cap[].action, mandate.sub == ECT.iss,
       ECT.inp_hash == sha256(body).

On failure any check returns HTTP 401/403 before the request reaches
FastMCP. Successful audits are appended to ``AUDIT_LOG_PATH`` so the
verifier CLI can reconstruct the server's view of the DAG.
"""

from __future__ import annotations

import argparse
import base64
import hashlib
import json
import logging
import os
import time
from pathlib import Path
from typing import Any

import uvicorn
from mcp.server.fastmcp import FastMCP
from starlette.middleware.base import BaseHTTPMiddleware
from starlette.requests import Request
from starlette.responses import JSONResponse, Response

from act.crypto import ACTKeyResolver, b64url_sha256
from act.errors import ACTError, ACTSignatureError
from act.verify import ACTVerifier

import ect as ect_pkg  # noqa: F401 (ensures package import)
from ect.verify import verify as ect_verify, VerifyOptions

from .http_sig import verify_request
from .keys import Identity, build_ect_key_resolver, build_key_registry, load_identities


LOG = logging.getLogger("poc.server")

AUDIT_LOG_PATH = Path(os.environ.get("POC_AUDIT_LOG", "keys/server-audit.jsonl"))
KEYS_DIR = os.environ.get("POC_KEYS_DIR", "keys")
SERVER_IDENTITY_NAME = "mcp-server"
AGENT_IDENTITY_NAME = "agent"


# ---- FastMCP tools ----------------------------------------------------------


def _build_mcp() -> FastMCP:
    """Fresh FastMCP instance with the two PoC tools registered.

    A new instance per ``build_app`` call keeps ``StreamableHTTPSessionManager``
    usable in tests that start the Starlette lifespan multiple times (the
    session manager is a single-use object).
    """
    from mcp.server.transport_security import TransportSecuritySettings

    mcp = FastMCP(
        "act-ect-poc",
        transport_security=TransportSecuritySettings(
            # Allow the test hostname ``testserver`` in addition to loopback.
            allowed_hosts=[
                "127.0.0.1:*", "localhost:*", "[::1]:*", "testserver",
            ],
            allowed_origins=[
                "http://127.0.0.1:*", "http://localhost:*",
                "http://[::1]:*", "http://testserver",
            ],
        ),
    )

    @mcp.tool()
    def search(query: str) -> list[str]:
        """Return three fake search hits for the query."""
        q = query.strip() or "empty"
        return [f"[{q}] result {i + 1}: lorem ipsum about {q}" for i in range(3)]

    @mcp.tool()
    def summarize(text: str) -> str:
        """Return a deterministic one-line summary of the text."""
        snippet = text.strip().replace("\n", " ")
        if len(snippet) > 120:
            snippet = snippet[:117] + "..."
        return f"Summary: {snippet}"

    return mcp


# ---- Auth middleware --------------------------------------------------------


class ActEctAuthMiddleware(BaseHTTPMiddleware):
    """Enforce ACT mandate + ECT + HTTP signature on every tool invocation."""

    def __init__(self, app, identities: dict[str, Identity]) -> None:
        super().__init__(app)
        self._identities = identities
        server = identities[SERVER_IDENTITY_NAME]
        agent = identities[AGENT_IDENTITY_NAME]

        registry = build_key_registry(identities)
        resolver = ACTKeyResolver(registry=registry)
        self._act_verifier = ACTVerifier(
            resolver,
            verifier_id=SERVER_IDENTITY_NAME,
            trusted_issuers={ident.name for ident in identities.values()},
        )
        self._ect_resolver = build_ect_key_resolver(identities)
        self._agent_public_key = agent.public_key
        self._server_name = server.name
        self._agent_name = agent.name

    async def dispatch(self, request: Request, call_next):
        # Only enforce on the MCP endpoint; let everything else through.
        if request.url.path != "/mcp" or request.method.upper() != "POST":
            return await call_next(request)

        try:
            auth_ctx = await self._authorize(request)
        except _AuthFailure as e:
            LOG.warning("auth rejected: %s", e.detail)
            return JSONResponse({"error": e.detail}, status_code=e.status)

        # Pass the verified context downstream so a handler could read it.
        request.state.act_ect = auth_ctx
        response: Response = await call_next(request)
        _append_audit(AUDIT_LOG_PATH, auth_ctx)
        return response

    async def _authorize(self, request: Request) -> dict[str, Any]:
        body = await request.body()

        # 1. Authorization: Bearer <act-mandate-compact>
        auth_hdr = request.headers.get("authorization", "")
        if not auth_hdr.lower().startswith("bearer "):
            raise _AuthFailure(401, "missing Authorization: Bearer")
        act_compact = auth_hdr[len("bearer ") :].strip()

        # 2. Wimse-ECT: <ect-compact>
        ect_compact = request.headers.get("wimse-ect", "").strip()
        if not ect_compact:
            raise _AuthFailure(401, "missing Wimse-ECT header")

        # 3. RFC 9421 HTTP signature (over method/target/content-digest/wimse-ect)
        sig_input = request.headers.get("signature-input", "")
        sig = request.headers.get("signature", "")
        content_digest_hdr = request.headers.get("content-digest", "")
        if not (sig_input and sig and content_digest_hdr):
            raise _AuthFailure(
                401, "missing Signature / Signature-Input / Content-Digest"
            )

        # Starlette's request.url gives us a URL; canonicalize target URI.
        # Use the path+query as the target since scheme/host differs behind
        # reverse proxies. For the PoC we match client and server on the
        # exact full URL the client signed.
        target_uri = str(request.url)

        try:
            parsed_sig = verify_request(
                method=request.method,
                target_uri=target_uri,
                body=body,
                wimse_ect_header=ect_compact,
                content_digest_header=content_digest_hdr,
                signature_input_header=sig_input,
                signature_header=sig,
                expected_audience=self._server_name,
                public_key=self._agent_public_key,
            )
        except ACTSignatureError as e:
            raise _AuthFailure(401, f"http-signature failed: {e}") from e

        # 4. ACT mandate
        try:
            mandate = self._act_verifier.verify_mandate(
                act_compact, check_sub=False
            )
        except ACTError as e:
            raise _AuthFailure(401, f"ACT mandate rejected: {e}") from e

        if mandate.sub != self._agent_name:
            raise _AuthFailure(
                403, f"mandate.sub {mandate.sub!r} != agent {self._agent_name!r}"
            )

        # 5. ECT
        try:
            ect_opts = VerifyOptions(
                verifier_id=self._server_name,
                resolve_key=self._ect_resolver,
            )
            parsed_ect = ect_verify(ect_compact, ect_opts)
        except Exception as e:  # ECT refimpl raises ValueError subclasses
            raise _AuthFailure(401, f"ECT rejected: {e}") from e

        if parsed_sig.keyid != parsed_ect.header.get("kid"):
            raise _AuthFailure(
                401,
                f"http-sig keyid {parsed_sig.keyid!r} != ect kid "
                f"{parsed_ect.header.get('kid')!r}",
            )

        # 6. Cross-check inp_hash binds to this body
        body_hash = b64url_sha256(body)
        if parsed_ect.payload.inp_hash and parsed_ect.payload.inp_hash != body_hash:
            raise _AuthFailure(401, "ECT.inp_hash does not match request body")

        # 7. exec_act must be authorised by mandate.cap
        cap_actions = {c.action for c in mandate.cap}
        if parsed_ect.payload.exec_act not in cap_actions:
            raise _AuthFailure(
                403,
                f"exec_act {parsed_ect.payload.exec_act!r} not in "
                f"mandate.cap {sorted(cap_actions)!r}",
            )

        # 8. ECT issuer should equal mandate subject (the executing agent)
        if parsed_ect.payload.iss != mandate.sub:
            raise _AuthFailure(
                403,
                f"ECT.iss {parsed_ect.payload.iss!r} != mandate.sub "
                f"{mandate.sub!r}",
            )

        return {
            "ts": int(time.time()),
            "mandate_jti": mandate.jti,
            "ect_jti": parsed_ect.payload.jti,
            "exec_act": parsed_ect.payload.exec_act,
            "pred": list(parsed_ect.payload.pred),
            "inp_hash": body_hash,
            "wimse_aud": parsed_sig.wimse_aud,
            "keyid": parsed_sig.keyid,
        }


class _AuthFailure(Exception):
    def __init__(self, status: int, detail: str) -> None:
        super().__init__(detail)
        self.status = status
        self.detail = detail


def _append_audit(path: Path, entry: dict[str, Any]) -> None:
    path.parent.mkdir(parents=True, exist_ok=True)
    with path.open("a", encoding="utf-8") as fh:
        fh.write(json.dumps(entry, separators=(",", ":")) + "\n")


# ---- ASGI assembly ----------------------------------------------------------


def build_app(identities: dict[str, Identity]):
    mcp = _build_mcp()
    app = mcp.streamable_http_app()
    app.add_middleware(ActEctAuthMiddleware, identities=identities)
    return app


def main() -> None:
    logging.basicConfig(
        level=os.environ.get("POC_LOG_LEVEL", "INFO"),
        format="%(asctime)s %(levelname)s %(name)s: %(message)s",
    )
    parser = argparse.ArgumentParser(description="ACT/ECT MCP PoC server")
    parser.add_argument("--host", default="127.0.0.1")
    parser.add_argument("--port", type=int, default=8765)
    parser.add_argument("--keys-dir", default=KEYS_DIR)
    args = parser.parse_args()

    identities = load_identities(args.keys_dir)
    app = build_app(identities)
    LOG.info(
        "serving MCP at http://%s:%d/mcp; audit=%s",
        args.host,
        args.port,
        AUDIT_LOG_PATH,
    )
    uvicorn.run(app, host=args.host, port=args.port, log_level="info")


if __name__ == "__main__":
    main()
