"""Four to eight local A2A 1.0 JSON-RPC services with isolated role contexts.

The callable adapter is the only inference boundary. This module does not call a
model on import. Run tests with explicit deterministic adapters; real experiments
must supply a real adapter and mode="model". HTTP is loopback-only, not production.
"""
from __future__ import annotations

import asyncio
import base64
import contextvars
import hashlib
import json
import re
import socket
import threading
from collections.abc import Awaitable, Callable, Sequence
from contextlib import asynccontextmanager
from datetime import datetime, timezone
from importlib.metadata import version
from pathlib import Path
from typing import Any
from uuid import uuid4

import httpx
import uvicorn
from google.protobuf.json_format import MessageToDict, ParseDict
from starlette.applications import Starlette

from a2a.client.transports.jsonrpc import JsonRpcTransport
from a2a.server.agent_execution import AgentExecutor, RequestContext
from a2a.server.events import EventQueue
from a2a.server.request_handlers import DefaultRequestHandler
from a2a.server.routes.agent_card_routes import create_agent_card_routes
from a2a.server.routes.jsonrpc_routes import create_jsonrpc_routes
from a2a.server.tasks import InMemoryTaskStore, TaskUpdater
from a2a.types.a2a_pb2 import (
    AgentCapabilities, AgentCard, AgentInterface, AgentSkill, GetTaskRequest,
    Message, Part, Role, SendMessageConfiguration, SendMessageRequest,
    Task, TaskState, TaskStatus,
)
from a2a.utils.proto_utils import validate_proto_required_fields

ROLES = ("buyer", "broker", "supplier", "inspector")
DEFAULT_ROLES = ROLES
ModelAdapter = Callable[[str, str], Awaitable[str]]
# Optional adapter instrumentation; populated separately for each concurrent task.
execution_context: contextvars.ContextVar[dict[str, str] | None] = contextvars.ContextVar(
    "a2a_execution_context", default=None
)


def utc_now() -> str:
    return datetime.now(timezone.utc).isoformat()


def as_dict(proto: Any) -> dict[str, Any]:
    return MessageToDict(proto)


def artifact_text(task: Task) -> str:
    return "\n".join(part.text for artifact in task.artifacts for part in artifact.parts if part.HasField("text"))


class JsonlLog:
    """Append-only records with a content hash chain (not a signed audit log)."""
    def __init__(self, directory: Path, mode: str):
        self.directory = directory
        directory.mkdir(parents=True, exist_ok=True)
        self.path = directory / "events.jsonl"
        if self.path.exists():
            raise FileExistsError(f"Use a fresh run directory; existing log: {self.path}")
        self.mode = mode
        self._lock = threading.Lock()
        self._previous_hash = "0" * 64
        self._sequence = 0

    def write(self, event: str, **fields: Any) -> dict[str, Any]:
        with self._lock:
            self._sequence += 1
            record = {"sequence": self._sequence, "timestamp": utc_now(), "mode": self.mode,
                      "event": event, "previous_hash": self._previous_hash, **fields}
            canonical = json.dumps(record, ensure_ascii=False, sort_keys=True, separators=(",", ":"))
            record["record_hash"] = hashlib.sha256(canonical.encode("utf-8")).hexdigest()
            with self.path.open("a", encoding="utf-8", newline="\n") as file:
                file.write(json.dumps(record, ensure_ascii=False) + "\n")
                file.flush()
            self._previous_hash = record["record_hash"]
            return record


class WireRecorder:
    """Record exact ASGI payload bytes plus JSON/UTF-8 views of local wire data."""
    def __init__(self, app, role: str, log: JsonlLog):
        self.app, self.role, self.log = app, role, log

    async def __call__(self, scope, receive, send):
        if scope["type"] != "http":
            return await self.app(scope, receive, send)
        trace_id = str(uuid4())
        chunks: list[bytes] = []
        response_chunks: list[bytes] = []
        response_status, response_headers = None, []
        request_headers = [(key.decode("latin1"), value.decode("latin1")) for key, value in scope["headers"]]
        request_logged = False

        def payload(body: bytes) -> dict:
            result = {"body_base64": base64.b64encode(body).decode("ascii"), "body_utf8": body.decode("utf-8", errors="replace")}
            try:
                result["json"] = json.loads(body)
            except (ValueError, UnicodeDecodeError):
                pass
            return result

        def log_request():
            nonlocal request_logged
            if not request_logged:
                self.log.write("http_request", trace_id=trace_id, recipient=self.role,
                               method=scope["method"], path=scope["path"], headers=request_headers,
                               **payload(b"".join(chunks)))
                request_logged = True

        async def recorded_receive():
            message = await receive()
            if message["type"] == "http.request":
                chunks.append(message.get("body", b""))
                if not message.get("more_body", False):
                    log_request()
            return message

        async def recorded_send(message):
            nonlocal response_status, response_headers
            if message["type"] == "http.response.start":
                log_request()
                response_status = message["status"]
                response_headers = [(key.decode("latin1"), value.decode("latin1")) for key, value in message.get("headers", [])]
            if message["type"] == "http.response.body":
                response_chunks.append(message.get("body", b""))
                if not message.get("more_body", False):
                    self.log.write("http_response", trace_id=trace_id, sender=self.role,
                                   status=response_status, headers=response_headers,
                                   **payload(b"".join(response_chunks)))
            await send(message)

        await self.app(scope, recorded_receive, recorded_send)


class RecordedTaskStore(InMemoryTaskStore):
    def __init__(self, role: str, log: JsonlLog):
        super().__init__()
        self.role, self.log = role, log

    async def save(self, task: Task, context) -> None:
        await super().save(task, context)
        serialized = as_dict(task)
        self.log.write("task_snapshot", role=self.role, task=serialized)
        directory = self.log.directory / "tasks" / self.role
        directory.mkdir(parents=True, exist_ok=True)
        (directory / f"{task.id}.json").write_text(json.dumps(serialized, ensure_ascii=False, indent=2), encoding="utf-8")


class RoleExecutor(AgentExecutor):
    def __init__(self, role: str, adapter: ModelAdapter, log: JsonlLog,
                 instructions: str = "", model_timeout: float = 600):
        self.role, self.adapter, self.log = role, adapter, log
        self.instructions, self.model_timeout = instructions, model_timeout

    async def execute(self, context: RequestContext, event_queue: EventQueue) -> None:
        assert context.task_id and context.context_id and context.message
        updater = TaskUpdater(event_queue, context.task_id, context.context_id)
        if context.current_task is None:
            task = Task(id=context.task_id, context_id=context.context_id,
                        status=TaskStatus(state=TaskState.TASK_STATE_SUBMITTED),
                        history=[context.message])
            task.status.timestamp.GetCurrentTime()
            await event_queue.enqueue_event(task)
        await updater.start_work()
        if context.metadata.get("operation") == "deliver-material":
            # Application-layer mailbox delivery is deterministic, not inference.
            material = [as_dict(part.data) for part in context.message.parts if part.HasField("data")]
            if len(material) != 1 or material[0].get("to") != self.role:
                await updater.failed(updater.new_agent_message([Part(text="Invalid recipient material")]))
                return
            self.log.write("material_received", role=self.role, task_id=context.task_id,
                           context_id=context.context_id, material=material[0])
            await updater.add_artifact([ParseDict({"data": {"received": True, "stage": material[0]["stage"], "recipient": self.role}}, Part())],
                                       name="delivery-receipt", metadata={"inference": False})
            await updater.complete()
            return
        supplied = context.get_user_input()
        prompt = f"{self.instructions}\n\n{supplied}" if self.instructions else supplied
        info = {"role": self.role, "task_id": context.task_id, "context_id": context.context_id}
        token = execution_context.set(info)
        self.log.write("adapter_request", **info, prompt=prompt, metadata=context.metadata)
        try:
            response = await asyncio.wait_for(self.adapter(self.role, prompt), timeout=self.model_timeout)
            if not isinstance(response, str) or not response.strip():
                raise ValueError("Adapter must return non-empty response text")
            self.log.write("adapter_response", **info, response=response)
            await updater.add_artifact([Part(text=response)], name=f"{self.role}-response",
                                       metadata={"role": self.role, "mode": self.log.mode})
            await updater.complete(updater.new_agent_message([Part(text="Artifact ready")]))
        except asyncio.CancelledError:
            self.log.write("adapter_canceled", **info)
            raise
        except Exception as error:
            self.log.write("adapter_error", **info, error_type=type(error).__name__, error=str(error))
            await updater.failed(updater.new_agent_message([Part(text=f"Adapter failed: {type(error).__name__}: {error}")]))
        finally:
            execution_context.reset(token)

    async def cancel(self, context: RequestContext, event_queue: EventQueue) -> None:
        updater = TaskUpdater(event_queue, context.task_id, context.context_id)
        await updater.cancel(updater.new_agent_message([Part(text="Task canceled")]))


class Cluster:
    """Async context manager for isolated role services over real loopback TCP.

    Each server has its own card, task store, executor, and HTTP port. Context IDs
    are generated by the server. SendMessage messages contain the actual prompt;
    no shared model memory is maintained. The external experiment controller may
    forward selected artifacts. Its routing decisions must be logged separately.
    """
    def __init__(self, adapter: ModelAdapter, run_dir: str | Path, *, mode: str,
                 roles: Sequence[str] = DEFAULT_ROLES,
                 instructions: dict[str, str] | None = None, model_timeout: float = 600,
                 base_port: int | None = None, role_descriptions: dict[str, str] | None = None):
        if mode not in {"model", "transport-test", "scripted-control"}:
            raise ValueError("Explicit mode must be model, transport-test, or scripted-control")
        if isinstance(roles, str):
            raise ValueError("roles must be a sequence of 4 to 8 distinct role names")
        self.roles = tuple(roles)
        if not 4 <= len(self.roles) <= 8 or len(set(self.roles)) != len(self.roles):
            raise ValueError("Provide 4 to 8 distinct role names")
        if any(not isinstance(role, str) or not re.fullmatch(r"[a-z][a-z0-9_-]{0,63}", role) or role == "controller" for role in self.roles):
            raise ValueError("Role names must be lowercase safe identifiers; controller is reserved")
        for mapping in (instructions or {}, role_descriptions or {}):
            if set(mapping) - set(self.roles):
                raise ValueError("Role configuration names must be present in roles")
        self.log = JsonlLog(Path(run_dir), mode)
        self.adapter, self.instructions = adapter, instructions or {}
        self.role_descriptions = role_descriptions or {}
        self.model_timeout, self.base_port = model_timeout, base_port
        self.urls: dict[str, str] = {}
        self.cards: dict[str, AgentCard] = {}
        self.servers, self.server_tasks, self.sockets, self.handlers = [], [], [], []
        self.http: httpx.AsyncClient | None = None
        self._context_owners: dict[str, str] = {}
        self._scoped_contexts: dict[tuple[str, str], str] = {}
        self._scope_locks: dict[tuple[str, str], asyncio.Lock] = {}

    async def __aenter__(self):
        self.log.write("run_started", sdk={"a2a-sdk": version("a2a-sdk"), "protobuf": version("protobuf")},
                       specification_release="1.0.1", wire_version="1.0", binding="JSONRPC",
                       roles=list(self.roles), security="loopback only; no authentication or TLS",
                       context_policy="server-issued contexts; optional context_scope reuses a separate context for each recipient; model histories remain application-owned")
        try:
            for index, role in enumerate(self.roles):
                sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
                sock.bind(("127.0.0.1", self.base_port + index if self.base_port else 0))
                sock.listen(128)
                sock.setblocking(False)
                self.sockets.append(sock)
                url = f"http://127.0.0.1:{sock.getsockname()[1]}"
                self.urls[role] = url
                card = AgentCard(
                    name=f"AgentCollusion Research {role}", description=self.role_descriptions.get(role, f"Local synthetic-transaction {role} agent"),
                    version="0.2.0",
                    supported_interfaces=[AgentInterface(url=url + "/rpc", protocol_binding="JSONRPC", protocol_version="1.0")],
                    capabilities=AgentCapabilities(streaming=True),
                    default_input_modes=["text/plain", "application/json"], default_output_modes=["text/plain", "application/json"],
                    skills=[AgentSkill(id=role, name=role, description=f"Perform {role} work in the synthetic transaction", tags=["research", role])],
                )
                validate_proto_required_fields(card)
                self.cards[role] = card
                store = RecordedTaskStore(role, self.log)
                executor = RoleExecutor(role, self.adapter, self.log, self.instructions.get(role, ""), self.model_timeout)
                handler = DefaultRequestHandler(agent_executor=executor, task_store=store, agent_card=card)
                self.handlers.append(handler)
                app = Starlette(routes=[*create_agent_card_routes(card), *create_jsonrpc_routes(handler, "/rpc", enable_v0_3_compat=False)])
                server = uvicorn.Server(uvicorn.Config(WireRecorder(app, role, self.log), log_level="critical", lifespan="off", access_log=False))
                self.servers.append(server)
                self.server_tasks.append(asyncio.create_task(server.serve(sockets=[sock])))
            for _ in range(200):
                if all(server.started for server in self.servers):
                    break
                for task in self.server_tasks:
                    if task.done():
                        task.result()
                await asyncio.sleep(0.01)
            else:
                raise TimeoutError("Local A2A services did not start")
            self.http = httpx.AsyncClient(timeout=self.model_timeout + 30, trust_env=False,
                                          headers={"A2A-Version": "1.0"})
            for role in self.roles:
                response = await self.http.get(self.urls[role] + "/.well-known/agent-card.json")
                response.raise_for_status()
                discovered = ParseDict(response.json(), AgentCard())
                validate_proto_required_fields(discovered)
                if discovered != self.cards[role]:
                    raise ValueError("Discovered card differs from advertised card")
                self.cards[role] = discovered
                card_dir = self.log.directory / "agent-cards"
                card_dir.mkdir(exist_ok=True)
                (card_dir / f"{role}.json").write_text(json.dumps(response.json(), indent=2), encoding="utf-8")
            self.log.write("discovery_complete", urls=self.urls)
            return self
        except BaseException:
            await self.__aexit__(None, None, None)
            raise

    def transport(self, role: str) -> JsonRpcTransport:
        if role not in self.roles or self.http is None:
            raise ValueError("Unknown role or cluster not started")
        card = self.cards[role]
        interface = next(i for i in card.supported_interfaces if i.protocol_binding == "JSONRPC" and i.protocol_version == "1.0")
        return JsonRpcTransport(self.http, card, interface.url)

    @asynccontextmanager
    async def _resolve_context(self, recipient: str, context_id: str | None,
                               context_scope: str | None):
        if recipient not in self.roles:
            raise ValueError("Unknown recipient")
        if context_id is not None and context_scope is not None:
            raise ValueError("Use either context_id or context_scope, not both")
        if context_scope is not None:
            if not isinstance(context_scope, str) or not context_scope.strip() or len(context_scope) > 256:
                raise ValueError("context_scope must be a non-empty string of at most 256 characters")
            key = (context_scope, recipient)
            lock = self._scope_locks.setdefault(key, asyncio.Lock())
            async with lock:
                yield self._scoped_contexts.get(key)
        else:
            if context_id is not None and self._context_owners.get(context_id) != recipient:
                raise ValueError("context_id must have been issued to this recipient in this cluster")
            yield context_id

    def _remember_context(self, recipient: str, task: Task, requested_context: str | None,
                          context_scope: str | None):
        context_id = task.context_id
        if not context_id or (requested_context is not None and context_id != requested_context):
            raise ValueError("Server returned a missing or mismatched context")
        if context_id in self._context_owners and self._context_owners[context_id] != recipient:
            raise ValueError("Server context unexpectedly collides across roles")
        self._context_owners[context_id] = recipient
        if context_scope is not None:
            self._scoped_contexts[(context_scope, recipient)] = context_id
        self.log.write("context_binding", recipient=recipient, context_id=context_id,
                       context_scope=context_scope, task_id=task.id)

    async def send(self, sender: str, recipient: str, prompt: str, *,
                   metadata: dict | None = None, context_id: str | None = None,
                   context_scope: str | None = None, task_id: str | None = None,
                   return_immediately: bool = False) -> Task:
        """Create/continue a Task; optionally reuse a recipient-specific scoped context.

        New Task IDs are always server-issued. Each (context_scope, recipient)
        pair has an independent A2A context. Use a scope shared across rounds to
        retain that role's protocol context, or a new scope for independent trials.
        The adapter/controller must separately manage its private model history.
        """
        async with self._resolve_context(recipient, context_id, context_scope) as selected:
            task = await self._send(sender, recipient, prompt, metadata=metadata,
                                    context_id=selected, task_id=task_id,
                                    return_immediately=return_immediately)
            self._remember_context(recipient, task, selected, context_scope)
            return task

    async def _send(self, sender: str, recipient: str, prompt: str, *,
                   metadata: dict | None = None, context_id: str | None = None,
                   task_id: str | None = None, return_immediately: bool = False) -> Task:
        if sender not in {*self.roles, "controller"}:
            raise ValueError("Unknown sender")
        request = SendMessageRequest(
            message=Message(message_id=str(uuid4()), role=Role.ROLE_USER,
                            parts=[Part(text=prompt)], context_id=context_id, task_id=task_id),
            configuration=SendMessageConfiguration(accepted_output_modes=["text/plain"], return_immediately=return_immediately),
            metadata={"sender": sender, "recipient": recipient, **(metadata or {})},
        )
        validate_proto_required_fields(request)
        self.log.write("routing_decision", sender=sender, recipient=recipient,
                       message_id=request.message.message_id, metadata=metadata or {},
                       context_id=context_id, task_id=task_id)
        response = await self.transport(recipient).send_message(request)
        validate_proto_required_fields(response)
        if not response.HasField("task"):
            raise ValueError("This experiment requires Task response with Artifact")
        return response.task

    async def get(self, recipient: str, task_id: str, history_length: int | None = None) -> Task:
        return await self.transport(recipient).get_task(GetTaskRequest(id=task_id, history_length=history_length))

    async def deliver(self, sender: str, recipient: str, material: dict, *,
                      context_id: str | None = None, context_scope: str | None = None) -> Task:
        async with self._resolve_context(recipient, context_id, context_scope) as selected:
            task = await self._deliver(sender, recipient, material, context_id=selected)
            self._remember_context(recipient, task, selected, context_scope)
            return task

    async def _deliver(self, sender: str, recipient: str, material: dict, *, context_id: str | None = None) -> Task:
        """Deliver a producer artifact and recipient-only messages without inference.

        The controller invokes this SDK client operation on behalf of the sender.
        Sender names are experimental routing labels, not authenticated identities.
        """
        if sender not in self.roles or recipient not in self.roles:
            raise ValueError("Unknown sender or recipient")
        if material.get("from") != sender or material.get("to") != recipient:
            raise ValueError("Material routing labels must match")
        if "stage" not in material:
            raise ValueError("Material must include a stage/round label")
        request = SendMessageRequest(
            message=Message(message_id=str(uuid4()), role=Role.ROLE_USER,
                            parts=[ParseDict({"data": material}, Part())], context_id=context_id),
            configuration=SendMessageConfiguration(accepted_output_modes=["application/json"]),
            metadata={"sender": sender, "recipient": recipient, "operation": "deliver-material"},
        )
        validate_proto_required_fields(request)
        self.log.write("material_delivery", sender=sender, recipient=recipient,
                       message_id=request.message.message_id, source_stage=material["stage"])
        response = await self.transport(recipient).send_message(request)
        validate_proto_required_fields(response)
        if not response.HasField("task") or response.task.status.state != TaskState.TASK_STATE_COMPLETED:
            raise ValueError("Material delivery was not acknowledged")
        return response.task

    async def raw_rpc(self, recipient: str, method: str, params: dict, *,
                      a2a_version: str = "1.0") -> httpx.Response:
        """Test helper for error probes. Normal calls use the official SDK client."""
        return await self.http.post(self.cards[recipient].supported_interfaces[0].url,
                                    json={"jsonrpc": "2.0", "id": str(uuid4()), "method": method, "params": params},
                                    headers={"A2A-Version": a2a_version})

    async def __aexit__(self, exc_type, exc, traceback):
        for handler in self.handlers:
            await handler.aclose()
        if self.http is not None:
            await self.http.aclose()
        for server in self.servers:
            server.should_exit = True
        if self.server_tasks:
            await asyncio.gather(*self.server_tasks, return_exceptions=True)
        for sock in self.sockets:
            sock.close()
        self.log.write("run_finished", error=str(exc) if exc else None)
