"""Session-scoped MCP tools for ACP orchestration signals.

ACP jobs use this bridge to report structured outcomes back to the
orchestrator without relying on brittle text parsing.
"""

from __future__ import annotations

import asyncio
import json
import os
import sys
import uuid
from dataclasses import dataclass, field
from typing import Any

import mcp.server
import mcp.types as mcp_types
from mcp.server.stdio import stdio_server

from otto.log import get_logger

_log = get_logger("otto.acp.orchestration_tools")

_SOCKET_ENV_VAR = "OTTO_ACP_ORCHESTRATION_SOCKET"
_SERVER_NAME = "otto-orchestration"
_KNOWN_TOOLS = frozenset({"complete_task", "report_blocker", "hand_off_to", "request_review"})


@dataclass(frozen=True, slots=True)
class ACPTaskCompletion:
    result: str
    changed_files: tuple[str, ...] = ()
    outcome_summary: str | None = None


@dataclass(frozen=True, slots=True)
class ACPHandoffRequest:
    agent_name: str
    context: dict[str, Any] = field(default_factory=dict)


@dataclass(frozen=True, slots=True)
class ACPReviewRequest:
    files: tuple[str, ...] = ()
    summary: str | None = None


@dataclass(frozen=True, slots=True)
class ACPOrchestrationSnapshot:
    completion: ACPTaskCompletion | None = None
    blocker: str | None = None
    handoff: ACPHandoffRequest | None = None
    review: ACPReviewRequest | None = None


def _coerce_text(value: Any) -> str | None:
    if not isinstance(value, str):
        return None
    normalized = value.strip()
    return normalized or None


def _coerce_context(value: Any) -> dict[str, Any]:
    if isinstance(value, dict):
        return dict(value)
    return {}


def _coerce_string_list(value: Any) -> tuple[str, ...]:
    if value is None:
        return ()
    if isinstance(value, str):
        normalized = value.strip()
        return (normalized,) if normalized else ()
    if not isinstance(value, list):
        return ()

    items: list[str] = []
    seen: set[str] = set()
    for entry in value:
        if not isinstance(entry, str):
            continue
        normalized = entry.strip()
        if not normalized or normalized in seen:
            continue
        seen.add(normalized)
        items.append(normalized)
    return tuple(items)


def _tool_schema(name: str) -> dict[str, Any]:
    if name == "complete_task":
        return {
            "type": "object",
            "properties": {
                "result": {"type": "string"},
                "changed_files": {
                    "type": "array",
                    "items": {"type": "string"},
                },
                "outcome_summary": {"type": "string"},
            },
            "required": ["result"],
            "additionalProperties": False,
        }
    if name == "report_blocker":
        return {
            "type": "object",
            "properties": {
                "description": {"type": "string"},
            },
            "required": ["description"],
            "additionalProperties": False,
        }
    if name == "hand_off_to":
        return {
            "type": "object",
            "properties": {
                "agent_name": {"type": "string"},
                "context": {"type": "object"},
            },
            "required": ["agent_name"],
            "additionalProperties": False,
        }
    if name == "request_review":
        return {
            "type": "object",
            "properties": {
                "files": {
                    "type": "array",
                    "items": {"type": "string"},
                },
                "summary": {"type": "string"},
            },
            "required": ["files"],
            "additionalProperties": False,
        }
    raise KeyError(name)


def _tool_description(name: str) -> str:
    return {
        "complete_task": "Record the final result for this stage.",
        "report_blocker": "Report that the task is blocked and needs human input.",
        "hand_off_to": "Hand work to another subagent with structured context.",
        "request_review": "Request a review pass for the given files.",
    }[name]


def _tool_to_mcp(name: str) -> mcp_types.Tool:
    return mcp_types.Tool(
        name=name,
        description=_tool_description(name),
        inputSchema=_tool_schema(name),
    )


class _SignalState:
    def __init__(self) -> None:
        self._completion: ACPTaskCompletion | None = None
        self._blocker: str | None = None
        self._handoff: ACPHandoffRequest | None = None
        self._review: ACPReviewRequest | None = None

    def snapshot(self) -> ACPOrchestrationSnapshot:
        return ACPOrchestrationSnapshot(
            completion=self._completion,
            blocker=self._blocker,
            handoff=self._handoff,
            review=self._review,
        )

    def call(self, name: str, arguments: dict[str, Any]) -> str:
        if name == "complete_task":
            result = _coerce_text(arguments.get("result"))
            if result is None:
                raise ValueError("result is required")
            self._completion = ACPTaskCompletion(
                result=result,
                changed_files=_coerce_string_list(arguments.get("changed_files")),
                outcome_summary=_coerce_text(arguments.get("outcome_summary")),
            )
            return "Recorded completion."

        if name == "report_blocker":
            description = _coerce_text(arguments.get("description"))
            if description is None:
                raise ValueError("description is required")
            self._blocker = description
            return "Recorded blocker."

        if name == "hand_off_to":
            agent_name = _coerce_text(arguments.get("agent_name"))
            if agent_name is None:
                raise ValueError("agent_name is required")
            self._handoff = ACPHandoffRequest(
                agent_name=agent_name,
                context=_coerce_context(arguments.get("context")),
            )
            self._review = None
            return "Recorded handoff."

        if name == "request_review":
            files = _coerce_string_list(arguments.get("files"))
            if not files:
                raise ValueError("files is required")
            self._review = ACPReviewRequest(
                files=files,
                summary=_coerce_text(arguments.get("summary")),
            )
            self._handoff = None
            return "Recorded review request."

        raise ValueError(f"Unknown orchestration tool: {name}")


class _SignalIPCServer:
    def __init__(
        self, socket_path: str, allowed_tools: tuple[str, ...], state: _SignalState
    ) -> None:
        self._socket_path = socket_path
        self._allowed_tools = allowed_tools
        self._state = state
        self._server: asyncio.AbstractServer | None = None

    async def start(self) -> None:
        try:
            os.unlink(self._socket_path)
        except FileNotFoundError:
            pass
        self._server = await asyncio.start_unix_server(
            self._handle_connection, path=self._socket_path
        )
        _log.info("acp orchestration ipc started", socket=self._socket_path)

    async def stop(self) -> None:
        if self._server is not None:
            self._server.close()
            await self._server.wait_closed()
            self._server = None
        try:
            os.unlink(self._socket_path)
        except FileNotFoundError:
            pass

    async def _handle_connection(
        self, reader: asyncio.StreamReader, writer: asyncio.StreamWriter
    ) -> None:
        try:
            while True:
                line = await reader.readline()
                if not line:
                    break
                request = json.loads(line.decode())
                method = request.get("method")
                if method == "list_tools":
                    response = {"tools": list(self._allowed_tools)}
                elif method == "call_tool":
                    name = str(request.get("name", ""))
                    if name not in self._allowed_tools:
                        response = {"error": f"Unknown orchestration tool: {name}"}
                    else:
                        try:
                            result = self._state.call(name, request.get("arguments", {}) or {})
                            response = {"result": result}
                        except Exception as exc:
                            response = {"error": str(exc)}
                else:
                    response = {"error": f"Unknown IPC method: {method}"}
                writer.write(json.dumps(response).encode() + b"\n")
                await writer.drain()
        except (ConnectionResetError, asyncio.IncompleteReadError):
            pass
        finally:
            writer.close()
            await writer.wait_closed()


class _SignalIPCClient:
    def __init__(self, socket_path: str) -> None:
        self._socket_path = socket_path
        self._reader: asyncio.StreamReader | None = None
        self._writer: asyncio.StreamWriter | None = None

    async def connect(self) -> None:
        self._reader, self._writer = await asyncio.open_unix_connection(self._socket_path)

    async def close(self) -> None:
        if self._writer is not None:
            self._writer.close()
            await self._writer.wait_closed()

    async def list_tools(self) -> tuple[str, ...]:
        response = await self._request({"method": "list_tools"})
        return tuple(
            name
            for name in response.get("tools", [])
            if isinstance(name, str) and name in _KNOWN_TOOLS
        )

    async def call_tool(self, name: str, arguments: dict[str, Any]) -> str:
        response = await self._request(
            {"method": "call_tool", "name": name, "arguments": arguments}
        )
        if "error" in response:
            raise RuntimeError(response["error"])
        return str(response.get("result", ""))

    async def _request(self, payload: dict[str, Any]) -> dict[str, Any]:
        assert self._reader is not None and self._writer is not None
        self._writer.write(json.dumps(payload).encode() + b"\n")
        await self._writer.drain()
        line = await self._reader.readline()
        if not line:
            raise ConnectionError("IPC server closed connection")
        return json.loads(line.decode())


class ACPOrchestrationBridge:
    """Expose per-session orchestration helpers to an ACP agent."""

    def __init__(self, *, allowed_tools: list[str] | tuple[str, ...] | None = None) -> None:
        raw_tools = tuple(_KNOWN_TOOLS) if allowed_tools is None else tuple(allowed_tools)
        self._allowed_tools = tuple(
            name for name in raw_tools if isinstance(name, str) and name in _KNOWN_TOOLS
        )
        self._state = _SignalState()
        self._server = mcp.server.Server(_SERVER_NAME)
        self._socket_path = f"/tmp/otto-acp-orch-{uuid.uuid4().hex[:12]}.sock"
        self._ipc_server: _SignalIPCServer | None = None
        self._register_handlers()

    @property
    def tool_names(self) -> list[str]:
        return sorted(self._allowed_tools)

    async def start(self) -> None:
        if self._ipc_server is not None:
            return
        self._ipc_server = _SignalIPCServer(self._socket_path, self._allowed_tools, self._state)
        await self._ipc_server.start()

    async def stop(self) -> None:
        if self._ipc_server is not None:
            await self._ipc_server.stop()
            self._ipc_server = None

    def snapshot(self) -> ACPOrchestrationSnapshot:
        return self._state.snapshot()

    def get_server_config(self) -> dict[str, Any]:
        return {
            "name": _SERVER_NAME,
            "command": sys.executable,
            "args": ["-m", "otto.acp.orchestration_tools"],
            "env": [{"name": _SOCKET_ENV_VAR, "value": self._socket_path}],
        }

    def _register_handlers(self) -> None:
        @self._server.list_tools()
        async def list_tools() -> list[mcp_types.Tool]:
            return [_tool_to_mcp(name) for name in self._allowed_tools]

        @self._server.call_tool()
        async def call_tool(
            name: str, arguments: dict[str, Any] | None = None
        ) -> list[mcp_types.TextContent]:
            if name not in self._allowed_tools:
                return [mcp_types.TextContent(type="text", text=f"Unknown tool: {name}")]
            try:
                result = self._state.call(name, arguments or {})
            except Exception as exc:
                return [mcp_types.TextContent(type="text", text=f"Tool error: {exc}")]
            return [mcp_types.TextContent(type="text", text=result)]

    async def run(self) -> None:
        options = self._server.create_initialization_options()
        async with stdio_server() as (read_stream, write_stream):
            await self._server.run(read_stream, write_stream, options)


async def _run_ipc_bridge(socket_path: str) -> None:
    client = _SignalIPCClient(socket_path)
    await client.connect()

    tool_names = await client.list_tools()
    server = mcp.server.Server(_SERVER_NAME)

    @server.list_tools()
    async def list_tools() -> list[mcp_types.Tool]:
        return [_tool_to_mcp(name) for name in tool_names]

    @server.call_tool()
    async def call_tool(
        name: str, arguments: dict[str, Any] | None = None
    ) -> list[mcp_types.TextContent]:
        try:
            result = await client.call_tool(name, arguments or {})
            return [mcp_types.TextContent(type="text", text=result)]
        except Exception as exc:
            return [mcp_types.TextContent(type="text", text=f"Tool error: {exc}")]

    options = server.create_initialization_options()
    try:
        async with stdio_server() as (read_stream, write_stream):
            await server.run(read_stream, write_stream, options)
    finally:
        await client.close()


def _run_standalone() -> None:
    socket_path = os.environ.get(_SOCKET_ENV_VAR)
    if socket_path:
        asyncio.run(_run_ipc_bridge(socket_path))
        return

    bridge = ACPOrchestrationBridge()
    asyncio.run(bridge.run())


if __name__ == "__main__":
    _run_standalone()
