"""Madpod control-plane WebSocket channel for madpod-agent."""

from __future__ import annotations

import asyncio
import collections
import json
import os
from pathlib import Path
from typing import Any

_MAX_TRACKED_DIRECTIVE_IDS = 10_000

import websockets
import websockets.connection
from websockets.asyncio.client import ClientConnection
from loguru import logger

from madpod_agent.bus.events import OutboundMessage
from madpod_agent.channels.base import BaseChannel


class MadpodChannel(BaseChannel):
    """
    Connects the madpod-agent to the Madpod control plane via a persistent WebSocket.

    Inbound (control plane → agent):
        assignment/correction/directive — task instruction, forwarded to AgentLoop via MessageBus
        ping       — keep-alive probe, reply with pong
        suspend    — write sentinel file to pause agent
        resume     — delete sentinel file to resume agent

    Outbound (agent → control plane):
        response   — agent reply to a directive
        heartbeat  — sent every 30 s while connected
    """

    name = "madpod"

    def __init__(self, config: Any, bus: Any) -> None:
        super().__init__(config, bus)
        self._ws: ClientConnection | None = None
        self._recv_task: asyncio.Task | None = None
        self._hb_task: asyncio.Task | None = None
        self._send_queue: asyncio.Queue[str] = asyncio.Queue()
        self._sentinel = Path(f"/tmp/madpod-{self.config.agent_id}.suspend")
        self._hook_server: asyncio.AbstractServer | None = None
        self._hook_socket = Path(
            os.environ.get("MADPOD_CHANNEL_SOCKET", f"/tmp/madpod-{self.config.agent_id}.sock")
        )
        self._current_directive_context: dict[str, Any] = {}
        self._accepted_directive_ids: collections.OrderedDict[str, None] = collections.OrderedDict()
        # Lost-in-the-middle fix: objectives extracted from directives
        self._current_task_summary: str | None = None
        self._current_round_goal: str | None = None
        self._current_remaining_items: list[str] | None = None

    # ------------------------------------------------------------------
    # Current objectives (lost-in-the-middle fix)
    # ------------------------------------------------------------------

    @property
    def current_task_summary(self) -> str | None:
        return self._current_task_summary

    @property
    def current_round_goal(self) -> str | None:
        return self._current_round_goal

    @property
    def current_remaining_items(self) -> list[str] | None:
        return self._current_remaining_items

    # ------------------------------------------------------------------
    # BaseChannel interface
    # ------------------------------------------------------------------

    async def start(self) -> None:
        self._running = True
        logger.info("MadpodChannel: starting (agent_id={})", self.config.agent_id)
        await self._start_hook_bridge()
        try:
            await self._connect_loop()
        finally:
            await self._stop_hook_bridge()

    async def stop(self) -> None:
        self._running = False
        self._cancel_tasks()
        await self._stop_hook_bridge()
        if self._ws and self._ws.state != websockets.connection.State.CLOSED:
            await self._ws.close()
        logger.info("MadpodChannel: stopped")

    async def send(self, msg: OutboundMessage) -> None:
        """Queue an outbound agent response to the control plane."""
        metadata = dict(msg.metadata or {})
        message_kind = str(metadata.get("message_kind") or "").strip().lower()
        outbound_type = message_kind if message_kind in {"acknowledgement", "activity", "deliverable", "intervention"} else "response"
        if outbound_type == "activity":
            payload = json.dumps({
                "type": "activity",
                "agent_id": self.config.agent_id,
                "content": msg.content,
                "summary": str(metadata.get("summary") or msg.content),
                "action": str(metadata.get("action") or "status_update"),
                "payload": metadata,
            })
        else:
            payload = json.dumps({
                "type": outbound_type,
                "agent_id": self.config.agent_id,
                "content": msg.content,
                "chat_id": msg.chat_id,
                "metadata": metadata,
            })
        await self._send_queue.put(payload)

    # ------------------------------------------------------------------
    # Connection loop with exponential backoff
    # ------------------------------------------------------------------

    async def _connect_loop(self) -> None:
        backoff = 1
        backoff_cap = 30

        while self._running:
            try:
                url = self._ws_url()
                headers = {"Authorization": f"Bearer {self.config.gateway_token}"}
                logger.info("MadpodChannel: connecting to {}", url)

                async with websockets.connect(url, additional_headers=headers) as ws:
                    self._ws = ws
                    backoff = 1  # reset on successful connect
                    logger.info("MadpodChannel: connected")

                    self._recv_task = asyncio.create_task(self._receive_loop(ws))
                    self._hb_task = asyncio.create_task(self._heartbeat_loop(ws))
                    send_task = asyncio.create_task(self._send_loop(ws))

                    done, pending = await asyncio.wait(
                        [self._recv_task, self._hb_task, send_task],
                        return_when=asyncio.FIRST_COMPLETED,
                    )
                    for t in pending:
                        t.cancel()
                    # propagate exception if any
                    for t in done:
                        if not t.cancelled() and t.exception():
                            logger.warning("MadpodChannel task error: {}", t.exception())

            except (websockets.ConnectionClosed, OSError, ConnectionRefusedError) as e:
                logger.warning("MadpodChannel: disconnected ({}), retrying in {}s", e, backoff)
            except Exception as e:
                logger.error("MadpodChannel: unexpected error: {}", e)

            if not self._running:
                break

            await asyncio.sleep(backoff)
            backoff = min(backoff * 2, backoff_cap)

    # ------------------------------------------------------------------
    # Receive loop
    # ------------------------------------------------------------------

    async def _receive_loop(self, ws: ClientConnection) -> None:
        async for raw in ws:
            try:
                msg = json.loads(raw)
            except json.JSONDecodeError:
                logger.warning("MadpodChannel: invalid JSON received")
                continue

            msg_type = msg.get("type")

            if msg_type in {"directive", "assignment", "correction"}:
                task_content = msg.get("task") or msg.get("content", "")
                sender = msg.get("sender_id", "orchestrator")
                payload = msg.get("payload") if isinstance(msg.get("payload"), dict) else {}
                message_id = str(payload.get("message_id") or msg.get("id") or "").strip()
                directive_id = str(payload.get("directive_id") or message_id).strip()
                self._record_directive_context(msg_type=msg_type, payload=payload, message_id=message_id)
                if directive_id and directive_id in self._accepted_directive_ids:
                    logger.info("MadpodChannel: duplicate {} ignored (directive_id={})", msg_type, directive_id)
                    await self._send_acknowledgement(
                        message_id=message_id,
                        payload=payload,
                        directive_id=directive_id,
                    )
                    continue
                metadata = dict(payload)
                if message_id and not str(metadata.get("message_id") or "").strip():
                    metadata["message_id"] = message_id
                if message_id and not str(metadata.get("in_reply_to") or "").strip():
                    metadata["in_reply_to"] = message_id
                if message_id and not str(metadata.get("assignment_message_id") or "").strip():
                    metadata["assignment_message_id"] = message_id
                if message_id and not str(metadata.get("directive_id") or "").strip():
                    metadata["directive_id"] = message_id
                metadata["directive_type"] = msg_type
                logger.info("MadpodChannel: {} received", msg_type)
                await self._handle_message(
                    sender_id=sender,
                    chat_id=self.config.agent_id,
                    content=task_content,
                    metadata=metadata,
                    session_key=self._directive_session_key(payload=metadata, message_id=message_id),
                )
                if directive_id:
                    self._accepted_directive_ids[directive_id] = None
                    if len(self._accepted_directive_ids) > _MAX_TRACKED_DIRECTIVE_IDS:
                        self._accepted_directive_ids.popitem(last=False)
                await self._send_acknowledgement(
                    message_id=message_id,
                    payload=payload,
                    directive_id=directive_id,
                )

            elif msg_type == "ping":
                await self._send_queue.put(json.dumps({
                    "type": "pong",
                    "agent_id": self.config.agent_id,
                }))

            elif msg_type == "pong":
                continue

            elif msg_type == "suspend":
                logger.info("MadpodChannel: suspend received — writing sentinel")
                self._sentinel.touch()

            elif msg_type == "resume":
                logger.info("MadpodChannel: resume received — removing sentinel")
                self._sentinel.unlink(missing_ok=True)

            else:
                logger.debug("MadpodChannel: unhandled message type '{}'", msg_type)

    async def _send_acknowledgement(
        self,
        *,
        message_id: str,
        payload: dict[str, Any],
        directive_id: str | None = None,
    ) -> None:
        if not message_id:
            return
        resolved_directive_id = str(directive_id or payload.get("directive_id") or message_id).strip() or message_id
        envelope = {
            "type": "acknowledgement",
            "agent_id": self.config.agent_id,
            "content": "Acknowledged. Working on it.",
            "metadata": {
                "message_id": f"ack-{resolved_directive_id}",
                "message_kind": "acknowledgement",
                "round_id": payload.get("round_id"),
                "subtask_id": payload.get("subtask_id"),
                "directive_id": resolved_directive_id,
                "in_reply_to": message_id,
            },
        }
        await self._send_queue.put(json.dumps(envelope))

    # ------------------------------------------------------------------
    # Send loop (drains queue → WebSocket)
    # ------------------------------------------------------------------

    async def _send_loop(self, ws: ClientConnection) -> None:
        while True:
            payload = await self._send_queue.get()
            try:
                await ws.send(payload)
            except websockets.ConnectionClosed:
                # Put it back and let _connect_loop reconnect
                await self._send_queue.put(payload)
                raise

    # ------------------------------------------------------------------
    # Heartbeat loop
    # ------------------------------------------------------------------

    async def _heartbeat_loop(self, ws: ClientConnection) -> None:
        while True:
            await asyncio.sleep(30)
            await self._send_queue.put(json.dumps({
                "type": "heartbeat",
                "agent_id": self.config.agent_id,
                "status": "running",
            }))

    # ------------------------------------------------------------------
    # Helpers
    # ------------------------------------------------------------------

    def _ws_url(self) -> str:
        base = self.config.control_plane_url.rstrip("/")
        return f"{base}/ws/agents/{self.config.agent_id}/channel"

    def _directive_session_key(
        self,
        *,
        payload: dict[str, Any],
        message_id: str,
    ) -> str | None:
        directive_id = str(payload.get("directive_id") or message_id or "").strip()
        if directive_id:
            return f"{self.name}:{self.config.agent_id}:{directive_id}"
        round_id = str(payload.get("round_id") or "").strip()
        if round_id:
            return f"{self.name}:{self.config.agent_id}:round:{round_id}"
        subtask_id = str(payload.get("subtask_id") or "").strip()
        if subtask_id:
            return f"{self.name}:{self.config.agent_id}:subtask:{subtask_id}"
        return None

    async def _start_hook_bridge(self) -> None:
        if self._hook_server is not None:
            return
        self._hook_socket.parent.mkdir(parents=True, exist_ok=True)
        self._hook_socket.unlink(missing_ok=True)
        self._hook_server = await asyncio.start_unix_server(
            self._handle_hook_connection,
            path=str(self._hook_socket),
        )
        logger.info("MadpodChannel: hook bridge listening on {}", self._hook_socket)

    async def _stop_hook_bridge(self) -> None:
        if self._hook_server is not None:
            self._hook_server.close()
            await self._hook_server.wait_closed()
            self._hook_server = None
        self._hook_socket.unlink(missing_ok=True)

    _HOOK_LINE_MAX_BYTES = 1_048_576  # 1 MB per JSON line

    async def _handle_hook_connection(
        self,
        reader: asyncio.StreamReader,
        writer: asyncio.StreamWriter,
    ) -> None:
        try:
            while True:
                raw = await reader.readline()
                if not raw:
                    break
                if len(raw) > self._HOOK_LINE_MAX_BYTES:
                    logger.warning("MadpodChannel: hook payload too large (%d bytes), skipping", len(raw))
                    continue
                try:
                    payload = json.loads(raw.decode("utf-8"))
                except Exception:
                    logger.warning("MadpodChannel: invalid hook payload")
                    continue
                await self._queue_hook_event(payload)
        finally:
            writer.close()
            try:
                await writer.wait_closed()
            except Exception:
                pass

    async def _queue_hook_event(self, event: dict[str, Any]) -> None:
        enriched = self._enrich_hook_event(event)
        await self._send_queue.put(json.dumps(enriched))

    def _record_directive_context(
        self,
        *,
        msg_type: str,
        payload: dict[str, Any],
        message_id: str,
    ) -> None:
        upstream_in_reply_to = str(payload.get("in_reply_to") or "").strip() or None
        self._current_directive_context = {
            "directive_type": msg_type,
            "assignment_message_id": message_id or None,
            "directive_id": str(payload.get("directive_id") or message_id or "").strip() or None,
            "round_id": str(payload.get("round_id") or "").strip() or None,
            "subtask_id": str(payload.get("subtask_id") or "").strip() or None,
            "upstream_in_reply_to": upstream_in_reply_to,
        }
        # Extract objectives for lost-in-the-middle injection
        self._current_task_summary = str(payload.get("task_summary") or payload.get("task_brief") or "").strip() or None
        self._current_round_goal = str(payload.get("round_goal") or payload.get("goal") or "").strip() or None
        raw_remaining = payload.get("remaining_items") or payload.get("expected_output") or payload.get("acceptance_criteria")
        if isinstance(raw_remaining, list):
            self._current_remaining_items = [str(item).strip() for item in raw_remaining if str(item).strip()] or None
        elif isinstance(raw_remaining, str) and raw_remaining.strip():
            self._current_remaining_items = [raw_remaining.strip()]
        else:
            self._current_remaining_items = None

    def _enrich_hook_event(self, event: dict[str, Any]) -> dict[str, Any]:
        enriched = dict(event or {})
        payload = dict(enriched.get("payload") or {}) if isinstance(enriched.get("payload"), dict) else {}
        context = dict(self._current_directive_context or {})

        assignment_message_id = str(context.get("assignment_message_id") or "").strip() or None
        directive_id = str(context.get("directive_id") or assignment_message_id or "").strip() or None
        reply_target = (
            str(enriched.get("in_reply_to") or "").strip()
            or str(payload.get("in_reply_to") or "").strip()
            or assignment_message_id
        )
        for key in ("round_id", "subtask_id"):
            value = (
                str(enriched.get(key) or "").strip()
                or str(payload.get(key) or "").strip()
                or str(context.get(key) or "").strip()
            )
            if value:
                enriched[key] = value
                payload[key] = value
        if reply_target:
            enriched["in_reply_to"] = reply_target
            payload["in_reply_to"] = reply_target
        if directive_id:
            enriched["directive_id"] = directive_id
            payload["directive_id"] = directive_id
        if assignment_message_id:
            payload["assignment_message_id"] = assignment_message_id
        upstream_in_reply_to = str(context.get("upstream_in_reply_to") or "").strip() or None
        if upstream_in_reply_to and upstream_in_reply_to != reply_target:
            payload["origin_in_reply_to"] = upstream_in_reply_to
        enriched["payload"] = payload
        return enriched

    def _cancel_tasks(self) -> None:
        for task in (self._recv_task, self._hb_task):
            if task and not task.done():
                task.cancel()

    def is_allowed(self, sender_id: str) -> bool:
        # Control plane is the only sender; access is secured by gateway token on WS handshake.
        return True
