"""
Shell Manager - Persistent Shell Session Management

Manages long-running shell subprocesses per agent session.
Each shell maintains a rolling output buffer and writes state to workspace.
"""

import subprocess
import threading
import time
import json
import os
import signal
import logging
from collections import deque
from datetime import datetime
from typing import Dict, List, Optional, Tuple
from pathlib import Path

logger = logging.getLogger(__name__)

MAX_SHELLS_PER_SESSION = 5
OUTPUT_BUFFER_SIZE = 5000  # Max lines to keep in memory


class ShellSession:
    """A single persistent shell subprocess with output buffering"""

    def __init__(self, shell_id: str, cwd: str, label: str = ""):
        self.shell_id = shell_id
        self.label = label or shell_id
        self.cwd = cwd
        self.created_at = datetime.now()
        self.last_activity = datetime.now()
        self.output_buffer: deque = deque(maxlen=OUTPUT_BUFFER_SIZE)
        self.process: Optional[subprocess.Popen] = None
        self._reader_thread: Optional[threading.Thread] = None
        self._stop_event = threading.Event()

    def start(self):
        """Start the shell subprocess"""
        if os.name == 'nt':
            # Windows: persistent PowerShell
            self.process = subprocess.Popen(
                ["powershell", "-NoProfile", "-NoLogo", "-ExecutionPolicy", "Bypass"],
                stdin=subprocess.PIPE,
                stdout=subprocess.PIPE,
                stderr=subprocess.STDOUT,
                text=True,
                cwd=self.cwd,
                bufsize=1,
                creationflags=subprocess.CREATE_NEW_PROCESS_GROUP
            )
        else:
            # Linux/Mac: persistent bash
            self.process = subprocess.Popen(
                ["/bin/bash", "--norc", "--noprofile"],
                stdin=subprocess.PIPE,
                stdout=subprocess.PIPE,
                stderr=subprocess.STDOUT,
                text=True,
                cwd=self.cwd,
                bufsize=1
            )

        # Start background reader thread
        self._reader_thread = threading.Thread(
            target=self._read_output,
            daemon=True,
            name=f"shell-reader-{self.shell_id}"
        )
        self._reader_thread.start()
        logger.info(f"[SHELL] Started shell '{self.shell_id}' (pid={self.process.pid}) in {self.cwd}")

    def _read_output(self):
        """Background thread: continuously read output into buffer"""
        try:
            while not self._stop_event.is_set() and self.process and self.process.poll() is None:
                line = self.process.stdout.readline()
                if line:
                    self.output_buffer.append(line.rstrip('\n\r'))
                    self.last_activity = datetime.now()
                else:
                    break
        except (ValueError, OSError):
            pass  # Pipe closed
        finally:
            # Capture any remaining output
            if self.process and self.process.stdout:
                try:
                    remaining = self.process.stdout.read()
                    if remaining:
                        for line in remaining.split('\n'):
                            if line.strip():
                                self.output_buffer.append(line.rstrip('\r'))
                except (ValueError, OSError):
                    pass

    def send_command(self, command: str, wait_ms: int = 3000) -> str:
        """Send command to shell and wait for output"""
        if not self.process or self.process.poll() is not None:
            return f"[Shell '{self.shell_id}' has exited with code {self.process.returncode if self.process else 'N/A'}]"

        # Record buffer size before command
        before_count = len(self.output_buffer)

        try:
            # Send command
            self.process.stdin.write(command + '\n')
            self.process.stdin.flush()
            self.last_activity = datetime.now()
        except (BrokenPipeError, OSError) as e:
            return f"[Shell error: {e}]"

        # Wait for output to accumulate
        time.sleep(wait_ms / 1000.0)

        # Collect new output since command was sent
        after_count = len(self.output_buffer)
        new_lines = list(self.output_buffer)[max(0, before_count - (OUTPUT_BUFFER_SIZE - after_count)):]

        # If buffer wrapped, just return last N lines
        if after_count >= OUTPUT_BUFFER_SIZE:
            new_lines = list(self.output_buffer)[-100:]

        if not new_lines:
            return "[No output yet — use SHELL_VIEW to check later]"

        output = '\n'.join(new_lines[-200:])  # Cap at 200 lines
        if len(new_lines) > 200:
            output = f"[... {len(new_lines) - 200} lines truncated ...]\n" + output

        return output

    def read_output(self, lines: int = 50) -> str:
        """Read last N lines from output buffer"""
        if not self.output_buffer:
            return "[No output in buffer]"

        recent = list(self.output_buffer)[-lines:]
        status = "running" if self.is_alive else f"exited (code={self.process.returncode})"
        header = f"[Shell '{self.shell_id}' — {status} — {len(self.output_buffer)} total lines]\n"
        return header + '\n'.join(recent)

    @property
    def is_alive(self) -> bool:
        return self.process is not None and self.process.poll() is None

    @property
    def status(self) -> str:
        if not self.process:
            return "not_started"
        if self.process.poll() is None:
            return "running"
        return f"exited:{self.process.returncode}"

    @property
    def pid(self) -> Optional[int]:
        return self.process.pid if self.process else None

    @property
    def last_output_line(self) -> str:
        if self.output_buffer:
            return str(self.output_buffer[-1])[:100]
        return ""

    def kill(self):
        """Terminate the shell process"""
        self._stop_event.set()
        if self.process and self.process.poll() is None:
            try:
                if os.name == 'nt':
                    self.process.terminate()
                else:
                    os.kill(self.process.pid, signal.SIGTERM)
                self.process.wait(timeout=3)
            except subprocess.TimeoutExpired:
                self.process.kill()
                self.process.wait(timeout=2)
            except (OSError, ProcessLookupError):
                pass
            logger.info(f"[SHELL] Killed shell '{self.shell_id}' (pid={self.pid})")

    def to_dict(self) -> dict:
        return {
            "id": self.shell_id,
            "label": self.label,
            "pid": self.pid,
            "status": self.status,
            "cwd": self.cwd,
            "created": self.created_at.isoformat(),
            "last_activity": self.last_activity.isoformat(),
            "last_output_line": self.last_output_line,
            "buffer_lines": len(self.output_buffer)
        }


class ShellManager:
    """Manages shell sessions per agent session"""

    _instance = None

    @classmethod
    def get_instance(cls) -> 'ShellManager':
        if cls._instance is None:
            cls._instance = cls()
        return cls._instance

    def __init__(self):
        self._sessions: Dict[str, Dict[str, ShellSession]] = {}
        self._lock = threading.Lock()
        self._counter: Dict[str, int] = {}  # session_id -> next shell number

    def create_shell(self, session_id: str, cwd: str, label: str = "",
                     workspace_path: str = "") -> Tuple[ShellSession, str]:
        """Create a new shell session. Returns (session, message)."""
        with self._lock:
            if session_id not in self._sessions:
                self._sessions[session_id] = {}
                self._counter[session_id] = 1

            shells = self._sessions[session_id]
            if len(shells) >= MAX_SHELLS_PER_SESSION:
                return None, f"Maximum {MAX_SHELLS_PER_SESSION} shells reached. Kill one first with SHELL_MANAGE."

            shell_id = f"shell-{self._counter[session_id]}"
            self._counter[session_id] += 1

            session = ShellSession(shell_id, cwd, label)
            session.start()
            shells[shell_id] = session

        self._write_workspace_file(session_id, workspace_path)
        return session, f"Created shell '{shell_id}' (pid={session.pid})"

    def get_or_create_default(self, session_id: str, cwd: str,
                              workspace_path: str = "") -> ShellSession:
        """Get existing default shell or create one"""
        with self._lock:
            shells = self._sessions.get(session_id, {})
            # Return first alive shell
            for shell in shells.values():
                if shell.is_alive:
                    return shell

        # No alive shells, create one
        session, _ = self.create_shell(session_id, cwd, label="default", workspace_path=workspace_path)
        return session

    def get_shell(self, session_id: str, shell_id: str) -> Optional[ShellSession]:
        """Get a specific shell"""
        with self._lock:
            return self._sessions.get(session_id, {}).get(shell_id)

    def list_shells(self, session_id: str) -> List[dict]:
        """List all shells for a session"""
        with self._lock:
            shells = self._sessions.get(session_id, {})
            return [s.to_dict() for s in shells.values()]

    def kill_shell(self, session_id: str, shell_id: str,
                   workspace_path: str = "") -> bool:
        """Kill a specific shell"""
        with self._lock:
            shells = self._sessions.get(session_id, {})
            session = shells.get(shell_id)
            if not session:
                return False
            session.kill()
            del shells[shell_id]

        self._write_workspace_file(session_id, workspace_path)
        return True

    def cleanup_session(self, session_id: str):
        """Kill all shells for a session (on disconnect)"""
        with self._lock:
            shells = self._sessions.pop(session_id, {})
            self._counter.pop(session_id, None)

        for shell in shells.values():
            try:
                shell.kill()
            except Exception as e:
                logger.warning(f"[SHELL] Error killing shell {shell.shell_id}: {e}")

        logger.info(f"[SHELL] Cleaned up {len(shells)} shell(s) for session {session_id}")

    def _write_workspace_file(self, session_id: str, workspace_path: str):
        """Write .shells.json to workspace for agent visibility"""
        if not workspace_path:
            return
        try:
            shells = self.list_shells(session_id)
            filepath = Path(workspace_path) / ".shells.json"
            filepath.write_text(json.dumps({"shells": shells}, indent=2, default=str))
        except Exception as e:
            logger.warning(f"[SHELL] Failed to write .shells.json: {e}")
