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

from __future__ import annotations

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

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] = {}

    # ------------------------------------------------------------------
    # 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()
                self._record_directive_context(msg_type=msg_type, payload=payload, message_id=message_id)
                await self._send_acknowledgement(message_id=message_id, payload=payload)
                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,
                )

            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]) -> None:
        if not message_id:
            return
        envelope = {
            "type": "acknowledgement",
            "agent_id": self.config.agent_id,
            "content": "Acknowledged. Working on it.",
            "metadata": {
                "message_id": os.urandom(8).hex(),
                "message_kind": "acknowledgement",
                "round_id": payload.get("round_id"),
                "subtask_id": payload.get("subtask_id"),
                "directive_id": payload.get("directive_id") or message_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"

    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)

    async def _handle_hook_connection(
        self,
        reader: asyncio.StreamReader,
        writer: asyncio.StreamWriter,
    ) -> None:
        try:
            while True:
                raw = await reader.readline()
                if not raw:
                    break
                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,
        }

    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
