"""Socket-based messaging between agents."""

import asyncio
import json
import logging
import os
import time
from typing import Callable, Dict, List, Optional

from .models import HubMessage
from .presence import _atomic_write, get_messages_dir, get_socket_dir

logger = logging.getLogger(__name__)


class AgentSocketServer:
    """Per-agent unix domain socket server for receiving messages.

    Each agent runs one of these. Other agents connect to deliver
    messages directly (peer-to-peer, no hub in the middle).
    """

    def __init__(
        self,
        agent_id: str,
        on_message: Callable,
        on_get_output: Optional[Callable] = None,
        on_shutdown: Optional[Callable] = None,
    ):
        self.agent_id = agent_id
        self.socket_path = get_socket_dir() / f"{agent_id}.sock"
        self._server: Optional[asyncio.AbstractServer] = None
        self._on_message = on_message
        self._on_get_output = on_get_output
        self._on_shutdown = on_shutdown
        self._identity: Optional[Dict] = None
        self._started_at: float = time.time()
        self._shutdown_requested = False

    async def start(self) -> str:
        """Start the socket server. Returns the socket path."""
        # Clean stale socket
        if self.socket_path.exists():
            self.socket_path.unlink()

        self._server = await asyncio.start_unix_server(
            self._handle_connection, path=str(self.socket_path)
        )
        os.chmod(str(self.socket_path), 0o600)
        logger.info(f"Agent socket listening: {self.socket_path}")
        return str(self.socket_path)

    async def _handle_connection(
        self, reader: asyncio.StreamReader, writer: asyncio.StreamWriter
    ):
        """Handle incoming connection."""
        try:
            while True:
                line = await reader.readline()
                if not line:
                    break

                try:
                    msg_data = json.loads(line.decode().strip())
                except (json.JSONDecodeError, UnicodeDecodeError):
                    continue

                action = msg_data.get("action", "")

                if action == "message":
                    msg = HubMessage.from_dict(msg_data)
                    await self._on_message(msg)
                    ack = json.dumps({"type": "ack", "id": msg.id}) + "\n"
                    writer.write(ack.encode())
                    await writer.drain()

                elif action == "ping":
                    pong = (
                        json.dumps(
                            {
                                "type": "pong",
                                "agent_id": self.agent_id,
                            }
                        )
                        + "\n"
                    )
                    writer.write(pong.encode())
                    await writer.drain()

                elif action == "get_context":
                    # Return recent conversation context for social awareness
                    lines_requested = msg_data.get("lines", 200)
                    context = await self._get_context(lines_requested)
                    resp = (
                        json.dumps(
                            {
                                "type": "context",
                                "content": context,
                            }
                        )
                        + "\n"
                    )
                    writer.write(resp.encode())
                    await writer.drain()

                elif action == "roster_update":
                    # Hub pushing updated roster
                    await self._on_message(
                        HubMessage(
                            action="roster_update",
                            content=json.dumps(msg_data.get("agents", [])),
                            from_designation="hub",
                        )
                    )
                    ack = json.dumps({"type": "ack"}) + "\n"
                    writer.write(ack.encode())
                    await writer.drain()

                elif action == "get_output":
                    lines_requested = msg_data.get("lines", 100)
                    output_lines: List[str] = []
                    if self._on_get_output:
                        try:
                            result = self._on_get_output(lines_requested)
                            if asyncio.iscoroutine(result):
                                result = await result
                            if isinstance(result, list):
                                output_lines = result
                        except Exception as exc:
                            logger.debug(f"get_output callback error: {exc}")
                    resp = json.dumps({"type": "output", "lines": output_lines}) + "\n"
                    writer.write(resp.encode())
                    await writer.drain()

                elif action == "get_frames":
                    # Return recent display frames (currently backed by ring buffer)
                    lines_requested = msg_data.get("lines", 100)
                    frame_lines: List[str] = []
                    if self._on_get_output:
                        try:
                            result = self._on_get_output(lines_requested)
                            if asyncio.iscoroutine(result):
                                result = await result
                            if isinstance(result, list):
                                frame_lines = result
                        except Exception as exc:
                            logger.debug(f"get_frames callback error: {exc}")
                    resp = (
                        json.dumps({"type": "frames", "lines": frame_lines}) + "\n"
                    )
                    writer.write(resp.encode())
                    await writer.drain()

                elif action == "subscribe":
                    # Acknowledge subscription request (live streaming future work)
                    resp = (
                        json.dumps(
                            {
                                "type": "subscribed",
                                "status": "subscribed",
                                "note": "live streaming not yet implemented",
                            }
                        )
                        + "\n"
                    )
                    writer.write(resp.encode())
                    await writer.drain()

                elif action == "get_status":
                    status_data = {
                        "type": "status",
                        "designation": "",
                        "state": "unknown",
                        "pid": os.getpid(),
                        "uptime": int(time.time() - self._started_at),
                        "current_task": "",
                    }
                    if self._identity:
                        status_data["designation"] = self._identity.get(
                            "designation", ""
                        )
                        status_data["state"] = self._identity.get("state", "unknown")
                        status_data["pid"] = self._identity.get("pid", os.getpid())
                        status_data["current_task"] = self._identity.get(
                            "current_task", ""
                        )
                    resp = json.dumps(status_data) + "\n"
                    writer.write(resp.encode())
                    await writer.drain()

                elif action == "shutdown":
                    reason = msg_data.get("reason", "")
                    logger.info(f"Shutdown requested: {reason or 'no reason given'}")
                    ack = json.dumps({"type": "ack"}) + "\n"
                    writer.write(ack.encode())
                    await writer.drain()
                    self._shutdown_requested = True
                    if self._on_shutdown:
                        try:
                            result = self._on_shutdown(reason)
                            if asyncio.iscoroutine(result):
                                await result
                        except Exception as exc:
                            logger.debug(f"shutdown callback error: {exc}")

        except Exception as e:
            logger.debug(f"Connection handler error: {e}")
        finally:
            writer.close()
            try:
                await writer.wait_closed()
            except Exception:
                pass

    async def _get_context(self, lines: int) -> str:
        """Get recent context - override in plugin integration."""
        return ""

    def set_identity(self, identity_dict: Dict) -> None:
        """Set the identity dict used for get_status responses."""
        self._identity = identity_dict

    @property
    def shutdown_requested(self) -> bool:
        """Whether a remote shutdown has been requested."""
        return self._shutdown_requested

    async def stop(self) -> None:
        """Stop the socket server and clean up."""
        if self._server:
            self._server.close()
            await self._server.wait_closed()
            self._server = None
        if self.socket_path.exists():
            try:
                self.socket_path.unlink()
            except OSError:
                pass
        logger.info("Agent socket stopped")


class AgentMessenger:
    """Send messages to other agents via their sockets."""

    @staticmethod
    async def send_to_agent(
        target_socket: str,
        message: HubMessage,
        timeout: float = 5.0,
    ) -> bool:
        """Send a message to an agent's socket.

        Returns True if delivered and acked.
        """
        try:
            reader, writer = await asyncio.wait_for(
                asyncio.open_unix_connection(target_socket),
                timeout=timeout,
            )

            msg_line = json.dumps(message.to_dict()) + "\n"
            writer.write(msg_line.encode())
            await writer.drain()

            # Wait for ack
            ack_line = await asyncio.wait_for(reader.readline(), timeout=timeout)
            writer.close()
            await writer.wait_closed()

            if ack_line:
                ack = json.loads(ack_line.decode().strip())
                if isinstance(ack, dict):
                    return ack.get("type") == "ack"
            return False

        except Exception as e:
            logger.debug(f"Send to {target_socket} failed: {e}")
            return False

    @staticmethod
    async def ping_agent(socket_path: str, timeout: float = 3.0) -> bool:
        """Ping an agent to check if it's alive."""
        try:
            reader, writer = await asyncio.wait_for(
                asyncio.open_unix_connection(socket_path),
                timeout=timeout,
            )
            ping = json.dumps({"action": "ping"}) + "\n"
            writer.write(ping.encode())
            await writer.drain()

            resp_line = await asyncio.wait_for(reader.readline(), timeout=timeout)
            writer.close()
            await writer.wait_closed()

            if resp_line:
                resp = json.loads(resp_line.decode().strip())
                if isinstance(resp, dict):
                    return resp.get("type") == "pong"
            return False
        except Exception:
            return False

    @staticmethod
    async def request_context(
        socket_path: str, lines: int = 200, timeout: float = 5.0
    ) -> str:
        """Request recent context from an agent."""
        try:
            reader, writer = await asyncio.wait_for(
                asyncio.open_unix_connection(socket_path),
                timeout=timeout,
            )
            req = json.dumps({"action": "get_context", "lines": lines}) + "\n"
            writer.write(req.encode())
            await writer.drain()

            resp_line = await asyncio.wait_for(reader.readline(), timeout=timeout)
            writer.close()
            await writer.wait_closed()

            if resp_line:
                resp = json.loads(resp_line.decode().strip())
                if isinstance(resp, dict):
                    content = resp.get("content", "")
                    return content if isinstance(content, str) else ""
            return ""
        except Exception:
            return ""

    @staticmethod
    async def request_output(
        socket_path: str, lines: int = 100, timeout: float = 5.0
    ) -> List[str]:
        """Request recent output lines from an agent."""
        try:
            reader, writer = await asyncio.wait_for(
                asyncio.open_unix_connection(socket_path),
                timeout=timeout,
            )
            req = json.dumps({"action": "get_output", "lines": lines}) + "\n"
            writer.write(req.encode())
            await writer.drain()

            resp_line = await asyncio.wait_for(reader.readline(), timeout=timeout)
            writer.close()
            await writer.wait_closed()

            if resp_line:
                resp = json.loads(resp_line.decode().strip())
                if isinstance(resp, dict) and resp.get("type") == "output":
                    result = resp.get("lines", [])
                    return result if isinstance(result, list) else []
            return []
        except Exception:
            return []

    @staticmethod
    async def request_status(socket_path: str, timeout: float = 3.0) -> dict:
        """Request status from an agent."""
        try:
            reader, writer = await asyncio.wait_for(
                asyncio.open_unix_connection(socket_path),
                timeout=timeout,
            )
            req = json.dumps({"action": "get_status"}) + "\n"
            writer.write(req.encode())
            await writer.drain()

            resp_line = await asyncio.wait_for(reader.readline(), timeout=timeout)
            writer.close()
            await writer.wait_closed()

            if resp_line:
                resp = json.loads(resp_line.decode().strip())
                if isinstance(resp, dict) and resp.get("type") == "status":
                    return resp
            return {}
        except Exception:
            return {}

    @staticmethod
    async def signal_shutdown(
        socket_path: str, reason: str = "", timeout: float = 3.0
    ) -> bool:
        """Signal an agent to shut down gracefully. Returns True if acked."""
        try:
            reader, writer = await asyncio.wait_for(
                asyncio.open_unix_connection(socket_path),
                timeout=timeout,
            )
            req = json.dumps({"action": "shutdown", "reason": reason}) + "\n"
            writer.write(req.encode())
            await writer.drain()

            resp_line = await asyncio.wait_for(reader.readline(), timeout=timeout)
            writer.close()
            await writer.wait_closed()

            if resp_line:
                resp = json.loads(resp_line.decode().strip())
                if isinstance(resp, dict):
                    return resp.get("type") == "ack"
            return False
        except Exception:
            return False

    @staticmethod
    async def request_frames(
        socket_path: str, lines: int = 100, timeout: float = 5.0
    ) -> List[str]:
        """Request recent display frames from an agent."""
        try:
            reader, writer = await asyncio.wait_for(
                asyncio.open_unix_connection(socket_path),
                timeout=timeout,
            )
            req = json.dumps({"action": "get_frames", "lines": lines}) + "\n"
            writer.write(req.encode())
            await writer.drain()

            resp_line = await asyncio.wait_for(reader.readline(), timeout=timeout)
            writer.close()
            await writer.wait_closed()

            if resp_line:
                resp = json.loads(resp_line.decode().strip())
                if isinstance(resp, dict) and resp.get("type") == "frames":
                    result = resp.get("lines", [])
                    return result if isinstance(result, list) else []
            return []
        except Exception:
            return []

    @staticmethod
    async def subscribe(
        socket_path: str, timeout: float = 5.0
    ) -> dict:
        """Subscribe to an agent's output stream. Returns ack dict."""
        try:
            reader, writer = await asyncio.wait_for(
                asyncio.open_unix_connection(socket_path),
                timeout=timeout,
            )
            req = json.dumps({"action": "subscribe"}) + "\n"
            writer.write(req.encode())
            await writer.drain()

            resp_line = await asyncio.wait_for(reader.readline(), timeout=timeout)
            writer.close()
            await writer.wait_closed()

            if resp_line:
                resp = json.loads(resp_line.decode().strip())
                if isinstance(resp, dict):
                    return resp
            return {}
        except Exception:
            return {}

    @staticmethod
    async def send_to_file(target_agent_id: str, message: HubMessage) -> None:
        """Fallback: write message to agent's filesystem mailbox."""
        msg_dir = get_messages_dir() / target_agent_id
        msg_dir.mkdir(parents=True, exist_ok=True)

        # Sequential naming for ordering
        existing = sorted(msg_dir.glob("*.json"))
        seq = len(existing) + 1
        filename = f"{seq:04d}-from-{message.from_designation}.json"

        _atomic_write(msg_dir / filename, message.to_dict())
        logger.info(f"Wrote message to mailbox: {filename}")

    @staticmethod
    def read_mailbox(agent_id: str) -> List[HubMessage]:
        """Read and consume messages from filesystem mailbox."""
        msg_dir = get_messages_dir() / agent_id
        if not msg_dir.exists():
            return []

        messages = []
        for f in sorted(msg_dir.glob("*.json")):
            try:
                with open(f) as fh:
                    data = json.load(fh)
                messages.append(HubMessage.from_dict(data))
                f.unlink()  # Consume after reading
            except Exception as e:
                logger.debug(f"Bad mailbox message {f}: {e}")
        return messages
