#!/usr/bin/env python3
"""Paste Finder files into an SSH terminal through an existing UU TCP mapping."""

import argparse
import collections
import errno
import fcntl
import json
import os
from pathlib import Path
import pty
import queue
import re
import select
import shlex
import signal
import stat
import struct
import subprocess
import sys
import termios
import threading
import time
import tty

START, END = b"\x1b[200~", b"\x1b[201~"
ENABLE, DISABLE = b"\x1b[?2004h", b"\x1b[?2004l"
MAX_PASTE = 8 * 1024 * 1024


class PasteError(Exception):
    """Messages must not include clipboard contents, paths or subprocess output."""


def prefix_tail(data, marker):
    for length in range(min(len(data), len(marker) - 1), 0, -1):
        if data.endswith(marker[:length]):
            return length
    return 0


class InputFrames:
    """Recognize bracketed pastes across arbitrary PTY read boundaries."""

    def __init__(self):
        self.buffer = b""
        self.pasting = False
        self.overflow = False

    def feed(self, data):
        self.buffer += data
        events = []
        while self.buffer:
            if self.pasting:
                end = self.buffer.find(END)
                if end < 0:
                    if len(self.buffer) > MAX_PASTE:
                        self.overflow = True
                        tail = prefix_tail(self.buffer, END)
                        self.buffer = self.buffer[-tail:] if tail else b""
                    break
                payload, self.buffer = self.buffer[:end], self.buffer[end + len(END):]
                events.append(("error" if self.overflow or len(payload) > MAX_PASTE else "paste", payload))
                self.pasting = self.overflow = False
            else:
                start = self.buffer.find(START)
                if start >= 0:
                    if start:
                        events.append(("data", self.buffer[:start]))
                    self.buffer = self.buffer[start + len(START):]
                    self.pasting = True
                else:
                    tail = prefix_tail(self.buffer, START)
                    size = len(self.buffer) - tail
                    if size:
                        events.append(("data", self.buffer[:size]))
                        self.buffer = self.buffer[size:]
                    break
        return events

    def flush_key(self):
        # A standalone Escape must not wait indefinitely for another key.
        if not self.pasting:
            result, self.buffer = self.buffer, b""
            return result
        return b""


class OutputModes:
    """Keep local paste framing on; retain the remote app's requested mode."""

    def __init__(self):
        self.buffer = b""
        self.bracketed = False

    def feed(self, data):
        self.buffer += data
        output = bytearray()
        while self.buffer:
            index = self.buffer.find(b"\x1b")
            if index < 0:
                output.extend(self.buffer)
                self.buffer = b""
                break
            output.extend(self.buffer[:index])
            self.buffer = self.buffer[index:]
            if self.buffer.startswith(b"\x1bc"):
                output.extend(b"\x1bc" + ENABLE)
                self.bracketed = False
                self.buffer = self.buffer[2:]
                continue
            match = re.match(rb"\x1b\[\?([0-9;]+)([hl])", self.buffer)
            if match:
                params = match[1].split(b";")
                if b"2004" in params:
                    self.bracketed = match[2] == b"h"
                    others = [p for p in params if p != b"2004"]
                    if others:
                        output.extend(b"\x1b[?" + b";".join(others) + match[2])
                    output.extend(ENABLE)
                else:
                    output.extend(match[0])
                self.buffer = self.buffer[match.end():]
            elif len(self.buffer) < 128 and re.fullmatch(rb"\x1b(?:\[(?:\?[0-9;]*)?)?", self.buffer):
                break
            else:
                output.extend(self.buffer[:1])
                self.buffer = self.buffer[1:]
        return bytes(output)


def paste_bytes(data, bracketed):
    return START + data + END if bracketed else data


def clipboard_paths(helper):
    try:
        result = subprocess.run([str(helper)], capture_output=True, timeout=3, check=True)
        paths = json.loads(result.stdout)
        if not isinstance(paths, list) or not all(isinstance(p, str) for p in paths):
            raise ValueError()
        return paths
    except (OSError, ValueError, subprocess.SubprocessError):
        raise PasteError("Cannot check copied files; paste was not sent.") from None


def copied_files(data, read_clipboard):
    try:
        text = data.decode("utf-8")
        paths = shlex.split(text)
    except (UnicodeError, ValueError):
        return None
    if not paths or not all(p.startswith("/") for p in paths):
        return None
    # Ordinary text that happens to name a local file must never upload it.
    # Ghostty 1.3.1 getOpinionatedStringContents escapes Finder file paths;
    # require the entire pasted path list to match actual file-URL clipboard data.
    if paths != read_clipboard():
        return None
    if len(paths) > 256:
        raise PasteError("Too many copied files (maximum 256).")
    for path in paths:
        if any(ord(c) < 32 or ord(c) == 127 for c in path):
            raise PasteError("A copied filename contains unsupported control characters.")
        try:
            if not stat.S_ISREG(os.stat(path).st_mode):
                raise PasteError("Copy files individually; folder pasting is not supported.")
        except OSError:
            raise PasteError("A copied file is no longer readable.") from None
    if len({Path(p).name for p in paths}) != len(paths):
        raise PasteError("Copied files must have distinct filenames; paste them separately.")
    return paths


class SCPMeter:
    """Reduce OpenSSH's private progress output to numeric batch counters."""

    # OpenSSH progressmeter.c writes: filename percent size rate time [ETA].
    # Anchor at the numeric suffix; filenames and errors are never displayed.
    record = re.compile(rb" (\d{1,3})%\s+(\d{1,20}(?:[KMGT]B)?)\s+"
                        rb"(\d{1,20}(?:\.\d{1,10})?[KMGT]B)/s\s+"
                        rb"(?:\d+:\d\d(?::\d\d)?(?: ETA)?|--:-- ETA|- stalled -)\s*$")

    def __init__(self, sizes, report):
        self.sizes, self.report = sizes, report
        self.buffer, self.index, self.finished = b"", 0, 0
        self.current, self.percent = 0, 0

    @staticmethod
    def amount(token):
        match = re.fullmatch(rb"(\d+(?:\.\d+)?)([KMGT])?B?", token)
        return int(float(match[1]) * 1024 ** (b" KMGT".index(match[2]) if match[2] else 0))

    def parse(self, line):
        match = self.record.search(line)
        if not match or self.index >= len(self.sizes):
            return
        percent = int(match[1])
        if percent > 100:
            return
        self.percent = max(self.percent, percent)
        size = self.sizes[self.index]
        self.current = max(self.current, min(size, self.amount(match[2])))
        if percent == 100:
            self.current = size
        self.report("uploading", self.finished + self.current, sum(self.sizes), self.amount(match[3]))

    def feed(self, data):
        for piece in re.split(rb"([\r\n])", data):
            if piece in (b"\r", b"\n"):
                self.parse(self.buffer)
                self.buffer = b""
                if piece == b"\n" and self.percent == 100:
                    self.finished += self.sizes[self.index]
                    self.index += 1
                    self.current = self.percent = 0
            else:
                # Keep only enough for a numeric suffix, not unlimited errors.
                self.buffer = (self.buffer + piece)[-512:]
        self.parse(self.buffer)


def run_metered(args, cancel, timeout, meter):
    master, slave = pty.openpty()
    fcntl.ioctl(slave, termios.TIOCSWINSZ, struct.pack("HHHH", 24, 120, 0, 0))

    try:
        # The upload runs on a worker thread, so avoid Python preexec_fn after
        # fork. Acquire the private controlling TTY in a fresh interpreter.
        launch = ("import fcntl,os,sys,termios;fcntl.ioctl(1,termios.TIOCSCTTY,0);"
                  "os.execv(sys.argv[1],sys.argv[1:])")
        with subprocess.Popen([sys.executable, "-c", launch, *args],
                              stdin=subprocess.DEVNULL, stdout=slave, stderr=slave,
                              start_new_session=True) as process:
            os.close(slave)
            slave = None
            deadline = time.monotonic() + timeout
            try:
                while True:
                    if cancel.is_set() or time.monotonic() > deadline:
                        raise PasteError("Upload cancelled or timed out; paste was not sent.")
                    if select.select([master], [], [], 0.1)[0]:
                        try:
                            data = os.read(master, 65536)
                        except OSError as error:
                            if error.errno != errno.EIO:
                                raise
                            data = b""
                        if not data:
                            break
                        meter.feed(data)
                    elif process.poll() is not None:
                        break
                try:
                    status = process.wait(timeout=max(0.1, deadline - time.monotonic()))
                except subprocess.TimeoutExpired:
                    raise PasteError("Upload timed out; paste was not sent.") from None
                if status:
                    raise PasteError("Upload failed; paste was not sent.")
            finally:
                if process.poll() is None:
                    os.killpg(process.pid, signal.SIGKILL)
                    process.wait()
    finally:
        os.close(master)
        if slave is not None:
            os.close(slave)
    return b""


def run_private(args, cancel, timeout, meter=None):
    # Neither SSH errors nor SCP output may disclose paths, clipboard/file data,
    # or credentials. Authentication never reads from the interactive terminal.
    if meter is not None:
        return run_metered(args, cancel, timeout, meter)
    with subprocess.Popen(args, stdin=subprocess.DEVNULL, stdout=subprocess.PIPE,
                          stderr=subprocess.DEVNULL, start_new_session=True) as process:
        deadline = time.monotonic() + timeout
        while True:
            if cancel.is_set() or time.monotonic() > deadline:
                os.killpg(process.pid, signal.SIGKILL)
                process.communicate()
                raise PasteError("Upload cancelled." if cancel.is_set() else "Upload timed out; paste was not sent.")
            try:
                stdout, _ = process.communicate(timeout=0.1)
                if process.returncode:
                    raise PasteError("Upload failed; check the UU mapping and SSH access. Paste was not sent.")
                return stdout
            except subprocess.TimeoutExpired:
                pass


class Uploader:
    def __init__(self, config, host):
        common = ["-F", str(config)] if config else []
        common += ["-o", "BatchMode=yes", "-o", "ConnectTimeout=15"]
        self.ssh = ["/usr/bin/ssh", *common, host]
        self.scp = ["/usr/bin/scp", "-q", *common]
        self.host = host

    def upload(self, paths, cancel, progress=None):
        sizes = [os.stat(path).st_size for path in paths] if progress else []
        if progress:
            progress("connecting", 0, sum(sizes), 0)
        script = ('umask 077; d="$HOME/.local/share/uuremote/uploads"; '
                  'mkdir -p "$d" && mktemp -d "$d/paste.XXXXXXXX"')
        try:
            directory = run_private([*self.ssh, script], cancel, 30).decode("utf-8").strip()
        except UnicodeError:
            raise PasteError("Invalid upload destination; paste was not sent.") from None
        if (not directory.startswith("/") or len(directory) > 4096 or
                any(ord(c) < 32 or ord(c) == 127 for c in directory) or
                not re.fullmatch(r"paste\.[A-Za-z0-9]+", directory.rsplit("/", 1)[-1])):
            raise PasteError("Invalid upload destination; paste was not sent.")
        try:
            # Use modern SCP's SFTP mode, never legacy SCP (-O). Unique batch
            # directories prevent overwrites; all local paths are absolute.
            if progress:
                progress("uploading", 0, sum(sizes), 0)
                scp = self.scp[:1] + self.scp[2:]  # Remove only our own -q.
                run_private([*scp, *paths, f"{self.host}:{directory}/"], cancel, 300,
                            SCPMeter(sizes, progress))
                progress("finishing", sum(sizes), sum(sizes), 0)
            else:
                run_private([*self.scp, *paths, f"{self.host}:{directory}/"], cancel, 300)
        except (PasteError, OSError):
            # Clean only this batch, with quoted remote shell syntax. Cancellation
            # of one upload must not stop another SSH/UU session.
            try:
                run_private([*self.ssh, "rm -rf -- " + shlex.quote(directory)], threading.Event(), 10)
            except (PasteError, OSError):
                pass
            raise
        return " ".join(shlex.quote(directory + "/" + Path(p).name) for p in paths).encode("utf-8")


def write_all(fd, data):
    view = memoryview(data)
    while view:
        view = view[os.write(fd, view):]


def display_size(size):
    for unit in ("B", "KiB", "MiB", "GiB", "TiB"):
        if size < 1024 or unit == "TiB":
            return f"{size:.0f} {unit}" if unit == "B" else f"{size:.1f} {unit}"
        size /= 1024


class UploadDisplay:
    """Session-local title and OSC 9;4 progress; terminal cells stay untouched."""

    # https://ghostty.org/docs/vt/osc/conemu: states 0=clear, 1=percent, 3=busy.

    def __init__(self, emit):
        self.emit, self.active = emit, False
        self.title, self.pending = b"UU SSH", b""
        self.state = ("connecting", 0, 0, 0)
        self.started = self.rendered = 0

    def observe(self, data):
        self.pending = (self.pending + data)[-8192:]
        for match in re.finditer(rb"\x1b\][02];([^\x07\x1b]*)(?:\x07|\x1b\\)", self.pending):
            self.title = match[1]
        # Retain only an incomplete OSC title across output reads.
        index = self.pending.rfind(b"\x1b]")
        if index < 0:
            self.pending = b"\x1b" if self.pending.endswith(b"\x1b") else b""
        else:
            tail = self.pending[index:]
            self.pending = tail if b"\x07" not in tail and b"\x1b\\" not in tail else b""

    def start(self, total):
        self.active = True
        self.started = time.monotonic()
        self.state = ("connecting", 0, total, 0)
        self.render(force=True)

    def update(self, stage, done, total, speed):
        self.state = stage, done, total, speed

    def render(self, force=False):
        now = time.monotonic()
        if not self.active or (not force and now - self.rendered < 0.2):
            return
        self.rendered = now
        stage, done, total, speed = self.state
        elapsed = int(now - self.started)
        if stage in ("connecting", "cancelling"):
            label = "Connecting" if stage == "connecting" else "Cancelling"
            text = f"UU upload: {label} | {display_size(total)} | {elapsed}s"
            indicator = b"\x1b]9;4;3\x1b\\"
        else:
            percent = min(100, int(done * 100 / total)) if total else 100
            label = "Finishing" if stage == "finishing" else "Uploading"
            text = (f"UU upload: {label} {percent}% | {display_size(done)}/{display_size(total)}"
                    f" | {display_size(speed)}/s | {elapsed}s")
            indicator = f"\x1b]9;4;1;{percent}\x1b\\".encode()
        self.emit(b"\x1b]2;" + text.encode() + b"\x1b\\" + indicator)

    def finish(self):
        if self.active:
            self.active = False
            self.emit(b"\x1b]9;4;0\x1b\\\x1b]2;" + self.title + b"\x1b\\")


def relay(command, read_clipboard, uploader):
    original = termios.tcgetattr(0)
    master, slave = pty.openpty()

    def resize(*_):
        try:
            size = fcntl.ioctl(0, termios.TIOCGWINSZ, bytes(8))
            fcntl.ioctl(master, termios.TIOCSWINSZ, size)
        except OSError:
            pass

    resize()

    def child_tty():
        os.setsid()
        fcntl.ioctl(0, termios.TIOCSCTTY, 0)

    child = subprocess.Popen(command, stdin=slave, stdout=slave, stderr=slave, preexec_fn=child_tty)
    os.close(slave)
    previous_resize = signal.signal(signal.SIGWINCH, resize)

    def shutdown(signum, _frame):
        raise SystemExit(128 + signum)

    previous_signals = {sig: signal.signal(sig, shutdown)
                        for sig in (signal.SIGHUP, signal.SIGTERM, signal.SIGINT)}
    frames, modes = InputFrames(), OutputModes()
    events, completed = collections.deque(), queue.Queue()
    updates = queue.Queue()
    display = UploadDisplay(lambda data: write_all(1, data))
    cancel, uploading, interrupted = threading.Event(), False, False
    worker = None
    last_input = time.monotonic()

    def upload_worker(paths):
        try:
            completed.put((uploader.upload(paths, cancel, lambda *args: updates.put(args)), None))
        except (OSError, PasteError):
            completed.put((None, "Upload failed or cancelled; paste was not sent."))

    def notice(message):
        write_all(2, ("\r\n[UU paste] " + message + "\r\n").encode())

    def interrupt_upload():
        nonlocal interrupted, frames
        if not uploading:
            return
        pending = list(events)
        positions = [(i, payload.rfind(b"\x03")) for i, (kind, payload) in enumerate(pending)
                     if kind == "data" and b"\x03" in payload]
        if positions:
            index, offset = positions[-1]
            cancel.set()
            interrupted = True
            _, done, total, _ = display.state
            display.update("cancelling", done, total, 0)
            display.render(force=True)
            frames = InputFrames()
            events.clear()
            events.extend(frames.feed(pending[index][1][offset + 1:]))
            events.extend(pending[index + 1:])
            write_all(master, b"\x03")

    try:
        tty.setraw(0)
        write_all(1, ENABLE)
        while True:
            interrupt_upload()
            while not updates.empty():
                update = updates.get_nowait()
                if not interrupted:
                    display.update(*update)
            if uploading:
                display.render()
                try:
                    data, error = completed.get_nowait()
                except queue.Empty:
                    pass
                else:
                    uploading = False
                    display.finish()
                    if error:
                        notice(error)
                    elif not interrupted:
                        write_all(master, paste_bytes(data, modes.bracketed))
            while events and not uploading:
                kind, data = events.popleft()
                if kind == "error":
                    notice("Paste exceeds 8 MiB; paste was not sent.")
                    continue
                if kind == "data":
                    write_all(master, data)
                    continue
                try:
                    files = copied_files(data, read_clipboard)
                except PasteError as error:
                    notice(str(error))
                    continue
                if files is None:
                    write_all(master, paste_bytes(data, modes.bracketed))
                else:
                    try:
                        total = sum(os.stat(path).st_size for path in files)
                    except OSError:
                        notice("A copied file is no longer readable; paste was not sent.")
                        continue
                    cancel.clear()
                    uploading, interrupted = True, False
                    display.start(total)
                    worker = threading.Thread(target=upload_worker, args=(files,), daemon=True)
                    worker.start()
            interrupt_upload()
            readable, _, _ = select.select([0, master], [], [], 0.025)
            if master in readable:
                try:
                    data = os.read(master, 65536)
                except OSError as error:
                    if error.errno != errno.EIO:
                        raise
                    data = b""
                if not data:
                    break
                display.observe(data)
                write_all(1, modes.feed(data))
            if 0 in readable:
                data = os.read(0, 65536)
                if not data:
                    break
                last_input = time.monotonic()
                events.extend(frames.feed(data))
                if sum(len(payload) for _, payload in events) > MAX_PASTE:
                    events.clear()
                    notice("Input queue is full; queued input was discarded.")
            elif time.monotonic() - last_input > 0.025:
                data = frames.flush_key()
                if data:
                    events.append(("data", data))
        if modes.buffer:
            write_all(1, modes.buffer)
    finally:
        cancel.set()
        signal.signal(signal.SIGWINCH, previous_resize)
        for sig, previous in previous_signals.items():
            signal.signal(sig, previous)
        termios.tcsetattr(0, termios.TCSADRAIN, original)
        display.finish()
        write_all(1, DISABLE)
        os.close(master)
        if child.poll() is None:
            try:
                child.wait(timeout=2)
            except subprocess.TimeoutExpired:
                os.killpg(child.pid, signal.SIGTERM)
                try:
                    child.wait(timeout=2)
                except subprocess.TimeoutExpired:
                    os.killpg(child.pid, signal.SIGKILL)
        status = child.wait()
        if worker is not None:
            worker.join(timeout=12)
    return status if status >= 0 else 128 - status


def main():
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("-F", "--ssh-config", type=Path)
    parser.add_argument("host", help="SSH alias for an existing UU loopback port mapping")
    parser.add_argument("command", nargs=argparse.REMAINDER)
    args = parser.parse_args()
    if not re.fullmatch(r"[A-Za-z0-9][A-Za-z0-9_.@-]*", args.host):
        parser.error("use an SSH host alias, without options or a port")
    common = ["-F", str(args.ssh_config)] if args.ssh_config else []
    command = ["/usr/bin/ssh", *common]
    if not (os.isatty(0) and os.isatty(1)):
        os.execv(command[0], [*command, args.host, *args.command])
    helper = Path(__file__).resolve().with_name("clipboard-files")
    if not helper.is_file():
        parser.error("run contrib/macos-ssh/install.sh on your Mac first")
    return relay([*command, "-tt", args.host, *args.command],
                 lambda: clipboard_paths(helper), Uploader(args.ssh_config, args.host))


if __name__ == "__main__":
    try:
        sys.exit(main())
    except (OSError, PasteError):
        print("UU SSH companion failed; check the UU mapping and SSH configuration.", file=sys.stderr)
        sys.exit(1)
