"""Terminal service - terminal I/O coordination."""

import asyncio
import logging
import re
from contextlib import suppress
from datetime import UTC, datetime
from typing import Any

from porterminal.domain import (
    PTYPort,
    RateLimitConfig,
    Session,
    TerminalDimensions,
    TokenBucketRateLimiter,
)

from ..ports.connection_port import ConnectionPort

logger = logging.getLogger(__name__)

# Terminal response sequences that should NOT be written to PTY.
# These are responses from the terminal emulator to queries from applications.
# If written to PTY, they get echoed back and displayed as garbage.
#
# Patterns:
#   \x1b[?...c  - Device Attributes (DA) response
#   \x1b[...R   - Cursor Position Report (CPR) response
TERMINAL_RESPONSE_PATTERN = re.compile(rb"\x1b\[\?[\d;]*c|\x1b\[[\d;]*R")

# Constants
HEARTBEAT_INTERVAL = 30  # seconds
HEARTBEAT_TIMEOUT = 300  # 5 minutes
PTY_READ_INTERVAL = 0.008  # ~120Hz polling
OUTPUT_BATCH_INTERVAL = 0.016  # ~60Hz output (batch writes for smoother rendering)
OUTPUT_BATCH_MAX_SIZE = 16384  # Flush if batch exceeds 16KB
INTERACTIVE_THRESHOLD = 64  # Bytes - flush immediately for small interactive data
MAX_INPUT_SIZE = 4096


class AsyncioClock:
    """Clock implementation using asyncio event loop time."""

    def now(self) -> float:
        return asyncio.get_running_loop().time()


class TerminalService:
    """Service for handling terminal I/O.

    Coordinates PTY reads, WebSocket writes, and message handling.
    Supports multiple clients connected to the same session.
    """

    def __init__(
        self,
        rate_limit_config: RateLimitConfig | None = None,
        max_input_size: int = MAX_INPUT_SIZE,
    ) -> None:
        self._rate_limit_config = rate_limit_config or RateLimitConfig()
        self._max_input_size = max_input_size

        # Multi-client support: track connections and read loops per session
        self._session_connections: dict[str, set[ConnectionPort]] = {}
        self._session_read_tasks: dict[str, asyncio.Task[None]] = {}

    # -------------------------------------------------------------------------
    # Multi-client connection tracking
    # -------------------------------------------------------------------------

    def _register_connection(self, session_id: str, connection: ConnectionPort) -> int:
        """Register a connection for a session. Returns connection count."""
        if session_id not in self._session_connections:
            self._session_connections[session_id] = set()
        self._session_connections[session_id].add(connection)
        return len(self._session_connections[session_id])

    def _unregister_connection(self, session_id: str, connection: ConnectionPort) -> int:
        """Unregister a connection. Returns remaining count."""
        if session_id not in self._session_connections:
            return 0
        self._session_connections[session_id].discard(connection)
        count = len(self._session_connections[session_id])
        if count == 0:
            del self._session_connections[session_id]
        return count

    async def _broadcast_output(self, session_id: str, data: bytes) -> None:
        """Broadcast PTY output to all connections for a session."""
        connections = self._session_connections.get(session_id, set())
        dead: list[ConnectionPort] = []
        for conn in list(connections):  # Copy to avoid mutation during iteration
            try:
                await conn.send_output(data)
            except Exception:
                dead.append(conn)
        for conn in dead:
            connections.discard(conn)

    async def _broadcast_message(self, session_id: str, message: dict[str, Any]) -> None:
        """Broadcast JSON message to all connections for a session."""
        connections = self._session_connections.get(session_id, set())
        dead: list[ConnectionPort] = []
        for conn in list(connections):
            try:
                await conn.send_message(message)
            except Exception:
                dead.append(conn)
        for conn in dead:
            connections.discard(conn)

    async def handle_session(
        self,
        session: Session[PTYPort],
        connection: ConnectionPort,
        skip_buffer: bool = False,
    ) -> None:
        """Handle terminal session I/O with multi-client support.

        Multiple clients can connect to the same session simultaneously.
        The first client starts the PTY read loop; the last client stops it.

        Args:
            session: Terminal session to handle.
            connection: Network connection to client.
            skip_buffer: Whether to skip sending buffered output.
        """
        session_id = str(session.id)
        clock = AsyncioClock()
        rate_limiter = TokenBucketRateLimiter(self._rate_limit_config, clock)

        # Register this connection
        connection_count = self._register_connection(session_id, connection)
        is_first_client = connection_count == 1

        logger.info(
            "Client connected session_id=%s connection_count=%d",
            session_id,
            connection_count,
        )

        try:
            # First client starts the shared PTY read loop
            if is_first_client:
                self._start_broadcast_read_loop(session, session_id)

            # Note: session_info is sent by the caller (app.py) to include tab_id

            # Replay buffered output to THIS connection only (not broadcast)
            if not skip_buffer and not session.output_buffer.is_empty:
                buffered = session.get_buffered_output()
                # Don't clear buffer - other clients may need it too
                if buffered:
                    await connection.send_output(buffered)

            # Start heartbeat for this connection
            heartbeat_task = asyncio.create_task(self._heartbeat_loop(connection))

            try:
                await self._handle_input_loop(session, connection, rate_limiter)
            finally:
                heartbeat_task.cancel()
                with suppress(asyncio.CancelledError):
                    await heartbeat_task

        finally:
            # Unregister this connection
            remaining = self._unregister_connection(session_id, connection)

            logger.info(
                "Client disconnected session_id=%s remaining_connections=%d",
                session_id,
                remaining,
            )

            # Last client: stop the read loop
            if remaining == 0:
                await self._stop_broadcast_read_loop(session_id)

    def _start_broadcast_read_loop(
        self,
        session: Session[PTYPort],
        session_id: str,
    ) -> None:
        """Start the PTY read loop that broadcasts to all clients."""
        if session_id in self._session_read_tasks:
            return  # Already running

        task = asyncio.create_task(self._read_pty_broadcast_loop(session, session_id))
        self._session_read_tasks[session_id] = task
        logger.debug("Started broadcast read loop session_id=%s", session_id)

    async def _stop_broadcast_read_loop(self, session_id: str) -> None:
        """Stop the PTY read loop for a session."""
        task = self._session_read_tasks.pop(session_id, None)
        if task and not task.done():
            task.cancel()
            with suppress(asyncio.CancelledError):
                await task
        logger.debug("Stopped broadcast read loop session_id=%s", session_id)

    async def _read_pty_broadcast_loop(
        self,
        session: Session[PTYPort],
        session_id: str,
    ) -> None:
        """Read from PTY and broadcast to all connected clients.

        Single loop per session, regardless of client count.

        Batching strategy:
        - Small data (<64 bytes): flush immediately for interactive responsiveness
        - Large data: batch for ~16ms to reduce WebSocket message frequency
        - Flush if batch exceeds 16KB to prevent memory buildup
        """
        # Check if PTY is alive at start
        if not session.pty_handle.is_alive():
            logger.error("PTY not alive at start session_id=%s", session.id)
            await self._broadcast_output(session_id, b"\r\n[PTY failed to start]\r\n")
            return

        batch_buffer: list[bytes] = []
        batch_size = 0
        last_flush_time = asyncio.get_running_loop().time()

        async def flush_batch() -> None:
            nonlocal batch_buffer, batch_size, last_flush_time
            if batch_buffer:
                combined = b"".join(batch_buffer)
                batch_buffer = []
                batch_size = 0
                last_flush_time = asyncio.get_running_loop().time()
                await self._broadcast_output(session_id, combined)

        def has_connections() -> bool:
            return (
                session_id in self._session_connections
                and len(self._session_connections[session_id]) > 0
            )

        while has_connections() and session.pty_handle.is_alive():
            try:
                data = session.pty_handle.read(4096)
                if data:
                    session.add_output(data)
                    session.touch(datetime.now(UTC))

                    # Small data (interactive): flush immediately for responsiveness
                    if len(data) < INTERACTIVE_THRESHOLD and not batch_buffer:
                        await self._broadcast_output(session_id, data)
                    else:
                        # Batch larger data
                        batch_buffer.append(data)
                        batch_size += len(data)

                        # Flush if batch is large enough
                        if batch_size >= OUTPUT_BATCH_MAX_SIZE:
                            await flush_batch()

            except Exception as e:
                logger.error("PTY read error session_id=%s: %s", session.id, e)
                await flush_batch()  # Flush any pending data
                await self._broadcast_output(session_id, f"\r\n[PTY error: {e}]\r\n".encode())
                break

            # Check if we should flush based on time
            current_time = asyncio.get_running_loop().time()
            if batch_buffer and (current_time - last_flush_time) >= OUTPUT_BATCH_INTERVAL:
                await flush_batch()

            await asyncio.sleep(PTY_READ_INTERVAL)

        # Flush any remaining data
        await flush_batch()

        # Notify all clients if PTY died
        if has_connections() and not session.pty_handle.is_alive():
            await self._broadcast_output(session_id, b"\r\n[Shell exited]\r\n")

    async def _heartbeat_loop(self, connection: ConnectionPort) -> None:
        """Send periodic heartbeat pings."""
        while connection.is_connected():
            await asyncio.sleep(HEARTBEAT_INTERVAL)
            try:
                await connection.send_message({"type": "ping"})
            except Exception:
                break

    async def _handle_input_loop(
        self,
        session: Session[PTYPort],
        connection: ConnectionPort,
        rate_limiter: TokenBucketRateLimiter,
    ) -> None:
        """Handle input from client."""
        while connection.is_connected():
            try:
                message = await connection.receive()
            except Exception:
                break

            if isinstance(message, bytes):
                await self._handle_binary_input(session, message, rate_limiter, connection)
            elif isinstance(message, dict):
                await self._handle_json_message(session, message, rate_limiter, connection)

    async def _handle_binary_input(
        self,
        session: Session[PTYPort],
        data: bytes,
        rate_limiter: TokenBucketRateLimiter,
        connection: ConnectionPort,
    ) -> None:
        """Handle binary terminal input."""
        if len(data) > self._max_input_size:
            await connection.send_message(
                {
                    "type": "error",
                    "message": "Input too large",
                }
            )
            return

        # Filter terminal response sequences before writing to PTY.
        # xterm.js generates these in response to DA/CPR queries.
        # If written back to PTY, they get echoed and displayed as garbage.
        filtered = TERMINAL_RESPONSE_PATTERN.sub(b"", data)
        if not filtered:
            return

        if rate_limiter.try_acquire(len(filtered)):
            session.pty_handle.write(filtered)
            session.touch(datetime.now(UTC))
        else:
            await connection.send_message(
                {
                    "type": "error",
                    "message": "Rate limit exceeded",
                }
            )
            logger.warning("Rate limit exceeded session_id=%s", session.id)

    async def _handle_json_message(
        self,
        session: Session[PTYPort],
        message: dict[str, Any],
        rate_limiter: TokenBucketRateLimiter,
        connection: ConnectionPort,
    ) -> None:
        """Handle JSON control message."""
        msg_type = message.get("type")

        if msg_type == "resize":
            await self._handle_resize(session, message)
        elif msg_type == "input":
            await self._handle_json_input(session, message, rate_limiter, connection)
        elif msg_type == "ping":
            await connection.send_message({"type": "pong"})
            session.touch(datetime.now(UTC))
        elif msg_type == "pong":
            session.touch(datetime.now(UTC))
        else:
            logger.warning("Unknown message type session_id=%s type=%s", session.id, msg_type)

    async def _handle_resize(
        self,
        session: Session[PTYPort],
        message: dict[str, Any],
    ) -> None:
        """Handle terminal resize message."""
        cols = int(message.get("cols", 120))
        rows = int(message.get("rows", 30))

        new_dims = TerminalDimensions.clamped(cols, rows)

        # Skip if same as current
        if session.dimensions == new_dims:
            return

        session.update_dimensions(new_dims)
        session.pty_handle.resize(new_dims)
        session.touch(datetime.now(UTC))

        logger.info(
            "Terminal resized session_id=%s cols=%d rows=%d",
            session.id,
            new_dims.cols,
            new_dims.rows,
        )

    async def _handle_json_input(
        self,
        session: Session[PTYPort],
        message: dict[str, Any],
        rate_limiter: TokenBucketRateLimiter,
        connection: ConnectionPort,
    ) -> None:
        """Handle JSON-encoded terminal input."""
        data = message.get("data", "")

        if len(data) > self._max_input_size:
            await connection.send_message(
                {
                    "type": "error",
                    "message": "Input too large",
                }
            )
            return

        if data:
            input_bytes = data.encode("utf-8")
            # Filter terminal response sequences
            filtered = TERMINAL_RESPONSE_PATTERN.sub(b"", input_bytes)
            if not filtered:
                return

            if rate_limiter.try_acquire(len(filtered)):
                session.pty_handle.write(filtered)
                session.touch(datetime.now(UTC))
            else:
                await connection.send_message(
                    {
                        "type": "error",
                        "message": "Rate limit exceeded",
                    }
                )
