"""E2B-specific terminal manager using native process management."""

from __future__ import annotations

import asyncio
from dataclasses import dataclass
from typing import TYPE_CHECKING, Any
import uuid

from anyenv.log import get_logger
from anyenv.process_manager import ProcessManagerProtocol, ProcessOutput
from anyenv.process_manager.process_manager import BaseTerminal


if TYPE_CHECKING:
    from pathlib import Path

    from e2b import AsyncSandbox
    from e2b.sandbox_async.commands.command_handle import AsyncCommandHandle


logger = get_logger(__name__)


@dataclass(kw_only=True)
class E2BTerminal(BaseTerminal):
    """Represents a terminal session using E2B's process management."""

    pid: int | None = None
    _handle: AsyncCommandHandle | None = None

    def is_running(self) -> bool:
        """Check if terminal is still running."""
        if self._exit_code is not None:
            return False
        if self._handle:
            # Check if handle has exit_code property (None means still running)
            return self._handle.exit_code is None
        return False

    def get_exit_code(self) -> int | None:
        """Get the exit code if available."""
        if self._exit_code is not None:
            return self._exit_code
        if self._handle and self._handle.exit_code is not None:
            # Get exit code from handle if available
            self._exit_code = self._handle.exit_code
            return self._exit_code
        return None

    def set_handle(self, handle: AsyncCommandHandle) -> None:
        """Set the E2B command handle."""
        self._handle = handle
        self.pid = handle.pid


class E2BTerminalManager(ProcessManagerProtocol):
    """Terminal manager that uses E2B's native process management."""

    def __init__(self, sandbox: AsyncSandbox) -> None:
        """Initialize with an E2B sandbox instance."""
        self.sandbox = sandbox
        self._terminals: dict[str, E2BTerminal] = {}

    async def start_process(
        self,
        command: str,
        args: list[str] | None = None,
        cwd: str | Path | None = None,
        env: dict[str, str] | None = None,
        output_limit: int | None = None,
    ) -> str:
        """Create a new terminal session using E2B's background commands."""
        terminal_id = f"e2b_term_{uuid.uuid4().hex[:8]}"
        args = args or []
        env = env or {}
        full_command = f"{command} {' '.join(args)}" if args else command
        terminal = E2BTerminal(
            terminal_id=terminal_id,
            command=command,
            args=args,
            cwd=str(cwd) if cwd else None,
            env=env,
            output_limit=output_limit or 1048576,
        )
        self._terminals[terminal_id] = terminal

        # Start the process using E2B's background execution
        try:
            # Create output handlers to collect output in real-time
            def on_stdout(data: str) -> None:
                terminal.add_output(data)

            def on_stderr(data: str) -> None:
                terminal.add_output(f"STDERR: {data}")

            # Start command in background with streaming handlers
            handle = await self.sandbox.commands.run(
                full_command,
                background=True,
                envs=env,
                cwd=str(cwd) if cwd else None,
                on_stdout=on_stdout,
                on_stderr=on_stderr,
            )

            terminal.set_handle(handle)
            msg = "Created E2B terminal %s (PID %s): %s"
            logger.info(msg, terminal_id, handle.pid, full_command)

        except Exception as e:
            # Clean up on failure
            self._terminals.pop(terminal_id, None)
            msg = f"Failed to create E2B terminal: {e}"
            logger.exception(msg)
            raise RuntimeError(msg) from e
        else:
            return terminal_id

    async def get_output(self, process_id: str) -> ProcessOutput:
        """Get current output from a process."""
        terminal = self.get_terminal(process_id)
        output = terminal.get_output()
        exit_code = terminal.get_exit_code()

        return ProcessOutput(stdout=output, stderr="", combined=output, exit_code=exit_code)

    async def wait_for_exit(self, process_id: str) -> int:
        """Wait for process to complete."""
        terminal = self.get_terminal(process_id)
        handle = terminal._handle  # noqa: SLF001
        try:
            if handle and handle.exit_code is None:
                # Wait for the E2B command to complete
                result = await handle.wait()
                terminal.set_exit_code(result.exit_code)

                # Add any final output from the final result
                # (streaming output is already collected via callbacks)

        except Exception:
            logger.exception("Error waiting for process %s", process_id)
            terminal.set_exit_code(1)

        return terminal.get_exit_code() or 0

    async def kill_process(self, process_id: str) -> None:
        """Kill a running process using E2B's process management."""
        terminal = self.get_terminal(process_id)
        try:
            # Kill the E2B process using the PID
            if terminal.pid and terminal.is_running():
                killed = await self.sandbox.commands.kill(terminal.pid)
                if killed:
                    terminal.set_exit_code(130)  # SIGINT exit code
                    msg = "Killed E2B process %s (PID %s)"
                    logger.info(msg, process_id, terminal.pid)
                else:
                    msg = "Failed to kill E2B terminal %s (PID %s) - process not found"
                    logger.warning(msg, process_id, terminal.pid)
                    terminal.set_exit_code(1)

        except Exception:
            logger.exception("Error killing process %s", process_id)
            terminal.set_exit_code(1)

    async def release_process(self, process_id: str) -> None:
        """Release process resources."""
        terminal = self.get_terminal(process_id)
        if terminal.is_running():
            await self.kill_process(process_id)
        del self._terminals[process_id]
        logger.info("Released process %s", process_id)

    # async def list_processes(self) -> dict[str, dict[str, Any]]:
    #     """List all tracked terminals and their status."""
    #     return {
    #         terminal_id: await self.get_process_info(terminal_id)
    #         for terminal_id in self._terminals
    #     }

    def get_terminal(self, terminal_id: str):
        """Get terminal by ID."""
        if terminal_id not in self._terminals:
            msg = f"Process {terminal_id} not found"
            raise ValueError(msg)

        return self._terminals[terminal_id]

    async def get_process_info(self, process_id: str) -> dict[str, Any]:
        """Get information about a specific process."""
        terminal = self.get_terminal(process_id)
        return {
            "terminal_id": process_id,
            "command": terminal.command,
            "args": terminal.args,
            "cwd": terminal.cwd,
            "pid": terminal.pid,
            "created_at": terminal.created_at.isoformat(),
            "is_running": terminal.is_running(),
            "exit_code": terminal.get_exit_code(),
            "output_limit": terminal.output_limit,
        }

    async def list_processes(self) -> list[str]:
        """List all tracked terminals."""
        return list(self._terminals.keys())

    async def get_sandbox_processes(self) -> dict[int, dict[str, Any]]:
        """Get all processes running in the E2B sandbox."""
        try:
            # Use E2B's list command to get all running processes
            processes = await self.sandbox.commands.list()
            result = {}
            for process_info in processes:
                result[process_info.pid] = {
                    "pid": process_info.pid,
                    "tag": process_info.tag,
                    "command": process_info.cmd,
                    "args": process_info.args,
                    "cwd": process_info.cwd,
                    "envs": process_info.envs,
                }
        except Exception:
            logger.exception("Error listing sandbox processes")
            return {}
        else:
            return result

    async def connect_to_process(
        self,
        pid: int,
        output_byte_limit: int = 1048576,
    ) -> str:
        """Connect to an existing process in the sandbox and manage it as a terminal."""
        terminal_id = f"e2b_conn_{uuid.uuid4().hex[:8]}"

        try:
            # Get process info first
            processes = await self.get_sandbox_processes()
            if pid not in processes:
                msg = f"Process {pid} not found in sandbox"
                raise ValueError(msg)  # noqa: TRY301

            process_info = processes[pid]

            # Create terminal for the existing process
            terminal = E2BTerminal(
                terminal_id=terminal_id,
                command=process_info["command"],
                args=process_info.get("args", []),
                cwd=process_info.get("cwd"),
                env=process_info.get("envs", {}),
                pid=pid,
                output_limit=output_byte_limit,
            )

            def on_stdout(data: str) -> None:
                terminal.add_output(data)

            def on_stderr(data: str) -> None:
                terminal.add_output(f"STDERR: {data}")

            # Connect to the existing process
            handle = await self.sandbox.commands.connect(
                pid=pid,
                on_stdout=on_stdout,
                on_stderr=on_stderr,
            )

            terminal.set_handle(handle)
            self._terminals[terminal_id] = terminal

            logger.info("Connected to E2B process %s as terminal %s", pid, terminal_id)

        except Exception as e:
            msg = f"Failed to connect to process {pid}: {e}"
            logger.exception(msg)
            raise RuntimeError(msg) from e
        else:
            return terminal_id

    async def send_stdin(self, terminal_id: str, data: str) -> None:
        """Send data to terminal stdin (if supported)."""
        terminal = self.get_terminal(terminal_id)
        if not terminal.pid:
            msg = f"Terminal {terminal_id} has no process ID"
            raise ValueError(msg)

        try:
            await self.sandbox.commands.send_stdin(terminal.pid, data)
            msg = "Sent stdin to terminal %s (PID %s): %r"
            logger.debug(msg, terminal_id, terminal.pid, data[:50])
        except Exception:
            logger.exception("Error sending stdin to terminal %s", terminal_id)
            raise

    async def cleanup(self) -> None:
        """Clean up all terminals."""
        logger.info("Cleaning up %s E2B terminals", len(self._terminals))
        if cleanup_tasks := [self.release_process(id_) for id_ in self._terminals]:
            await asyncio.gather(*cleanup_tasks, return_exceptions=True)
        logger.info("E2B terminal cleanup completed")
