"""Terminal command handlers."""

import asyncio
import logging
import time
from typing import Any, Dict, List, Optional

from .base import AsyncHandler
from .runtime_user import get_default_runtime_user, wrap_shell_command
from .session import SessionManager

logger = logging.getLogger(__name__)


class TerminalStartHandler(AsyncHandler):
    """Handler for starting new terminal sessions."""
    
    @property
    def command_name(self) -> str:
        return "terminal_start"
    
    async def execute(self, message: Dict[str, Any]) -> Dict[str, Any]:
        """Start a new terminal session."""
        shell = message.get("shell")
        cwd = message.get("cwd")
        project_id = message.get("project_id")
        
        session_manager = self.context.get("session_manager")
        if not session_manager:
            raise RuntimeError("Session manager not available")
        
        session_info = await session_manager.create_session(
            shell=shell,
            cwd=cwd,
            project_id=project_id,
            run_as_user=get_default_runtime_user(message),
        )
        
        # Start background watcher for process exit
        asyncio.create_task(self._watch_process_exit(session_info["terminal_id"]))
        
        return {
            "event": "terminal_started",
            "terminal_id": session_info["terminal_id"],
            "channel": session_info["channel"],
            "pid": session_info["pid"],
            "shell": session_info.get("shell"),
            "cwd": session_info.get("cwd"),
            "project_id": session_info.get("project_id"),
        }
    
    async def _watch_process_exit(self, terminal_id: str) -> None:
        """Watch for process exit and send notification."""
        session_manager = self.context.get("session_manager")
        if not session_manager:
            return
        
        session = session_manager.get_session(terminal_id)
        if not session:
            return
        
        await session.proc.wait()
        
        await self.send_response({
            "event": "terminal_exit",
            "terminal_id": terminal_id,
            "returncode": session.proc.returncode,
            "project_id": session.project_id,
        }, project_id=session.project_id)
        
        # Only cleanup session if it still exists (not already removed by stop handler)
        if session_manager.get_session(terminal_id):
            session_manager.remove_session(terminal_id)


class TerminalSendHandler(AsyncHandler):
    """Handler for sending data to terminal sessions."""
    
    @property
    def command_name(self) -> str:
        return "terminal_send"
    
    async def execute(self, message: Dict[str, Any]) -> Dict[str, Any]:
        """Send data to a terminal session."""
        terminal_id = message.get("terminal_id")
        data = message.get("data", "")
        
        if not terminal_id:
            raise ValueError("terminal_id is required")
        
        session_manager = self.context.get("session_manager")
        if not session_manager:
            raise RuntimeError("Session manager not available")
        
        session = session_manager.get_session(terminal_id)
        if not session:
            raise ValueError(f"terminal_id {terminal_id} not found")
        
        await session.write(data)
        
        # No response expected for terminal_send
        return {"event": "terminal_send_ack"}


class TerminalStopHandler(AsyncHandler):
    """Handler for stopping terminal sessions."""
    
    @property
    def command_name(self) -> str:
        return "terminal_stop"
    
    async def execute(self, message: Dict[str, Any]) -> Dict[str, Any]:
        """Stop a terminal session."""
        terminal_id = message.get("terminal_id")
        
        if not terminal_id:
            logger.error("terminal_stop: Missing terminal_id in message")
            raise ValueError("terminal_id is required")
        
        logger.info("terminal_stop: Processing stop request for terminal_id=%s", terminal_id)
        
        session_manager = self.context.get("session_manager")
        if not session_manager:
            logger.error("terminal_stop: Session manager not available in context")
            raise RuntimeError("Session manager not available")
        
        # Remove session from manager first
        session = session_manager.remove_session(terminal_id)
        if not session:
            logger.warning("terminal_stop: Terminal %s not found, may have already been stopped", terminal_id)
            # Send completion event immediately for not found terminals
            asyncio.create_task(self._send_not_found_completion(terminal_id, None))
            return {
                "event": "terminal_stopped",
                "terminal_id": terminal_id,
                "status": "not_found",
                "message": "Terminal was not found or already stopped",
                "project_id": None,
            }
        
        logger.info("terminal_stop: Found session for terminal %s (PID: %s), starting background stop process", 
                   terminal_id, getattr(session.proc, 'pid', 'unknown'))
        
        # Start stop process in background without blocking the control channel
        asyncio.create_task(self._stop_session_safely(session, terminal_id, session.project_id))
        
        return {
            "event": "terminal_stopped",
            "terminal_id": terminal_id,
            "status": "stopping",
            "message": "Terminal stop process initiated",
            "project_id": session.project_id,
        }
    
    async def _stop_session_safely(self, session, terminal_id: str, project_id: Optional[str] = None) -> None:
        """Safely stop a session in the background with timeout and error handling."""
        logger.info("terminal_stop: Starting background stop process for terminal %s", terminal_id)
        
        try:
            # Attempt graceful stop with timeout
            await asyncio.wait_for(session.stop(), timeout=10.0)
            logger.info("terminal_stop: Successfully stopped terminal %s", terminal_id)
            
            # Send success notification
            await self.send_response({
                "event": "terminal_stop_completed",
                "terminal_id": terminal_id,
                "status": "success",
                "message": "Terminal stopped successfully",
                "project_id": project_id,
            }, project_id=project_id)
            
        except asyncio.TimeoutError:
            logger.warning("terminal_stop: Stop timeout for terminal %s, forcing kill", terminal_id)
            
            # Force kill the process
            try:
                if hasattr(session.proc, 'kill'):
                    session.proc.kill()
                    logger.info("terminal_stop: Force killed terminal %s", terminal_id)
                elif hasattr(session.proc, 'terminate'):
                    session.proc.terminate()
                    logger.info("terminal_stop: Force terminated terminal %s", terminal_id)
            except Exception as kill_exc:
                logger.error("terminal_stop: Failed to force kill terminal %s: %s", terminal_id, kill_exc)
            
            # Send timeout notification
            await self.send_response({
                "event": "terminal_stop_completed",
                "terminal_id": terminal_id,
                "status": "timeout",
                "message": "Terminal stop timed out, process was force killed",
                "project_id": project_id,
            }, project_id=project_id)
            
        except Exception as exc:
            logger.exception("terminal_stop: Error stopping terminal %s: %s", terminal_id, exc)
            
            # Send error notification
            await self.send_response({
                "event": "terminal_stop_completed",
                "terminal_id": terminal_id,
                "status": "error",
                "message": f"Error stopping terminal: {str(exc)}",
                "project_id": project_id,
            }, project_id=project_id)

    async def _send_not_found_completion(self, terminal_id: str, project_id: Optional[str] = None) -> None:
        """Send completion event for not found terminals."""
        await self.send_response({
            "event": "terminal_stop_completed",
            "terminal_id": terminal_id,
            "status": "not_found",
            "message": "Terminal was not found or already stopped",
            "project_id": project_id,
        }, project_id=project_id)


class TerminalListHandler(AsyncHandler):
    """Handler for listing terminal sessions."""
    
    @property
    def command_name(self) -> str:
        return "terminal_list"
    
    async def handle(self, message: Dict[str, Any], reply_channel: Optional[str] = None) -> None:
        """Handle the command by executing it and sending the response to the requesting client session."""
        logger.info("handler: Processing command %s with reply_channel=%s",
                   self.command_name, reply_channel)

        try:
            response = await self.execute(message)
            logger.info("handler: Command %s executed successfully", self.command_name)

            # Automatically copy request_id if present in the incoming message
            if "request_id" in message and "request_id" not in response:
                response["request_id"] = message["request_id"]

            # Get the source client session from the message
            source_client_session = message.get("source_client_session")
            project_id = response.get("project_id")

            logger.info("handler: %s response project_id=%s, source_client_session=%s",
                       self.command_name, project_id, source_client_session)

            # Send response only to the requesting client session
            if source_client_session:
                # Add client_sessions field to target only the requesting session
                response["client_sessions"] = [source_client_session]

                import json
                logger.info("handler: 📤 SENDING EVENT '%s' (via direct control_channel.send)", response.get("event", "unknown"))
                logger.info("handler: 📤 FULL EVENT PAYLOAD: %s", json.dumps(response, indent=2, default=str))

                await self.control_channel.send(response)
            else:
                # Fallback to original behavior if no source_client_session
                await self.send_response(response, reply_channel, project_id)
        except Exception as exc:
            logger.exception("handler: Error in command %s: %s", self.command_name, exc)
            await self.send_error(str(exc), reply_channel, message.get("project_id"))
    
    async def execute(self, message: Dict[str, Any]) -> Dict[str, Any]:
        """List all terminal sessions."""
        session_manager = self.context.get("session_manager")
        if not session_manager:
            raise RuntimeError("Session manager not available")

        # Accept project_id argument: None (default) = only no project, 'all' = all, else = filter by project_id
        requested_project_id = message.get("project_id")
        logger.info("terminal_list: requested_project_id=%r (type: %s)", requested_project_id, type(requested_project_id))

        if requested_project_id == "all":
            logger.info("terminal_list: Using 'all' mode to list all terminals")
            sessions = session_manager.list_sessions(project_id="all")
        else:
            logger.info("terminal_list: Filtering by project_id=%r", requested_project_id)
            sessions = session_manager.list_sessions(project_id=requested_project_id)

        logger.info("terminal_list: Found %d sessions, returning with project_id=%r", len(sessions), requested_project_id)
        return {
            "event": "terminal_list",
            "sessions": sessions,
            "project_id": requested_project_id,
        }


class TerminalExecHandler(AsyncHandler):
    """Handler for running one-off commands and returning structured results."""

    OUTPUT_FLUSH_INTERVAL_S = 1.0

    @property
    def command_name(self) -> str:
        return "terminal_exec"

    async def execute(self, message: Dict[str, Any]) -> Dict[str, Any]:
        """Run a command on the host and capture stdout, stderr, and the exit code."""
        command = message.get("command")
        if not command:
            raise ValueError("command is required for terminal_exec")

        cwd = message.get("cwd")
        env_spec = message.get("env")
        timeout_param = message.get("timeout")
        automation_task_id = message.get("automation_task_id")
        automation_step_index = message.get("automation_step_index")
        timeout_value: Optional[float] = None
        if timeout_param is not None:
            try:
                timeout_value = float(timeout_param)
                if timeout_value <= 0:
                    raise ValueError("timeout must be a positive number of seconds")
            except (TypeError, ValueError) as exc:
                raise ValueError("timeout must be a positive number") from exc

        env = None
        if env_spec is not None:
            if not isinstance(env_spec, dict):
                raise ValueError("env must be an object/dict of environment variables")
            env = {str(k): str(v) for k, v in env_spec.items()}

        project_id = message.get("project_id")

        pending_stdout: List[str] = []
        pending_stderr: List[str] = []
        final_stdout: List[str] = []
        final_stderr: List[str] = []
        pending_lock = asyncio.Lock()
        stop_event = asyncio.Event()

        async def _append_pending(bucket: List[str], text: str) -> None:
            async with pending_lock:
                bucket.append(text)

        async def _read_stream(stream, pending: List[str], archive: List[str]) -> None:
            while True:
                chunk = await stream.read(1024)
                if not chunk:
                    break
                text = chunk.decode(errors="replace")
                await _append_pending(pending, text)
                archive.append(text)

        async def _flush_pending() -> None:
            async with pending_lock:
                stdout_chunk = "".join(pending_stdout)
                stderr_chunk = "".join(pending_stderr)
                pending_stdout.clear()
                pending_stderr.clear()
            if not stdout_chunk and not stderr_chunk:
                return
            payload: Dict[str, Any] = {
                "event": "terminal_exec_output",
                "command": command,
            }
            if stdout_chunk:
                payload["stdout"] = stdout_chunk
            if stderr_chunk:
                payload["stderr"] = stderr_chunk
            if project_id is not None:
                payload["project_id"] = project_id
            if automation_task_id is not None:
                payload["automation_task_id"] = automation_task_id
            if automation_step_index is not None:
                payload["automation_step_index"] = automation_step_index
            await self._broadcast_terminal_exec_payload(payload)

        async def _streaming_flusher() -> None:
            try:
                while True:
                    try:
                        await asyncio.wait_for(stop_event.wait(), timeout=self.OUTPUT_FLUSH_INTERVAL_S)
                        break
                    except asyncio.TimeoutError:
                        await _flush_pending()
                await _flush_pending()
            except asyncio.CancelledError:
                await _flush_pending()
                raise

        start = time.monotonic()
        process = await asyncio.create_subprocess_shell(
            wrap_shell_command(command, get_default_runtime_user(message)),
            cwd=cwd,
            env=env,
            stdout=asyncio.subprocess.PIPE,
            stderr=asyncio.subprocess.PIPE,
        )

        stdout_task = asyncio.create_task(_read_stream(process.stdout, pending_stdout, final_stdout))
        stderr_task = asyncio.create_task(_read_stream(process.stderr, pending_stderr, final_stderr))
        flusher_task = asyncio.create_task(_streaming_flusher())

        try:
            await asyncio.wait_for(process.wait(), timeout=timeout_value)
        except asyncio.TimeoutError:
            process.kill()
            await process.wait()
            raise TimeoutError(
                f"terminal_exec timed out after {timeout_value}s" if timeout_value is not None else "terminal_exec timed out"
            )
        finally:
            stop_event.set()
            await asyncio.gather(stdout_task, stderr_task, flusher_task, return_exceptions=True)

        duration = time.monotonic() - start
        stdout = "".join(final_stdout)
        stderr = "".join(final_stderr)

        response: Dict[str, Any] = {
            "event": "terminal_exec_result",
            "command": command,
            "returncode": process.returncode,
            "stdout": stdout,
            "stderr": stderr,
            "duration_s": round(duration, 3),
        }
        if automation_task_id is not None:
            response["automation_task_id"] = automation_task_id
        if automation_step_index is not None:
            response["automation_step_index"] = automation_step_index
        if cwd is not None:
            response["cwd"] = cwd
        if project_id is not None:
            response["project_id"] = project_id
        return response

    async def _broadcast_terminal_exec_payload(self, payload: Dict[str, Any]) -> None:
        """Send terminal_exec events to all currently connected client sessions."""
        target_sessions = self._gather_client_session_channels()
        enhanced_payload = dict(payload)
        if target_sessions:
            enhanced_payload["client_sessions"] = target_sessions
        enhanced_payload["bypass_session_gate"] = True
        await self.control_channel.send(enhanced_payload)

    def _gather_client_session_channels(self) -> List[str]:
        manager = self.context.get("client_session_manager")
        if not manager:
            return []
        return list(manager.get_sessions().keys())
