"""
IPC Server for ESPal daemon.

Handles Unix domain socket connections, client session management,
and message routing between clients and the daemon core.
"""

import os
import socket
import threading
import logging
import time
from typing import Optional, Callable, Any
from dataclasses import dataclass, field

from .buffer import RingBuffer, BufferReader, BufferEntry
from .protocol import (
    serialize, parse_command, parse_json_line, MessageFramer, ProtocolError,
    SubscribeCommand, BuildCommand, UploadCommand, StatusCommand,
    ShutdownCommand, UnsubscribeCommand,
    DataMessage, MarkerMessage, StatusMessage, ErrorMessage,
    SubscribedMessage, ShutdownNotice, CommandType
)


logger = logging.getLogger(__name__)


@dataclass
class ClientSession:
    """Tracks state for a connected client."""
    client_id: int
    socket: socket.socket
    address: str
    connected_at: float
    subscribed: bool = False
    subscribe_mode: Optional[str] = None
    reader: Optional[BufferReader] = None
    lock: threading.Lock = field(default_factory=threading.Lock)
    running: bool = True

    def send(self, msg: Any) -> bool:
        """
        Send a message to this client.

        Args:
            msg: Message object with to_dict() method

        Returns:
            True if sent successfully, False otherwise
        """
        with self.lock:
            if not self.running:
                return False
            try:
                data = serialize(msg)
                self.socket.sendall(data)
                return True
            except (OSError, BrokenPipeError) as e:
                logger.debug(f"Client {self.client_id}: send failed: {e}")
                self.running = False
                return False

    def close(self) -> None:
        """Close the client connection."""
        with self.lock:
            self.running = False
            try:
                self.socket.shutdown(socket.SHUT_RDWR)
            except OSError:
                pass
            try:
                self.socket.close()
            except OSError:
                pass


class IPCServer:
    """
    Unix domain socket server for daemon IPC.

    Accepts client connections, parses commands, and dispatches
    to registered handlers.
    """

    def __init__(
        self,
        socket_path: str,
        buffer: RingBuffer,
        on_client_connect: Optional[Callable[[ClientSession], None]] = None,
        on_client_disconnect: Optional[Callable[[ClientSession], None]] = None
    ):
        """
        Initialise the IPC server.

        Args:
            socket_path: Path for the Unix domain socket
            buffer: Ring buffer for data streaming
            on_client_connect: Callback when client connects
            on_client_disconnect: Callback when client disconnects
        """
        self._socket_path = socket_path
        self._buffer = buffer
        self._on_connect = on_client_connect
        self._on_disconnect = on_client_disconnect

        self._server_socket: Optional[socket.socket] = None
        self._clients: dict[int, ClientSession] = {}
        self._clients_lock = threading.Lock()
        self._next_client_id = 1

        self._running = False
        self._accept_thread: Optional[threading.Thread] = None

        # Command handlers
        self._command_handlers: dict[type, Callable] = {}

        # Streaming threads for subscribed clients
        self._stream_threads: dict[int, threading.Thread] = {}

    @property
    def socket_path(self) -> str:
        """Path to the Unix socket."""
        return self._socket_path

    @property
    def client_count(self) -> int:
        """Number of connected clients."""
        with self._clients_lock:
            return len(self._clients)

    @property
    def is_running(self) -> bool:
        """Whether the server is running."""
        return self._running

    def register_handler(self, cmd_type: type, handler: Callable) -> None:
        """
        Register a command handler.

        Args:
            cmd_type: Command class (e.g., BuildCommand)
            handler: Callable(session, command) -> response or None
        """
        self._command_handlers[cmd_type] = handler

    def start(self) -> None:
        """Start the IPC server."""
        if self._running:
            return

        # Remove stale socket file
        if os.path.exists(self._socket_path):
            os.unlink(self._socket_path)

        # Create socket
        self._server_socket = socket.socket(socket.AF_UNIX, socket.SOCK_STREAM)
        self._server_socket.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
        self._server_socket.bind(self._socket_path)
        self._server_socket.listen(5)

        # Set socket permissions (user only)
        os.chmod(self._socket_path, 0o600)

        self._running = True

        # Start accept thread
        self._accept_thread = threading.Thread(target=self._accept_loop, daemon=True)
        self._accept_thread.start()

        logger.info(f"IPC server started on {self._socket_path}")

    def stop(self) -> None:
        """Stop the IPC server and disconnect all clients."""
        if not self._running:
            return

        self._running = False

        # Close server socket to unblock accept()
        if self._server_socket:
            try:
                self._server_socket.close()
            except OSError:
                pass
            self._server_socket = None

        # Notify and disconnect all clients
        with self._clients_lock:
            for session in list(self._clients.values()):
                try:
                    session.send(ShutdownNotice(reason="server_shutdown"))
                except Exception:
                    pass
                session.close()
            self._clients.clear()

        # Remove socket file
        if os.path.exists(self._socket_path):
            try:
                os.unlink(self._socket_path)
            except OSError:
                pass

        # Wait for accept thread
        if self._accept_thread and self._accept_thread.is_alive():
            self._accept_thread.join(timeout=2.0)

        logger.info("IPC server stopped")

    def broadcast(self, msg: Any, subscribed_only: bool = True) -> int:
        """
        Broadcast a message to clients.

        Args:
            msg: Message object to send
            subscribed_only: If True, only send to subscribed clients

        Returns:
            Number of clients that received the message
        """
        sent = 0
        with self._clients_lock:
            for session in list(self._clients.values()):
                if subscribed_only and not session.subscribed:
                    continue
                if session.send(msg):
                    sent += 1
        return sent

    def get_clients(self) -> list[ClientSession]:
        """Get list of connected clients."""
        with self._clients_lock:
            return list(self._clients.values())

    def _accept_loop(self) -> None:
        """Accept incoming client connections."""
        while self._running:
            try:
                client_socket, _ = self._server_socket.accept()
                self._handle_new_client(client_socket)
            except OSError:
                # Socket closed, server stopping
                break
            except Exception as e:
                if self._running:
                    logger.error(f"Accept error: {e}")

    def _handle_new_client(self, client_socket: socket.socket) -> None:
        """Handle a new client connection."""
        with self._clients_lock:
            client_id = self._next_client_id
            self._next_client_id += 1

        session = ClientSession(
            client_id=client_id,
            socket=client_socket,
            address=f"client-{client_id}",
            connected_at=time.time()
        )

        with self._clients_lock:
            self._clients[client_id] = session

        logger.info(f"Client {client_id} connected")

        if self._on_connect:
            try:
                self._on_connect(session)
            except Exception as e:
                logger.error(f"on_connect callback error: {e}")

        # Start client handler thread
        thread = threading.Thread(
            target=self._client_loop,
            args=(session,),
            daemon=True
        )
        thread.start()

    def _client_loop(self, session: ClientSession) -> None:
        """Handle messages from a client."""
        framer = MessageFramer()
        session.socket.settimeout(1.0)  # Allow periodic check of running flag

        try:
            while session.running and self._running:
                try:
                    data = session.socket.recv(4096)
                    if not data:
                        break  # Client disconnected

                    for line in framer.feed(data):
                        self._handle_message(session, line)

                except socket.timeout:
                    continue
                except OSError:
                    break

        finally:
            self._cleanup_client(session)

    def _handle_message(self, session: ClientSession, line: bytes) -> None:
        """Parse and handle a single message from a client."""
        try:
            data = parse_json_line(line)
            command = parse_command(data)
            self._dispatch_command(session, command)

        except ProtocolError as e:
            logger.warning(f"Client {session.client_id}: protocol error: {e}")
            session.send(ErrorMessage(code="PROTOCOL_ERROR", message=str(e)))

        except Exception as e:
            logger.error(f"Client {session.client_id}: error handling message: {e}")
            session.send(ErrorMessage(code="INTERNAL_ERROR", message=str(e)))

    def _dispatch_command(self, session: ClientSession, command: CommandType) -> None:
        """Dispatch a command to the appropriate handler."""
        cmd_type = type(command)

        # Built-in handlers
        if isinstance(command, SubscribeCommand):
            self._handle_subscribe(session, command)
            return

        if isinstance(command, UnsubscribeCommand):
            self._handle_unsubscribe(session)
            return

        # Registered handlers
        handler = self._command_handlers.get(cmd_type)
        if handler:
            try:
                response = handler(session, command)
                if response:
                    session.send(response)
            except Exception as e:
                logger.error(f"Handler error for {cmd_type.__name__}: {e}")
                session.send(ErrorMessage(
                    code="HANDLER_ERROR",
                    message=f"Error handling {cmd_type.__name__}: {e}"
                ))
        else:
            session.send(ErrorMessage(
                code="NO_HANDLER",
                message=f"No handler for command: {cmd_type.__name__}"
            ))

    def _handle_subscribe(self, session: ClientSession, cmd: SubscribeCommand) -> None:
        """Handle a subscribe command."""
        try:
            # Create appropriate reader based on mode
            if cmd.mode == "live":
                reader = BufferReader(self._buffer)
            elif cmd.mode == "since":
                if not cmd.marker:
                    session.send(ErrorMessage(
                        code="MISSING_MARKER",
                        message="'since' mode requires marker parameter"
                    ))
                    return
                try:
                    reader = BufferReader.from_marker(self._buffer, cmd.marker)
                except KeyError:
                    session.send(ErrorMessage(
                        code="MARKER_NOT_FOUND",
                        message=f"Marker '{cmd.marker}' not found or evicted"
                    ))
                    return
            elif cmd.mode == "tail":
                lines = cmd.lines or 100
                reader = BufferReader.from_tail(self._buffer, lines)
            elif cmd.mode == "replay":
                if cmd.marker:
                    try:
                        reader = BufferReader.from_marker(self._buffer, cmd.marker)
                    except KeyError:
                        session.send(ErrorMessage(
                            code="MARKER_NOT_FOUND",
                            message=f"Marker '{cmd.marker}' not found or evicted"
                        ))
                        return
                else:
                    # Replay from beginning
                    start_seq, _ = self._buffer.sequence_range
                    reader = BufferReader(self._buffer, start_seq)
            else:
                session.send(ErrorMessage(
                    code="INVALID_MODE",
                    message=f"Unknown subscribe mode: {cmd.mode}"
                ))
                return

            session.subscribed = True
            session.subscribe_mode = cmd.mode
            session.reader = reader

            session.send(SubscribedMessage(mode=cmd.mode, position=reader.position))

            # Start streaming thread
            self._start_stream_thread(session, is_replay=(cmd.mode == "replay"))

        except Exception as e:
            logger.error(f"Subscribe error: {e}")
            session.send(ErrorMessage(code="SUBSCRIBE_ERROR", message=str(e)))

    def _handle_unsubscribe(self, session: ClientSession) -> None:
        """Handle an unsubscribe command."""
        session.subscribed = False
        session.reader = None

        # Stop streaming thread if running
        if session.client_id in self._stream_threads:
            # Thread will exit on next iteration when session.subscribed is False
            pass

    def _start_stream_thread(self, session: ClientSession, is_replay: bool = False) -> None:
        """Start a thread to stream data to a subscribed client."""
        def stream_loop():
            reader = session.reader
            if not reader:
                return

            try:
                while session.subscribed and session.running and self._running:
                    # Read available entries
                    entries = reader.read_available()

                    for seq, entry in entries:
                        if not session.running:
                            break
                        self._send_entry(session, entry)

                    if is_replay and not entries:
                        # Replay mode: disconnect when caught up
                        break

                    if not is_replay:
                        # Live mode: wait for more data
                        self._buffer.wait_for_data(reader.position, timeout=1.0)

            except Exception as e:
                logger.error(f"Stream error for client {session.client_id}: {e}")

            finally:
                if session.client_id in self._stream_threads:
                    del self._stream_threads[session.client_id]

        thread = threading.Thread(target=stream_loop, daemon=True)
        self._stream_threads[session.client_id] = thread
        thread.start()

    def _send_entry(self, session: ClientSession, entry: BufferEntry) -> bool:
        """Send a buffer entry to a client."""
        if entry.source == "marker":
            msg = MarkerMessage(name=entry.marker or "", timestamp=entry.timestamp)
        else:
            msg = DataMessage(
                source=entry.source,
                timestamp=entry.timestamp,
                data=entry.data
            )
        return session.send(msg)

    def _cleanup_client(self, session: ClientSession) -> None:
        """Clean up a disconnected client."""
        with self._clients_lock:
            if session.client_id in self._clients:
                del self._clients[session.client_id]

        session.close()

        if self._on_disconnect:
            try:
                self._on_disconnect(session)
            except Exception as e:
                logger.error(f"on_disconnect callback error: {e}")

        logger.info(f"Client {session.client_id} disconnected")
