"""Local PTY-backed shell sessions for the mock workspace backend."""

from __future__ import annotations

import codecs
import fcntl
import os
import shlex
import shutil
import signal
import struct
import subprocess
import termios
import threading
import time
from pathlib import Path
from typing import IO, Literal

ShellState = Literal["running", "stopped"]

SHELL_SIGNAL_NAMES = ("HUP", "INT", "TERM", "KILL")
_SHELL_SIGNAL_MAP = {
    "HUP": signal.SIGHUP,
    "INT": signal.SIGINT,
    "TERM": signal.SIGTERM,
    "KILL": signal.SIGKILL,
}

_LOCAL_SHELLS: dict[str, "LocalShellSession"] = {}
_LOCAL_SHELLS_LOCK = threading.Lock()


def _shell_argv(*, interactive: bool) -> list[str]:
    shell_program = shutil.which("bash") or "/bin/sh"
    argv = [shell_program]
    if shell_program.endswith("bash"):
        argv.extend(["--noprofile", "--norc"])
    if interactive:
        argv.append("-i")
    return argv


def _redact_text(text: str, redact_values: list[str]) -> str:
    redacted = text
    for secret_value in sorted(
        {item for item in redact_values if item != ""},
        key=len,
        reverse=True,
    ):
        redacted = redacted.replace(secret_value, "[REDACTED]")
    return redacted


def _set_pty_size(fd: int, rows: int, cols: int) -> None:
    winsize = struct.pack("HHHH", rows, cols, 0, 0)
    fcntl.ioctl(fd, termios.TIOCSWINSZ, winsize)


class LocalShellSession:
    """Host-local interactive shell used by the mock backend."""

    def __init__(
        self,
        *,
        shell_id: str,
        cwd: Path,
        display_cwd: str,
        cols: int,
        rows: int,
        env_overrides: dict[str, str] | None = None,
        redact_values: list[str] | None = None,
    ) -> None:
        self.shell_id = shell_id
        self.cwd = display_cwd
        self.cols = cols
        self.rows = rows
        self.started_at = time.time()
        self.ended_at: float | None = None
        self.exit_code: int | None = None
        self.state: ShellState = "running"
        self.pid: int | None = None
        self._lock = threading.RLock()
        self._output = ""
        self._master_fd: int | None = None
        self._input_pipe: IO[bytes] | None = None
        self._output_pipe: IO[bytes] | None = None
        self._reader: threading.Thread | None = None
        self._waiter: threading.Thread | None = None
        self._decoder = codecs.getincrementaldecoder("utf-8")("replace")
        self._redact_values = list(redact_values or [])
        env = os.environ.copy()
        env.update(
            {
                "TERM": env.get("TERM", "xterm-256color"),
                "PS1": "pyro$ ",
                "PROMPT_COMMAND": "",
            }
        )
        if env_overrides is not None:
            env.update(env_overrides)

        process: subprocess.Popen[bytes]
        try:
            master_fd, slave_fd = os.openpty()
        except OSError:
            process = subprocess.Popen(  # noqa: S603
                _shell_argv(interactive=False),
                stdin=subprocess.PIPE,
                stdout=subprocess.PIPE,
                stderr=subprocess.STDOUT,
                cwd=str(cwd),
                env=env,
                text=False,
                close_fds=True,
                preexec_fn=os.setsid,
            )
            self._input_pipe = process.stdin
            self._output_pipe = process.stdout
        else:
            try:
                _set_pty_size(slave_fd, rows, cols)
                process = subprocess.Popen(  # noqa: S603
                    _shell_argv(interactive=True),
                    stdin=slave_fd,
                    stdout=slave_fd,
                    stderr=slave_fd,
                    cwd=str(cwd),
                    env=env,
                    text=False,
                    close_fds=True,
                    preexec_fn=os.setsid,
                )
            except Exception:
                os.close(master_fd)
                raise
            finally:
                os.close(slave_fd)
            self._master_fd = master_fd

        self._process = process
        self.pid = process.pid
        self._reader = threading.Thread(target=self._reader_loop, daemon=True)
        self._waiter = threading.Thread(target=self._waiter_loop, daemon=True)
        self._reader.start()
        self._waiter.start()

    def summary(self) -> dict[str, object]:
        with self._lock:
            return {
                "shell_id": self.shell_id,
                "cwd": self.cwd,
                "cols": self.cols,
                "rows": self.rows,
                "state": self.state,
                "started_at": self.started_at,
                "ended_at": self.ended_at,
                "exit_code": self.exit_code,
                "pid": self.pid,
            }

    def read(self, *, cursor: int, max_chars: int) -> dict[str, object]:
        with self._lock:
            redacted_output = _redact_text(self._output, self._redact_values)
            clamped_cursor = min(max(cursor, 0), len(redacted_output))
            output = redacted_output[clamped_cursor : clamped_cursor + max_chars]
            next_cursor = clamped_cursor + len(output)
            payload = self.summary()
            payload.update(
                {
                    "cursor": clamped_cursor,
                    "next_cursor": next_cursor,
                    "output": output,
                    "truncated": next_cursor < len(redacted_output),
                }
            )
            return payload

    def write(self, text: str, *, append_newline: bool) -> dict[str, object]:
        if self._process.poll() is not None:
            self._refresh_process_state()
        with self._lock:
            if self.state != "running":
                raise RuntimeError(f"shell {self.shell_id} is not running")
            master_fd = self._master_fd
            input_pipe = self._input_pipe
        payload = text + ("\n" if append_newline else "")
        try:
            if master_fd is not None:
                os.write(master_fd, payload.encode("utf-8"))
            else:
                if input_pipe is None:
                    raise RuntimeError(f"shell {self.shell_id} transport is unavailable")
                input_pipe.write(payload.encode("utf-8"))
                input_pipe.flush()
        except OSError as exc:
            self._refresh_process_state()
            raise RuntimeError(f"failed to write to shell {self.shell_id}: {exc}") from exc
        result = self.summary()
        result.update({"input_length": len(text), "append_newline": append_newline})
        return result

    def send_signal(self, signal_name: str) -> dict[str, object]:
        signal_name = signal_name.upper()
        signum = _SHELL_SIGNAL_MAP.get(signal_name)
        if signum is None:
            raise ValueError(f"unsupported shell signal: {signal_name}")
        if self._process.poll() is not None:
            self._refresh_process_state()
        with self._lock:
            if self.state != "running" or self.pid is None:
                raise RuntimeError(f"shell {self.shell_id} is not running")
            pid = self.pid
        try:
            os.killpg(pid, signum)
        except ProcessLookupError as exc:
            self._refresh_process_state()
            raise RuntimeError(f"shell {self.shell_id} is not running") from exc
        result = self.summary()
        result["signal"] = signal_name
        return result

    def close(self) -> dict[str, object]:
        if self._process.poll() is None and self.pid is not None:
            try:
                os.killpg(self.pid, signal.SIGHUP)
            except ProcessLookupError:
                pass
            try:
                self._process.wait(timeout=5)
            except subprocess.TimeoutExpired:
                try:
                    os.killpg(self.pid, signal.SIGKILL)
                except ProcessLookupError:
                    pass
                self._process.wait(timeout=5)
        else:
            self._refresh_process_state()
        self._close_master_fd()
        if self._reader is not None:
            self._reader.join(timeout=1)
        if self._waiter is not None:
            self._waiter.join(timeout=1)
        result = self.summary()
        result["closed"] = True
        return result

    def _reader_loop(self) -> None:
        master_fd = self._master_fd
        output_pipe = self._output_pipe
        if master_fd is None and output_pipe is None:
            return
        while True:
            try:
                if master_fd is not None:
                    chunk = os.read(master_fd, 65536)
                else:
                    if output_pipe is None:
                        break
                    chunk = os.read(output_pipe.fileno(), 65536)
            except OSError:
                break
            if chunk == b"":
                break
            decoded = self._decoder.decode(chunk)
            if decoded:
                with self._lock:
                    self._output += decoded
        decoded = self._decoder.decode(b"", final=True)
        if decoded:
            with self._lock:
                self._output += decoded

    def _waiter_loop(self) -> None:
        exit_code = self._process.wait()
        with self._lock:
            self.state = "stopped"
            self.exit_code = exit_code
            self.ended_at = time.time()

    def _refresh_process_state(self) -> None:
        exit_code = self._process.poll()
        if exit_code is None:
            return
        with self._lock:
            if self.state == "running":
                self.state = "stopped"
                self.exit_code = exit_code
                self.ended_at = time.time()

    def _close_master_fd(self) -> None:
        with self._lock:
            master_fd = self._master_fd
            self._master_fd = None
            input_pipe = self._input_pipe
            self._input_pipe = None
            output_pipe = self._output_pipe
            self._output_pipe = None
        if input_pipe is not None:
            input_pipe.close()
        if output_pipe is not None:
            output_pipe.close()
        if master_fd is None:
            return
        try:
            os.close(master_fd)
        except OSError:
            pass


def create_local_shell(
    *,
    workspace_id: str,
    shell_id: str,
    cwd: Path,
    display_cwd: str,
    cols: int,
    rows: int,
    env_overrides: dict[str, str] | None = None,
    redact_values: list[str] | None = None,
) -> LocalShellSession:
    session_key = f"{workspace_id}:{shell_id}"
    with _LOCAL_SHELLS_LOCK:
        if session_key in _LOCAL_SHELLS:
            raise RuntimeError(f"shell {shell_id} already exists in workspace {workspace_id}")
        session = LocalShellSession(
            shell_id=shell_id,
            cwd=cwd,
            display_cwd=display_cwd,
            cols=cols,
            rows=rows,
            env_overrides=env_overrides,
            redact_values=redact_values,
        )
        _LOCAL_SHELLS[session_key] = session
        return session


def get_local_shell(*, workspace_id: str, shell_id: str) -> LocalShellSession:
    session_key = f"{workspace_id}:{shell_id}"
    with _LOCAL_SHELLS_LOCK:
        try:
            return _LOCAL_SHELLS[session_key]
        except KeyError as exc:
            raise ValueError(
                f"shell {shell_id!r} does not exist in workspace {workspace_id!r}"
            ) from exc


def remove_local_shell(*, workspace_id: str, shell_id: str) -> LocalShellSession | None:
    session_key = f"{workspace_id}:{shell_id}"
    with _LOCAL_SHELLS_LOCK:
        return _LOCAL_SHELLS.pop(session_key, None)


def shell_signal_names() -> tuple[str, ...]:
    return SHELL_SIGNAL_NAMES


def shell_signal_arg_help() -> str:
    return ", ".join(shlex.quote(name) for name in SHELL_SIGNAL_NAMES)
