#!/usr/bin/env python3
"""Keep browser-owning gptbridge launchers outside an agent tmp jail."""

from __future__ import annotations

import os
import pwd
import shutil
import subprocess
import sys
import tempfile
from pathlib import Path

MARKER = "GPTBRIDGE_TRUSTED_RELAUNCH"
JAIL_MARKER = "TMPJAIL_ACTIVE"
TARGETS = {"ask-gpt": "ask_gpt.py", "solwebd": "solwebd.py"}
ENV_NAMES = {
    "USER", "LOGNAME", "LANG", "LC_ALL", "LC_CTYPE", "TERM", "COLORTERM",
    "NO_COLOR", "FORCE_COLOR", "ASKGPT_DISPATCH_GAP", "SOLWEBD_PORT",
    "SOLWEBD_TOKEN_FILE", "SOLWEBD_TURN_TIMEOUT", "GPTBRIDGE_HOME",
    "GPTBRIDGE_STATE_DIR",
}
SAFE_PATH = "/usr/local/bin:/usr/bin:/bin"
ISOLATED_BOOTSTRAP = (
    "import runpy,sys;root=sys.argv.pop(1);script=sys.argv.pop(1);"
    "sys.path.insert(0,root);sys.argv[0]=script;"
    "runpy.run_path(script,run_name='__main__')"
)


def in_tmp_jail() -> bool:
    return os.environ.get(JAIL_MARKER) == "1" or Path("/tmp/.tmpjail").exists()


def _command(target: str, args: list[str]) -> list[str]:
    script = TARGETS.get(target)
    if script is None:
        raise SystemExit(f"trusted-launch: unsupported target: {target}")
    root = Path(__file__).resolve().parent
    return [str(Path(sys.executable).resolve()), "-I", "-c", ISOLATED_BOOTSTRAP,
            str(root), str(root / script), *args]


def _resume(argv: list[str]) -> None:
    if in_tmp_jail():
        raise SystemExit("trusted-launch: relaunch remained inside tmp jail")
    if len(argv) < 2:
        raise SystemExit("trusted-launch: invalid resume request")
    env_path = Path(argv[0])
    try:
        entries = env_path.read_bytes().split(b"\0")
    except OSError as exc:
        raise SystemExit(f"trusted-launch: cannot restore environment: {exc}") from exc
    finally:
        env_path.unlink(missing_ok=True)
    env = {"HOME": pwd.getpwuid(os.getuid()).pw_dir, "PATH": SAFE_PATH, MARKER: "1"}
    for entry in entries:
        if not entry:
            continue
        name, separator, value = entry.partition(b"=")
        if not separator:
            raise SystemExit("trusted-launch: malformed environment")
        decoded_name = os.fsdecode(name)
        if decoded_name in ENV_NAMES or decoded_name.startswith("LC_"):
            env[decoded_name] = os.fsdecode(value)
    command = _command(argv[1], argv[2:])
    os.execve(command[0], command, env)


def _save_environment(directory: Path) -> Path:
    directory.mkdir(mode=0o700, parents=True, exist_ok=True)
    fd, raw_path = tempfile.mkstemp(prefix="env-", dir=directory)
    path = Path(raw_path)
    try:
        os.fchmod(fd, 0o600)
        with os.fdopen(fd, "wb") as handle:
            names = ENV_NAMES | {key for key in os.environ if key.startswith("LC_")}
            for name in sorted(names):
                if name in os.environ:
                    handle.write(os.fsencode(name) + b"=" + os.fsencode(os.environ[name]) + b"\0")
    except BaseException:
        os.close(fd)
        path.unlink(missing_ok=True)
        raise
    return path


def _stage_attachments(args: list[str], directory: Path) -> tuple[list[str], list[Path]]:
    rewritten = list(args)
    staged: list[Path] = []
    index = 0
    while index < len(rewritten):
        argument = rewritten[index]
        if argument not in {"-a", "--attach"}:
            index += 1
            continue
        if index + 1 >= len(rewritten):
            raise SystemExit(f"trusted-launch: {argument} requires a file")
        source = Path(rewritten[index + 1])
        if not source.is_file():
            raise SystemExit(f"trusted-launch: attachment is not a file: {source}")
        directory.mkdir(mode=0o700, parents=True, exist_ok=True)
        fd, raw_path = tempfile.mkstemp(prefix="attachment-", suffix=source.suffix, dir=directory)
        destination = Path(raw_path)
        try:
            with source.open("rb") as input_file, os.fdopen(fd, "wb") as output_file:
                shutil.copyfileobj(input_file, output_file)
            os.chmod(destination, 0o600)
        except BaseException:
            os.close(fd)
            destination.unlink(missing_ok=True)
            for path in staged:
                path.unlink(missing_ok=True)
            raise
        staged.append(destination)
        rewritten[index + 1] = str(destination)
        index += 2
    return rewritten, staged


def launch(target: str, args: list[str]) -> int:
    command = _command(target, args)
    if not in_tmp_jail():
        return os.execve(command[0], command, os.environ)
    if os.environ.get(MARKER):
        raise SystemExit("trusted-launch: relaunch remained inside tmp jail")
    runtime = Path.home() / ".overdeck" / "gptbridge" / "trusted-launch"
    staged_args, staged_paths = _stage_attachments(args, runtime)
    env_path = _save_environment(runtime)
    mode = "--pty" if all(os.isatty(fd) for fd in (0, 1, 2)) else "--pipe"
    relaunch = [
        "/usr/bin/systemd-run", "--user", "--wait", "--collect", "--quiet",
        "--service-type=exec", "--same-dir", mode,
        str(Path(sys.executable).resolve()), "-I", "-c", ISOLATED_BOOTSTRAP,
        str(Path(__file__).resolve().parent), str(Path(__file__).resolve()),
        "--resume", str(env_path), target, *staged_args,
    ]
    try:
        return subprocess.run(relaunch, check=False).returncode
    except OSError as exc:
        raise SystemExit(f"trusted-launch: relaunch failed: {exc}") from exc
    finally:
        env_path.unlink(missing_ok=True)
        for path in staged_paths:
            path.unlink(missing_ok=True)


def main() -> int:
    if len(sys.argv) > 1 and sys.argv[1] == "--resume":
        _resume(sys.argv[2:])
        return 125
    if len(sys.argv) < 2:
        raise SystemExit("trusted-launch: missing target")
    return launch(sys.argv[1], sys.argv[2:])


if __name__ == "__main__":
    raise SystemExit(main())