"""Client session: Support multiple concurrent command runs from a single connection.

This module enables advanced client usage patterns:
- Multiple concurrent commands from same client connection
- Proper multiplexing of output between runs
- Resource pooling (single SSH connection per host)
- Interactive shell-like behavior
"""

import asyncio
import json
import logging
from dataclasses import dataclass, field
from typing import Dict, Optional, Callable, Any

from rpipe.protocol import (
    read_frame,
    pack_frame,
    RunRequest,
    TYPE_RUN_REQ,
    TYPE_STDIN,
    TYPE_CANCEL,
    TYPE_META,
    TYPE_STDOUT,
    TYPE_STDERR,
    TYPE_EXIT,
    TYPE_ERROR,
)

logger = logging.getLogger(__name__)


@dataclass
class RunResult:
    """Result from a single command execution."""

    run_id: int
    exit_code: Optional[int] = None
    stdout: bytes = field(default_factory=bytes)
    stderr: bytes = field(default_factory=bytes)
    error: Optional[str] = None
    completed: bool = False


class ClientSession:
    """Manages a session with the daemon, supporting multiple concurrent runs."""

    def __init__(self, reader, writer, default_host: Optional[str] = None):
        """
        Initialize client session.

        Args:
            reader: asyncio StreamReader connected to daemon
            writer: asyncio StreamWriter connected to daemon
            default_host: Default remote host for commands
        """
        self.reader = reader
        self.writer = writer
        self.default_host = default_host
        self._next_run_id = 1
        self._runs: Dict[int, RunResult] = {}
        self._run_metadata: Dict[int, Dict[str, Any]] = {}
        self._output_callbacks: Dict[int, Dict[str, Callable]] = {}

    def _alloc_run_id(self) -> int:
        """Allocate a unique run ID for this session."""
        rid = self._next_run_id
        self._next_run_id += 1
        return rid

    def set_output_callback(
        self,
        run_id: int,
        on_stdout: Optional[Callable[[bytes], None]] = None,
        on_stderr: Optional[Callable[[bytes], None]] = None,
    ) -> None:
        """
        Set callbacks for output from a run.

        Args:
            run_id: Run ID to set callbacks for
            on_stdout: Callback for stdout data
            on_stderr: Callback for stderr data
        """
        if run_id not in self._output_callbacks:
            self._output_callbacks[run_id] = {}

        if on_stdout:
            self._output_callbacks[run_id]["stdout"] = on_stdout
        if on_stderr:
            self._output_callbacks[run_id]["stderr"] = on_stderr

    async def run_command(
        self,
        cmd: str,
        host: Optional[str] = None,
        username: Optional[str] = None,
        cwd: Optional[str] = None,
        env: Optional[Dict[str, str]] = None,
        timeout_s: int = 60,
        policy: Optional[str] = None,
        read_only: bool = True,
        stdin_mode: str = "none",
    ) -> int:
        """
        Submit a command to execute on the daemon.

        Args:
            cmd: Command to execute
            host: Remote host (uses default_host if not specified)
            username: SSH username
            cwd: Working directory on remote
            env: Environment variables
            timeout_s: Command timeout in seconds
            policy: Capability policy to use
            read_only: Whether to use read-only mode (v0 compat)
            stdin_mode: "none" or "pipe"

        Returns:
            Run ID for this command

        Raises:
            ValueError: If host not specified and no default
        """
        if not host:
            host = self.default_host
        if not host:
            raise ValueError("host required (no default specified)")

        # Allocate run ID
        run_id = self._alloc_run_id()

        # Create run result tracker
        self._runs[run_id] = RunResult(run_id=run_id)
        self._run_metadata[run_id] = {
            "cmd": cmd,
            "host": host,
            "username": username,
            "started_at": None,
        }

        # Build RUN_REQ
        req = RunRequest(
            host=host,
            cmd=cmd,
            cwd=cwd,
            env=env,
            username=username,
            timeout_s=timeout_s,
            policy=policy,
            read_only=read_only,
            stdin_mode=stdin_mode,
        )

        # Send RUN_REQ frame
        req_payload = json.dumps(req.to_dict()).encode("utf-8")
        frame = pack_frame(TYPE_RUN_REQ, run_id, req_payload)
        self.writer.write(frame)
        await self.writer.drain()

        logger.debug(f"run_id={run_id} submitted: {cmd} on {host}")

        return run_id

    def get_result(self, run_id: int) -> RunResult:
        """
        Get result of a completed run (synchronous).

        Args:
            run_id: Run ID to get results for

        Returns:
            RunResult with exit code and output

        Raises:
            KeyError: If run_id not found
            RuntimeError: If run not completed
        """
        if run_id not in self._runs:
            raise KeyError(f"run_id {run_id} not found")

        result = self._runs[run_id]
        if not result.completed:
            raise RuntimeError(f"run_id {run_id} still executing")

        return result

    async def send_stdin(self, run_id: int, data: bytes) -> None:
        """
        Send stdin data to a running command.

        Args:
            run_id: Run ID to send data to
            data: Bytes to send
        """
        frame = pack_frame(TYPE_STDIN, run_id, data)
        self.writer.write(frame)
        await self.writer.drain()

    async def cancel_command(self, run_id: int, hard: bool = False) -> None:
        """
        Cancel a running command.

        Args:
            run_id: Run ID to cancel
            hard: If True, send SIGKILL (vs SIGTERM)
        """
        # For now, just send CANCEL frame. Hard/soft is daemon decision.
        frame = pack_frame(TYPE_CANCEL, run_id, b"")
        self.writer.write(frame)
        await self.writer.drain()

    async def pump_responses(self) -> None:
        """
        Pump responses from daemon, processing frames and updating run results.

        This should be called as a background task or in the main loop.
        It reads frames and routes them to the appropriate run.
        """
        try:
            while True:
                ftype, run_id, payload = await read_frame(self.reader)

                if run_id == 0:
                    # Session-level frame (not tied to specific run)
                    logger.warning(f"Received session-level frame type {ftype}")
                    continue

                if run_id not in self._runs:
                    logger.warning(f"Received frame for unknown run_id {run_id}")
                    continue

                result = self._runs[run_id]

                if ftype == TYPE_META:
                    # Metadata about the run
                    meta = json.loads(payload.decode("utf-8"))
                    self._run_metadata[run_id].update(meta)
                    logger.debug(f"run_id={run_id} META: {meta}")

                elif ftype == TYPE_STDOUT:
                    # Stdout data
                    result.stdout += payload
                    callback = self._output_callbacks.get(run_id, {}).get("stdout")
                    if callback:
                        callback(payload)
                    logger.debug(f"run_id={run_id} STDOUT: {len(payload)} bytes")

                elif ftype == TYPE_STDERR:
                    # Stderr data
                    result.stderr += payload
                    callback = self._output_callbacks.get(run_id, {}).get("stderr")
                    if callback:
                        callback(payload)
                    logger.debug(f"run_id={run_id} STDERR: {len(payload)} bytes")

                elif ftype == TYPE_EXIT:
                    # Process completed
                    if len(payload) >= 4:
                        import struct

                        (exit_code,) = struct.unpack("!i", payload[:4])
                        result.exit_code = exit_code
                    result.completed = True
                    logger.info(f"run_id={run_id} completed with exit_code={result.exit_code}")

                elif ftype == TYPE_ERROR:
                    # Error occurred
                    error_dict = json.loads(payload.decode("utf-8"))
                    result.error = error_dict.get("error", "Unknown error")
                    result.completed = True
                    logger.error(f"run_id={run_id} error: {result.error}")

                else:
                    logger.warning(f"run_id={run_id} unknown frame type {ftype}")

        except asyncio.CancelledError:
            logger.debug("pump_responses cancelled")
            raise
        except Exception as e:
            logger.error(f"pump_responses error: {e}")
            # Mark all pending runs as errored
            for result in self._runs.values():
                if not result.completed:
                    result.error = str(e)
                    result.completed = True

    async def wait_for_run(self, run_id: int, timeout: Optional[float] = None) -> RunResult:
        """
        Wait for a specific run to complete.

        Args:
            run_id: Run ID to wait for
            timeout: Maximum time to wait (None for no timeout)

        Returns:
            RunResult when completed

        Raises:
            TimeoutError: If timeout exceeded
            KeyError: If run_id not found
        """
        if run_id not in self._runs:
            raise KeyError(f"run_id {run_id} not found")

        async def wait():
            while not self._runs[run_id].completed:
                await asyncio.sleep(0.01)
            return self._runs[run_id]

        try:
            return await asyncio.wait_for(wait(), timeout=timeout)
        except asyncio.TimeoutError:
            raise TimeoutError(f"Timeout waiting for run_id {run_id}")

    def get_active_runs(self) -> list[int]:
        """Get list of active (not completed) run IDs."""
        return [rid for rid, result in self._runs.items() if not result.completed]

    def close(self) -> None:
        """Close the session connection."""
        try:
            self.writer.close()
        except Exception:
            pass

    async def async_close(self) -> None:
        """Asynchronously close the session."""
        try:
            await self.writer.wait_closed()
        except Exception:
            pass
