from __future__ import annotations

import os
import pty
import select
import shutil
import signal
import subprocess
import sys
import termios
import tty
from collections.abc import Callable
from dataclasses import dataclass

from dxcli_sessions import TMUX_WHEEL_SCROLL_SPEED

FilterHandler = Callable[[str], None]

CURSOR_WHEEL_SEQUENCES = (
    (b"\x1b[A", "up"),
    (b"\x1b[B", "down"),
)
MAX_SEQUENCE_LENGTH = max(len(sequence) for sequence, _direction in CURSOR_WHEEL_SEQUENCES)
ESCAPE_FLUSH_SECONDS = 0.03
CURSOR_WHEEL_COALESCE_SECONDS = 0.16


@dataclass
class PendingCursorSequence:
    sequence: bytes
    direction: str
    received_at: float


class TerminalInputFilter:
    def __init__(
        self,
        handler: FilterHandler,
        *,
        intercept_cursor_wheel: bool = False,
        clock: Callable[[], float] | None = None,
    ) -> None:
        self._handler = handler
        self._intercept_cursor_wheel = intercept_cursor_wheel
        self._clock = clock or _monotonic
        self._buffer = bytearray()
        self._pending_cursor: PendingCursorSequence | None = None
        self._cursor_wheel_stream_until = -1.0

    def feed(self, data: bytes) -> bytes:
        now = self._clock()
        output = b""
        self._buffer.extend(data)
        output += self._drain(flush=False, now=now)
        return output

    def flush(self) -> bytes:
        return self.flush_pending_cursor(now=self._clock()) + self._drain(flush=True, now=self._clock())

    def flush_pending_cursor(self, *, now: float | None = None, force: bool = False) -> bytes:
        pending = self._pending_cursor
        if pending is None:
            return b""
        current_time = self._clock() if now is None else now
        if not force and current_time - pending.received_at < CURSOR_WHEEL_COALESCE_SECONDS:
            return b""
        self._pending_cursor = None
        return pending.sequence

    def has_pending_cursor(self) -> bool:
        return self._pending_cursor is not None

    def pending_cursor_age(self, *, now: float | None = None) -> float | None:
        pending = self._pending_cursor
        if pending is None:
            return None
        current_time = self._clock() if now is None else now
        return current_time - pending.received_at

    def _drain(self, *, flush: bool, now: float) -> bytes:
        output = bytearray()
        while self._buffer:
            mouse_result = self._match_sgr_mouse()
            if mouse_result == "partial" and not flush:
                break
            if mouse_result in {"up", "down"}:
                self._handler(mouse_result)
                continue
            if isinstance(mouse_result, bytes):
                output.extend(mouse_result)
                continue

            matched = False
            waiting_for_more = False
            intercept_sequences = CURSOR_WHEEL_SEQUENCES if self._intercept_cursor_wheel else ()
            for sequence, direction in intercept_sequences:
                if self._buffer.startswith(sequence):
                    del self._buffer[: len(sequence)]
                    if now < self._cursor_wheel_stream_until:
                        self._handler(direction)
                        self._cursor_wheel_stream_until = now + CURSOR_WHEEL_COALESCE_SECONDS
                    elif self._pending_cursor is None and not flush:
                        self._pending_cursor = PendingCursorSequence(sequence, direction, now)
                    elif self._pending_cursor is not None:
                        self._handler(self._pending_cursor.direction)
                        self._handler(direction)
                        self._pending_cursor = PendingCursorSequence(sequence, direction, now)
                        self._pending_cursor = None
                        self._cursor_wheel_stream_until = now + CURSOR_WHEEL_COALESCE_SECONDS
                    else:
                        output.extend(sequence)
                    matched = True
                    break
                if sequence.startswith(self._buffer):
                    waiting_for_more = True
            if matched:
                continue
            if waiting_for_more and not flush:
                break

            output.append(self._buffer.pop(0))
            if len(self._buffer) > MAX_SEQUENCE_LENGTH:
                output.extend(self._buffer[:-MAX_SEQUENCE_LENGTH])
                del self._buffer[:-MAX_SEQUENCE_LENGTH]
        return bytes(output)

    def _match_sgr_mouse(self) -> str | bytes | None:
        prefix = b"\x1b[<"
        if not self._buffer.startswith(prefix):
            return None
        terminator_index = -1
        for index in range(len(prefix), len(self._buffer)):
            if self._buffer[index] in (ord("M"), ord("m")):
                terminator_index = index
                break
        if terminator_index < 0:
            return "partial"

        raw = bytes(self._buffer[: terminator_index + 1])
        del self._buffer[: terminator_index + 1]
        try:
            code = int(raw[len(prefix):].split(b";", 1)[0])
        except ValueError:
            return raw
        wheel_code = code & 0b11
        if code >= 64 and wheel_code == 0:
            return "up"
        if code >= 64 and wheel_code == 1:
            return "down"
        return raw


def run_with_terminal_filter(command: list[str]) -> int:
    if not _filter_enabled():
        return subprocess.call(command)
    if not sys.stdin.isatty() or not sys.stdout.isatty():
        return subprocess.call(command)

    pid, master_fd = pty.fork()
    if pid == 0:
        try:
            os.execvpe(command[0], command, os.environ)
        except FileNotFoundError:
            os._exit(127)

    stdin_fd = sys.stdin.fileno()
    stdout_fd = sys.stdout.fileno()
    previous_termios = termios.tcgetattr(stdin_fd)
    previous_winch = signal.getsignal(signal.SIGWINCH)
    last_buffered_at: float | None = None

    def scroll_handler(direction: str) -> None:
        _scroll_tmux(direction)

    input_filter = TerminalInputFilter(scroll_handler, intercept_cursor_wheel=_intercept_cursor_wheel())

    def resize_child(_signum: int | None = None, _frame: object | None = None) -> None:
        _copy_window_size(stdin_fd, master_fd)

    try:
        resize_child()
        signal.signal(signal.SIGWINCH, resize_child)
        tty.setraw(stdin_fd)
        while True:
            try:
                child_pid, status = os.waitpid(pid, os.WNOHANG)
            except ChildProcessError:
                return 0
            if child_pid == pid:
                return _wait_status_to_exit_code(status)

            timeout = _next_timeout(input_filter, last_buffered_at)
            readable, _writable, _errors = select.select([stdin_fd, master_fd], [], [], timeout)
            if stdin_fd in readable:
                data = os.read(stdin_fd, 4096)
                if not data:
                    break
                forwarded = input_filter.feed(data)
                last_buffered_at = _monotonic() if not forwarded and data.startswith(b"\x1b") else None
                if forwarded:
                    os.write(master_fd, forwarded)
            else:
                now = _monotonic()
                if input_filter.has_pending_cursor():
                    forwarded = input_filter.flush_pending_cursor(now=now)
                    if forwarded:
                        os.write(master_fd, forwarded)
                if last_buffered_at is not None and now - last_buffered_at >= ESCAPE_FLUSH_SECONDS:
                    forwarded = input_filter.flush()
                    last_buffered_at = None
                    if forwarded:
                        os.write(master_fd, forwarded)
            if master_fd in readable:
                try:
                    output = os.read(master_fd, 4096)
                except OSError:
                    break
                if not output:
                    break
                os.write(stdout_fd, output)
    finally:
        termios.tcsetattr(stdin_fd, termios.TCSADRAIN, previous_termios)
        signal.signal(signal.SIGWINCH, previous_winch)
        try:
            os.close(master_fd)
        except OSError:
            pass

    _, status = os.waitpid(pid, 0)
    return _wait_status_to_exit_code(status)


def _filter_enabled() -> bool:
    raw = os.getenv("DXCLI_TERMINAL_INPUT_FILTER", "1").strip().lower()
    return raw not in {"0", "false", "no", "off"}


def _intercept_cursor_wheel() -> bool:
    raw = os.getenv("DXCLI_FINALSHELL_CURSOR_WHEEL_FILTER", "1").strip().lower()
    return raw not in {"0", "false", "no", "off"}


def _scroll_tmux(direction: str) -> None:
    tmux_bin = shutil.which("tmux")
    pane = os.getenv("TMUX_PANE", "").strip()
    if not tmux_bin or not pane:
        return
    if direction == "up":
        subprocess.run(
            [
                tmux_bin,
                "copy-mode",
                "-e",
                "-t",
                pane,
                ";",
                "send-keys",
                "-X",
                "-t",
                pane,
                "-N",
                str(TMUX_WHEEL_SCROLL_SPEED),
                "scroll-up",
            ],
            capture_output=True,
            text=True,
            check=False,
        )
        return

    if not _pane_in_mode(tmux_bin, pane):
        return
    subprocess.run(
        [
            tmux_bin,
            "send-keys",
            "-X",
            "-t",
            pane,
            "-N",
            str(TMUX_WHEEL_SCROLL_SPEED),
            "scroll-down",
        ],
        capture_output=True,
        text=True,
        check=False,
    )
    if _scroll_position(tmux_bin, pane) == "0":
        subprocess.run(
            [tmux_bin, "send-keys", "-X", "-t", pane, "cancel"],
            capture_output=True,
            text=True,
            check=False,
        )


def _next_timeout(
    input_filter: TerminalInputFilter,
    last_buffered_at: float | None,
) -> float:
    candidates = [ESCAPE_FLUSH_SECONDS]
    now = _monotonic()
    if input_filter.has_pending_cursor():
        age = input_filter.pending_cursor_age(now=now) or 0.0
        candidates.append(max(0.0, CURSOR_WHEEL_COALESCE_SECONDS - age))
    if last_buffered_at is not None:
        candidates.append(max(0.0, ESCAPE_FLUSH_SECONDS - (now - last_buffered_at)))
    return min(candidates)


def _pane_in_mode(tmux_bin: str, pane: str) -> bool:
    result = subprocess.run(
        [tmux_bin, "display-message", "-p", "-t", pane, "#{pane_in_mode}"],
        capture_output=True,
        text=True,
        check=False,
    )
    return result.stdout.strip() == "1"


def _scroll_position(tmux_bin: str, pane: str) -> str:
    result = subprocess.run(
        [tmux_bin, "display-message", "-p", "-t", pane, "#{scroll_position}"],
        capture_output=True,
        text=True,
        check=False,
    )
    return result.stdout.strip()


def _copy_window_size(source_fd: int, target_fd: int) -> None:
    try:
        import fcntl
        import struct

        size = fcntl.ioctl(source_fd, termios.TIOCGWINSZ, b"\0" * 8)
        rows, columns, x_pixels, y_pixels = struct.unpack("HHHH", size)
        if rows and columns:
            fcntl.ioctl(target_fd, termios.TIOCSWINSZ, struct.pack("HHHH", rows, columns, x_pixels, y_pixels))
    except OSError:
        pass


def _wait_status_to_exit_code(status: int) -> int:
    if os.WIFEXITED(status):
        return os.WEXITSTATUS(status)
    if os.WIFSIGNALED(status):
        return 128 + os.WTERMSIG(status)
    return 1


def _monotonic() -> float:
    import time

    return time.monotonic()
