"""Daemon: Unix socket server, request handler, process orchestrator."""

import asyncio
import json
import logging
import os
import struct
import time
from pathlib import Path
from typing import Dict, Optional

from rpipe.audit import get_audit_logger
from rpipe.execution import CommandExecutionError, validate_command_execution
from rpipe.protocol import (
    FramingError,
    RunRequest,
    read_frame,
    pack_frame,
    TYPE_RUN_REQ,
    TYPE_STDIN,
    TYPE_CANCEL,
    TYPE_META,
    TYPE_STDOUT,
    TYPE_STDERR,
    TYPE_EXIT,
    TYPE_ERROR,
)
from rpipe.policy import Policy, validate_command
from rpipe.capability import (
    CapabilitySet,
    DEFAULT_READ_ONLY_CAPABILITIES,
    DEFAULT_WRITE_CAPABLE_CAPABILITIES,
    DEFAULT_RESTRICTIVE_CAPABILITIES,
)
from rpipe.policy_capabilities import validate_command_against_capabilities
from rpipe.ssh_backend import AsyncSSHBackend, SSHPool, RunHandle, SSHError

logger = logging.getLogger(__name__)


class DaemonError(Exception):
    """Daemon-level error."""
    pass


class RPipeDaemon:
    """
    Unix socket daemon: listens for RUN_REQ frames, spawns SSH processes, streams output.

    Current Architecture (v1):
    - Single-threaded event loop with concurrent client handling
    - Multiple concurrent client connections (asyncio tasks)
    - Multiple concurrent SSH channels (per host, per client)
    - Client connection tracking and lifecycle management

    Guarantees:
    - Each client receives all frames for their run(s) in correct order
    - No frame loss or corruption
    - Timeout is enforced; process is killed after timeout_s
    - Concurrent clients don't interfere with each other
    """

    def __init__(
        self,
        sock_path: str,
        ssh_pool: Optional[SSHPool] = None,
        known_hosts: Optional[str] = None,
        audit_log_dir: Optional[str] = None,
    ):
        """
        Initialize daemon.

        Args:
            sock_path: Path to Unix domain socket
            ssh_pool: SSHPool instance (created if None)
            known_hosts: Path to SSH known_hosts file for host key verification
                        Required for secure v1 deployment
            audit_log_dir: Directory for audit logs. If None, logs to stderr only.
        """
        self.sock_path = sock_path
        self.pool = ssh_pool or SSHPool(known_hosts=known_hosts)
        self.backend = AsyncSSHBackend(self.pool)
        self.audit_logger = get_audit_logger()  # Get or create global audit logger
        self._next_run_id = 1
        self._runs: Dict[int, RunHandle] = {}
        self._runs_lock = asyncio.Lock()
        self._client_tasks: set = set()  # Track active client connection tasks
        self._client_tasks_lock = asyncio.Lock()

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

    def _load_policy(self, policy_name: Optional[str] = None) -> CapabilitySet:
        """
        Load policy by name or return default.

        Args:
            policy_name: Policy preset name ("read_only", "write_capable", "restrictive")
                        or None for default

        Returns:
            CapabilitySet for the requested policy

        Raises:
            DaemonError: If policy name is unknown
        """
        if not policy_name or policy_name == "read_only":
            return CapabilitySet(DEFAULT_READ_ONLY_CAPABILITIES)
        elif policy_name == "write_capable":
            return CapabilitySet(DEFAULT_WRITE_CAPABLE_CAPABILITIES)
        elif policy_name == "restrictive":
            return CapabilitySet(DEFAULT_RESTRICTIVE_CAPABILITIES)
        else:
            raise DaemonError(f"Unknown policy: {policy_name}")

    async def _handle_client_wrapper(
        self, reader: asyncio.StreamReader, writer: asyncio.StreamWriter
    ) -> None:
        """
        Wrapper for handling a single client connection.

        Manages client lifecycle: tracks connection, ensures cleanup on errors.
        Allows multiple concurrent client connections.

        Args:
            reader: Client reader stream
            writer: Client writer stream
        """
        # Track this client task
        task = asyncio.current_task()
        if task:
            async with self._client_tasks_lock:
                self._client_tasks.add(task)

        try:
            # Handle the actual client request
            await self.handle_client(reader, writer)
        finally:
            # Clean up the task tracking
            if task:
                async with self._client_tasks_lock:
                    self._client_tasks.discard(task)

            # Ensure streams are closed
            try:
                writer.close()
                await writer.wait_closed()
            except Exception:
                pass

    async def handle_client(self, reader: asyncio.StreamReader, writer: asyncio.StreamWriter) -> None:
        """
        Handle one client connection: receive RUN_REQ, spawn process, stream output.

        Current limitation: Serial per-client. Next client waits until this completes.

        Args:
            reader: Client reader
            writer: Client writer
        """
        try:
            # Read RUN_REQ frame
            ftype, run_id, payload = await read_frame(reader)

            # First frame must be RUN_REQ with run_id=0
            if ftype != TYPE_RUN_REQ or run_id != 0:
                await self._send_error(writer, 0, "first frame must be RUN_REQ with run_id=0")
                return

            # Parse request
            try:
                req_dict = json.loads(payload.decode("utf-8"))
                req = RunRequest.from_dict(req_dict)
            except (json.JSONDecodeError, KeyError, ValueError) as e:
                await self._send_error(writer, 0, f"invalid RUN_REQ: {e}")
                return

            # Validate cwd/env (security: prevent shell injection)
            try:
                validate_command_execution(cwd=req.cwd, env=req.env)
            except CommandExecutionError as e:
                await self._send_error(writer, 0, f"unsafe cwd/env: {e}")
                return

            # Allocate run ID first (for logging)
            rid = self._alloc_run_id()

            # Validate command against policy
            # Try v1 capability-based validation first, fall back to v0 token-based
            try:
                if req.policy:
                    # Use new capability-based validation
                    cap_set = self._load_policy(req.policy)
                    validate_command_against_capabilities(req.cmd, cap_set)
                else:
                    # Fall back to old token-based validation
                    policy = Policy(read_only=req.read_only)
                    validate_command(req.cmd, policy)
            except Exception as e:
                # Log policy rejection
                self.audit_logger.log_command_denied(
                    run_id=rid,
                    cmd=req.cmd,
                    binary=req.cmd.split()[0] if req.cmd else "",
                    reason=str(e),
                    policy=req.policy,
                    host=req.host,
                    username=req.username,
                )
                await self._send_error(writer, rid, f"command rejected by policy: {e}")
                return

            # Log command approval
            self.audit_logger.log_command_allowed(
                run_id=rid,
                cmd=req.cmd,
                binary=req.cmd.split()[0] if req.cmd else "",
                host=req.host,
                policy=req.policy,
                username=req.username,
                cwd=req.cwd,
            )

            # Send META frame with run_id
            meta = {
                "run_id": rid,
                "host": req.host,
                "cmd": req.cmd,
                "cwd": req.cwd,
                "pty": req.pty,
                "stdin_mode": req.stdin_mode,
                "timeout_s": req.timeout_s,
                "read_only": req.read_only,
                "started_at": time.time(),
            }
            writer.write(pack_frame(TYPE_META, rid, json.dumps(meta).encode("utf-8")))
            await writer.drain()

            # Spawn remote process
            try:
                run = await self.backend.run(
                    rid,
                    req.host,
                    req.cmd,
                    username=req.username,
                    cwd=req.cwd,
                    env=req.env,
                    timeout_s=req.timeout_s,
                )
            except SSHError as e:
                # Log SSH execution error
                self.audit_logger.log_command_error(
                    run_id=rid,
                    cmd=req.cmd,
                    binary=req.cmd.split()[0] if req.cmd else "",
                    error=str(e),
                    error_type="ssh_error",
                    host=req.host,
                    username=req.username,
                )
                await self._send_error(writer, rid, f"process execution failed: {e}")
                return

            # Track active run
            async with self._runs_lock:
                self._runs[rid] = run

            # Start output streaming tasks
            stdout_task = asyncio.create_task(self._pump_stream(writer, rid, run.proc.stdout, TYPE_STDOUT))
            stderr_task = asyncio.create_task(self._pump_stream(writer, rid, run.proc.stderr, TYPE_STDERR))

            # Start stdin relaying task (if enabled)
            stdin_task = None
            if req.stdin_mode == "pipe":
                stdin_task = asyncio.create_task(self._pump_stdin(reader, rid, run))

            # Wait for process with timeout
            start_time = time.time()
            try:
                exit_code = await asyncio.wait_for(run.wait(), timeout=req.timeout_s)
            except asyncio.TimeoutError:
                logger.warning(f"run_id={rid} timeout after {req.timeout_s}s; killing with SIGKILL")
                # Log timeout event
                self.audit_logger.log_command_error(
                    run_id=rid,
                    cmd=req.cmd,
                    binary=req.cmd.split()[0] if req.cmd else "",
                    error=f"Timeout after {req.timeout_s}s",
                    error_type="timeout",
                    host=req.host,
                    username=req.username,
                )
                run.cancel(hard=True)  # SIGKILL
                try:
                    exit_code = await asyncio.wait_for(run.wait(), timeout=5.0)
                except asyncio.TimeoutError:
                    logger.error(f"run_id={rid} SIGKILL failed; giving up")
                    exit_code = -1

            # Calculate duration
            duration_s = time.time() - start_time

            # Wait for stream tasks to finish
            try:
                await asyncio.gather(stdout_task, stderr_task, return_exceptions=True)
            except Exception as e:
                logger.error(f"run_id={rid} error waiting for streams: {e}")

            if stdin_task:
                stdin_task.cancel()

            # Send EXIT frame with exit code
            exit_payload = struct.pack("!i", int(exit_code))
            writer.write(pack_frame(TYPE_EXIT, rid, exit_payload))
            await writer.drain()

            # Log command completion
            self.audit_logger.log_command_completed(
                run_id=rid,
                exit_code=int(exit_code),
                duration_s=duration_s,
            )

            logger.info(f"run_id={rid} completed with exit_code={exit_code}")

        except Exception as e:
            logger.exception(f"handle_client error: {e}")
            await self._send_error(writer, 0, f"internal error: {e}")

        finally:
            try:
                writer.close()
                await writer.wait_closed()
            except Exception as e:
                logger.warning(f"error closing client: {e}")

    async def _pump_stream(
        self,
        writer: asyncio.StreamWriter,
        run_id: int,
        stream,
        frame_type: int,
    ) -> None:
        """
        Stream data from remote process to client.

        Reads chunks from stream, encodes as frames, sends to client.

        Args:
            writer: Client writer
            run_id: Run identifier
            stream: asyncssh stream (stdout or stderr)
            frame_type: TYPE_STDOUT or TYPE_STDERR
        """
        try:
            while True:
                chunk = await stream.read(8192)  # 8KB chunks
                if not chunk:
                    break

                # Handle both bytes and str (depending on asyncssh config)
                if isinstance(chunk, str):
                    chunk = chunk.encode("utf-8")

                writer.write(pack_frame(frame_type, run_id, chunk))
                try:
                    await writer.drain()
                except Exception as e:
                    logger.warning(f"run_id={run_id} drain error: {e}")
                    break

        except Exception as e:
            logger.error(f"run_id={run_id} pump_stream error: {e}")

    async def _pump_stdin(
        self,
        reader: asyncio.StreamReader,
        run_id: int,
        run: RunHandle,
    ) -> None:
        """
        Relay stdin from client to remote process.

        Reads STDIN and CANCEL frames from client, forwards to process.

        Args:
            reader: Client reader
            run_id: Run identifier
            run: RunHandle for the process
        """
        try:
            while True:
                ftype, frame_run_id, payload = await read_frame(reader)

                # Only process frames for this run
                if frame_run_id != run_id:
                    continue

                if ftype == TYPE_STDIN:
                    run.proc.stdin.write(payload)
                    try:
                        await run.proc.stdin.drain()
                    except Exception as e:
                        logger.warning(f"run_id={run_id} stdin write error: {e}")
                        return

                elif ftype == TYPE_CANCEL:
                    logger.info(f"run_id={run_id} cancel requested")
                    run.cancel(hard=False)  # SIGTERM
                    return

        except asyncio.CancelledError:
            pass
        except Exception as e:
            logger.error(f"run_id={run_id} pump_stdin error: {e}")
        finally:
            try:
                run.proc.stdin.write_eof()
            except Exception:
                pass

    async def _send_error(self, writer: asyncio.StreamWriter, run_id: int, msg: str) -> None:
        """Send ERROR frame to client."""
        payload = json.dumps({"error": msg}).encode("utf-8")
        writer.write(pack_frame(TYPE_ERROR, run_id, payload))
        try:
            await writer.drain()
        except Exception as e:
            logger.warning(f"error sending error frame: {e}")


async def serve(sock_path: str, known_hosts: Optional[str] = None, audit_log_dir: Optional[str] = None) -> None:
    """
    Start daemon server.

    Handles multiple concurrent client connections efficiently.

    Args:
        sock_path: Path to Unix domain socket
        known_hosts: Path to SSH known_hosts file for host key verification.
                    If None (insecure), will warn at startup.
                    REQUIRED: Use ~/.ssh/known_hosts or /etc/ssh/ssh_known_hosts for production.
        audit_log_dir: Directory for audit logs. If None, logs to stderr only.
    """
    # Initialize audit logging
    if audit_log_dir:
        from rpipe.audit import init_audit_logger
        init_audit_logger(audit_log_dir)

    # Create socket directory
    sock_dir = os.path.dirname(sock_path)
    if sock_dir:
        Path(sock_dir).mkdir(parents=True, exist_ok=True)

    # Clean up stale socket
    try:
        os.unlink(sock_path)
    except FileNotFoundError:
        pass

    # Create daemon with optional host key verification
    # For v1+, pass a real known_hosts file path
    daemon = RPipeDaemon(sock_path, known_hosts=known_hosts)

    # Start server with concurrent client handling
    # Note: start_unix_server creates a new asyncio.Task for each client connection,
    # so multiple clients are handled concurrently by the event loop
    server = await asyncio.start_unix_server(
        daemon._handle_client_wrapper, path=sock_path
    )

    # Set strict permissions (owner only)
    try:
        os.chmod(sock_path, 0o600)
    except Exception as e:
        logger.warning(f"failed to set socket permissions: {e}")

    logger.info(f"rpipe-daemon listening on {sock_path}")

    # Host key verification warning
    if not known_hosts:
        logger.error("⚠️  CRITICAL SECURITY WARNING: Host key verification is disabled (known_hosts=None)")
        logger.error("⚠️  This makes the daemon VULNERABLE to MITM (man-in-the-middle) attacks!")
        logger.error("⚠️  For production use, configure known_hosts with: ~/.ssh/known_hosts or /etc/ssh/ssh_known_hosts")
    else:
        logger.info(f"Host key verification enabled: {known_hosts}")

    async with server:
        try:
            await server.serve_forever()
        except KeyboardInterrupt:
            logger.info("daemon shutting down")
        finally:
            # Wait for all client tasks to complete
            if daemon._client_tasks:
                logger.info(f"waiting for {len(daemon._client_tasks)} client tasks to finish...")
                await asyncio.gather(*daemon._client_tasks, return_exceptions=True)

            await pool.close_all()
            logger.info("daemon stopped")


if __name__ == "__main__":
    logging.basicConfig(
        level=logging.INFO,
        format="%(asctime)s [%(levelname)s] %(name)s: %(message)s",
    )

    import sys

    # Parse command-line arguments
    sock = os.path.join(os.path.expanduser("~"), ".rpipe", "daemon.sock")
    known_hosts = None
    audit_log_dir = None

    # Usage: python -m rpipe.daemon [socket_path] [known_hosts_path] [audit_log_dir]
    if len(sys.argv) > 1:
        sock = sys.argv[1]
    if len(sys.argv) > 2:
        known_hosts = sys.argv[2]
    if len(sys.argv) > 3:
        audit_log_dir = sys.argv[3]

    # Check environment variables as fallback
    if not known_hosts:
        known_hosts = os.environ.get("RPIPE_KNOWN_HOSTS")

    if not audit_log_dir:
        audit_log_dir = os.environ.get("RPIPE_AUDIT_LOG_DIR")

    # Try default system SSH known_hosts if not specified
    if not known_hosts:
        for default_path in [os.path.expanduser("~/.ssh/known_hosts"), "/etc/ssh/ssh_known_hosts"]:
            if os.path.exists(default_path):
                known_hosts = default_path
                break

    asyncio.run(serve(sock, known_hosts=known_hosts, audit_log_dir=audit_log_dir))
