"""Mint and verify ACT mandates/records and ECT payloads for the PoC.

Thin convenience wrappers around ietf-act and ietf-ect that fix the
PoC's shape (one user issues to one agent, one MCP-server audience)
while leaving all cryptography and claim validation to the refimpls.
"""

from __future__ import annotations

import time
import uuid
from dataclasses import dataclass
from typing import Optional

from act.crypto import b64url_sha256, sign as act_sign
from act.token import (
    ACTMandate,
    ACTRecord,
    Capability,
    TaskClaim,
    encode_jws,
)

import ect
from ect.create import CreateOptions, create as ect_create
from ect.types import Payload as ECTPayload

from .keys import Identity


# Capability action names the agent may exercise on the MCP server.
# mcp.session.* covers JSON-RPC plumbing (initialize, tools/list, ping, …)
# that precedes real tool invocations.
MCP_CAPS = (
    "mcp.session.initialize",
    "mcp.session.list_tools",
    "mcp.session.other",
    "mcp.search",
    "mcp.summarize",
)

# Map tool-name → exec_act for real tool calls.
TOOL_ACTION = {
    "search": "mcp.search",
    "summarize": "mcp.summarize",
}


def exec_act_for_rpc_method(method: str, tool_name: str | None = None) -> str:
    """Return the ACT/ECT exec_act string for a JSON-RPC call.

    Rules:
      * ``tools/call`` dispatches to ``TOOL_ACTION[tool_name]``
      * ``initialize`` and ``tools/list`` have dedicated ``mcp.session.*`` actions
      * anything else collapses to ``mcp.session.other`` so the mandate can
        still authorise session-level JSON-RPC without enumerating every
        method on both sides.
    """
    if method == "tools/call":
        if tool_name is None or tool_name not in TOOL_ACTION:
            raise ValueError(f"unknown tool: {tool_name!r}")
        return TOOL_ACTION[tool_name]
    if method == "initialize":
        return "mcp.session.initialize"
    if method == "tools/list":
        return "mcp.session.list_tools"
    return "mcp.session.other"


@dataclass
class MintedMandate:
    compact: str
    mandate: ACTMandate


def mint_mandate(
    *,
    user: Identity,
    agent: Identity,
    audience: str,
    purpose: str,
    ttl_seconds: int = 900,
    now: Optional[int] = None,
) -> MintedMandate:
    """User issues a Phase 1 ACT mandate to the agent.

    Reference: ACT §3.1, §4.2.
    """
    iat = int(now if now is not None else time.time())
    mandate = ACTMandate(
        alg="ES256",
        kid=user.kid,
        iss=user.name,
        sub=agent.name,
        aud=audience,
        iat=iat,
        exp=iat + ttl_seconds,
        jti=str(uuid.uuid4()),
        wid=agent.name,
        task=TaskClaim(purpose=purpose, created_by=user.name),
        cap=[Capability(action=a) for a in MCP_CAPS],
    )
    mandate.validate()
    signature = act_sign(user.private_key, mandate.signing_input())
    compact = encode_jws(mandate, signature)
    return MintedMandate(compact=compact, mandate=mandate)


@dataclass
class MintedRecord:
    compact: str
    record: ACTRecord


def mint_exec_record(
    *,
    agent: Identity,
    mandate: ACTMandate,
    exec_act: str,
    pred_jtis: list[str],
    inp_body: bytes,
    out_body: bytes,
    now: Optional[int] = None,
    status: str = "completed",
) -> MintedRecord:
    """Agent mints a Phase 2 ACT execution record after a tool call.

    ``pred_jtis`` should include the mandate jti and any prior exec record
    jtis for this run (DAG semantics, ACT §4.3 ``pred``).

    Reference: ACT §3.2, §4.3.
    """
    exec_ts = int(now if now is not None else time.time())

    record = ACTRecord.from_mandate(
        mandate,
        kid=agent.kid,
        exec_act=exec_act,
        pred=pred_jtis,
        exec_ts=exec_ts,
        status=status,
        inp_hash=b64url_sha256(inp_body),
        out_hash=b64url_sha256(out_body),
    )
    # The record's jti MUST equal the mandate's jti per ACT §3.2 / §4 (the
    # Phase 2 token records the same task). ``from_mandate`` already copies
    # mandate.jti, so no override here.
    record.validate()
    signature = act_sign(agent.private_key, record.signing_input())
    compact = encode_jws(record, signature)
    return MintedRecord(compact=compact, record=record)


@dataclass
class MintedECT:
    compact: str
    payload: ECTPayload


def mint_ect(
    *,
    agent: Identity,
    audience: str,
    exec_act: str,
    pred_jtis: list[str],
    inp_body: bytes,
    now: Optional[int] = None,
    ttl_seconds: int = 300,
) -> MintedECT:
    """Agent mints an ECT binding the MCP request body to the execution.

    The ECT's ``jti`` identifies this specific invocation; ``pred`` links
    back to the mandate (and earlier ECT invocations, if any); the request
    body is hashed into ``inp_hash``.

    Reference: draft-nennemann-wimse-ect-01 §4.
    """
    iat = int(now if now is not None else time.time())
    payload = ECTPayload(
        iss=agent.name,
        aud=[audience],
        iat=iat,
        exp=iat + ttl_seconds,
        jti=str(uuid.uuid4()),
        exec_act=exec_act,
        pred=pred_jtis,
        wid=agent.name,
        inp_hash=b64url_sha256(inp_body),
    )
    # ECT refimpl wants a real ES256 key and the same kid shape as ACT.
    compact = ect_create(
        payload,
        agent.private_key,
        CreateOptions(key_id=agent.kid),
    )
    # Reflect the defaulted iat/exp back onto our payload copy for logging.
    return MintedECT(compact=compact, payload=payload)
