"""Standalone guest-only tool installer, submitted through the execution agent.

Importing this module on the host has no side effects. It uses only Python 3.10
stdlib; all package commands, downloads, and filesystem changes run in the guest.
"""

from __future__ import annotations

import fcntl
import hashlib
import json
import os
import pwd
import re
import shlex
import shutil
import stat
import subprocess
import sys
import tarfile
import tempfile
import time
import urllib.request
from pathlib import Path


ROOT = Path("/opt/uvm-tools")
LOCK = Path("/run/uvm-tools.lock")
BIN_DIR = Path("/usr/local/bin")
PROFILE_DIR = Path("/etc/profile.d")
PATH = "/usr/local/sbin:/usr/local/bin:/usr/sbin:/usr/bin:/sbin:/bin"
LOG_LIMIT = 16384
PHP_KEY = "B8DC7E53946656EFBCE4C1DD71DAEAAB4AD4CAB6"


class ToolError(Exception):
    pass


def atomic_write(path, data, mode=0o644):
    descriptor, name = tempfile.mkstemp(prefix=".uvm-tools-", dir=path.parent)
    try:
        with os.fdopen(descriptor, "wb") as stream:
            os.fchmod(stream.fileno(), mode)
            stream.write(data)
            stream.flush()
            os.fsync(stream.fileno())
        os.replace(name, path)
        directory = os.open(path.parent, os.O_RDONLY | os.O_DIRECTORY)
        try:
            os.fsync(directory)
        finally:
            os.close(directory)
    finally:
        Path(name).unlink(missing_ok=True)


def safe_directory(path):
    for parent in reversed((path, *path.parents)):
        if not parent.exists():
            parent.mkdir(mode=0o755)
        metadata = parent.lstat()
        if not stat.S_ISDIR(metadata.st_mode) or metadata.st_uid != os.geteuid() or metadata.st_mode & 0o022:
            raise ToolError(f"tool installation needs a root-owned directory: {parent}")


def regular_file(path):
    metadata = path.lstat()
    if not stat.S_ISREG(metadata.st_mode) or metadata.st_uid != os.geteuid() or metadata.st_nlink != 1 or metadata.st_mode & 0o022:
        raise ToolError(f"refusing to replace an unmanaged tool file: {path}")


def download(url, destination, digest=None, *, max_bytes=1024 * 1024 * 1024):
    if not url.startswith("https://"):
        raise ToolError("tool downloads require HTTPS")
    checksum = hashlib.sha256()
    size = 0
    with urllib.request.urlopen(url, timeout=30) as response, destination.open("wb") as output:
        if not response.geturl().startswith("https://"):
            raise ToolError("tool download redirected away from HTTPS")
        while True:
            chunk = response.read(65536)
            if not chunk:
                break
            size += len(chunk)
            if size > max_bytes:
                raise ToolError("tool download exceeds its size limit")
            checksum.update(chunk)
            output.write(chunk)
    if digest is not None and checksum.hexdigest() != digest:
        destination.unlink(missing_ok=True)
        raise ToolError("tool download checksum did not match the pinned release")


def extract_archive(archive_path, destination):
    """Bound extraction and check each path, including previously extracted links."""
    root = destination.resolve()
    with tarfile.open(archive_path) as archive:
        members = archive.getmembers()
        if len(members) > 50000 or sum(member.size for member in members) > 4 * 1024**3:
            raise ToolError("tool archive exceeds its extraction limit")
        for member in members:
            name = Path(member.name)
            target = destination / name
            if name.is_absolute() or ".." in name.parts or not target.resolve().is_relative_to(root):
                raise ToolError("unsafe path in tool archive")
            if not (member.isfile() or member.isdir() or member.issym() or member.islnk()):
                raise ToolError("unsupported file in tool archive")
            if member.issym() or member.islnk():
                link = (target.parent if member.issym() else destination) / member.linkname
                if Path(member.linkname).is_absolute() or not link.resolve().is_relative_to(root):
                    raise ToolError("unsafe link in tool archive")
            member.mode &= 0o755
            member.uid = member.gid = 0
            member.uname = member.gname = "root"
            options = {"filter": "data"} if hasattr(tarfile, "data_filter") else {}
            archive.extract(member, destination, **options)


def platform_info():
    values = {}
    for line in Path("/etc/os-release").read_text().splitlines():
        if "=" in line:
            key, value = line.split("=", 1)
            values[key] = value.strip('"')
    arch = os.uname().machine
    return {"os": values.get("PRETTY_NAME", "Unknown OS"), "architecture": arch,
            "supported": values.get("ID") == "ubuntu" and values.get("VERSION_ID") == "22.04"
                         and arch in ("x86_64", "aarch64"),
            "boot_id": Path("/proc/sys/kernel/random/boot_id").read_text().strip()}


def guest_busy():
    try:
        descriptor = os.open(LOCK, os.O_RDONLY | os.O_NOFOLLOW)
    except FileNotFoundError:
        return False
    try:
        try:
            fcntl.flock(descriptor, fcntl.LOCK_EX | fcntl.LOCK_NB)
            return False
        except BlockingIOError:
            return True
    finally:
        os.close(descriptor)


def probe(tool, username):
    account = pwd.getpwnam(username)
    environment = {"PATH": PATH, "HOME": account.pw_dir, "USER": username, "LOGNAME": username,
                   "LANG": "C.UTF-8", "LC_ALL": "C.UTF-8", "PYTHONDONTWRITEBYTECODE": "1"}
    result = {"key": tool["key"], "id": tool["id"], "version": tool["version"], "installed": False,
              "detected_version": None, "error": None, "usage": tool["usage"]}
    for index, command in enumerate(tool["checks"]):
        if shutil.which(command[0], path=PATH) is None:
            return result
        try:
            process = subprocess.run(command, stdin=subprocess.DEVNULL, stdout=subprocess.PIPE,
                                     stderr=subprocess.STDOUT, timeout=3, env=environment, cwd=account.pw_dir,
                                     user=account.pw_uid, group=account.pw_gid,
                                     extra_groups=os.getgrouplist(username, account.pw_gid))
            output = process.stdout[:4096].decode("utf-8", errors="replace")
        except (OSError, subprocess.TimeoutExpired) as error:
            result["error"] = f"could not verify {command[0]}: {error}"[:512]
            return result
        if process.returncode != 0:
            result["error"] = f"{command[0]} verification failed: {output}"[:512]
            return result
        match = re.search(tool["match"], output)
        if index == 0:
            if not match:
                result["error"] = f"could not determine the installed {tool['id']} version"
                return result
            result["detected_version"] = match.group(1)
            if not result["detected_version"].startswith(tool["prefix_match"]):
                return result
        elif tool["id"] == "jdk" and (not match or not match.group(1).startswith("21.")):
            return result
    result["installed"] = True
    return result


class Installer:
    def __init__(self, username):
        self.username = username
        self.log = bytearray()
        self.truncated = False
        self.lock = None
        self.managed = {}

    def append(self, data):
        self.log.extend(data)
        if len(self.log) > LOG_LIMIT:
            del self.log[:-LOG_LIMIT]
            self.truncated = True

    def run(self, command):
        self.append(("+ " + shlex.join(map(str, command)) + "\n").encode())
        environment = {**os.environ, "PATH": PATH, "LC_ALL": "C.UTF-8", "DEBIAN_FRONTEND": "noninteractive"}
        with subprocess.Popen(command, stdin=subprocess.DEVNULL, stdout=subprocess.PIPE,
                              stderr=subprocess.STDOUT, env=environment, pass_fds=(self.lock,)) as process:
            while True:
                chunk = process.stdout.read(8192)
                if not chunk:
                    break
                self.append(chunk)
            if process.wait() != 0:
                raise ToolError(f"{command[0]} failed with exit code {process.returncode}; see the installation output")

    def remember(self, path, value):
        name = str(path)
        self.managed[name] = list(dict.fromkeys([*self.managed.get(name, []), value]))[-2:]
        atomic_write(ROOT / "managed.json", json.dumps(self.managed).encode(), 0o600)

    def write_managed(self, path, data, mode=0o644):
        safe_directory(path.parent)
        digest = hashlib.sha256(data).hexdigest()
        if path.exists() or path.is_symlink():
            regular_file(path)
            current = hashlib.sha256(path.read_bytes()).hexdigest()
            if current == digest:
                return
            if current not in self.managed.get(str(path), []):
                raise ToolError(f"customized file needs attention before installing: {path}")
        # Record intent first so an interrupted publication can be retried.
        self.remember(path, digest)
        atomic_write(path, data, mode)

    def link(self, name, target):
        path = BIN_DIR / name
        safe_directory(path.parent)
        target = str(target)
        if path.exists() or path.is_symlink():
            if path.is_symlink() and os.readlink(path) == target:
                return
            if not path.is_symlink() or os.readlink(path) not in self.managed.get(str(path), []):
                raise ToolError(f"customized executable needs attention before installing: {path}")
        self.remember(path, target)
        temporary = path.with_name(".uvm-" + name)
        temporary.unlink(missing_ok=True)
        temporary.symlink_to(target)
        os.replace(temporary, path)

    def apt(self, packages, *, refresh=False):
        stamp = ROOT / "apt-updated"
        if refresh or not stamp.exists() or time.time() - stamp.stat().st_mtime > 300:
            self.run(["/usr/bin/apt-get", "-o", "DPkg::Lock::Timeout=300", "-o", "Acquire::Retries=2",
                      "-o", "APT::Update::Error-Mode=any", "update"])
            atomic_write(stamp, b"updated\n", 0o600)
        pending = ROOT / "apt-pending"
        if pending.exists():
            # Complete configuration left by an interrupted tool installation.
            self.run(["/usr/bin/dpkg", "--configure", "-a"])
        atomic_write(pending, b"pending\n", 0o600)
        self.run(["/usr/bin/apt-get", "-o", "DPkg::Lock::Timeout=300", "-o", "Acquire::Retries=2",
                  "-o", "Dpkg::Options::=--force-confold", "-y", "--no-install-recommends", "install", *packages])
        pending.unlink()

    def php_repository(self, temporary):
        self.apt(["ca-certificates", "gnupg"])
        key = temporary / "php.asc"
        download(f"https://keyserver.ubuntu.com/pks/lookup?op=get&search=0x{PHP_KEY}", key, max_bytes=65536)
        result = subprocess.run(["gpg", "--batch", "--homedir", str(temporary), "--with-colons", "--show-keys", str(key)],
                                capture_output=True, check=True)
        fingerprints = [line.split(":")[9] for line in result.stdout.decode().splitlines() if line.startswith("fpr:")]
        if not fingerprints or fingerprints[0] != PHP_KEY:
            raise ToolError("PHP repository signing key did not match its pinned fingerprint")
        binary_key = temporary / "php.gpg"
        self.run(["gpg", "--batch", "--homedir", str(temporary), "--dearmor", "--output", str(binary_key), str(key)])
        self.write_managed(Path("/etc/apt/keyrings/uvm-php.gpg"), binary_key.read_bytes())
        self.write_managed(Path("/etc/apt/sources.list.d/uvm-php.list"),
                           b"deb [signed-by=/etc/apt/keyrings/uvm-php.gpg] https://ppa.launchpadcontent.net/ondrej/php/ubuntu jammy main\n")

    def payload(self, tool, temporary):
        prefix = ROOT / tool["release"]
        marker = prefix / ".uvm-tools.json"
        identity = {"release": tool["release"], "sha256": tool["sha256"]}
        if prefix.exists() or prefix.is_symlink():
            safe_directory(prefix)
            if not marker.exists():
                raise ToolError(f"installation directory is already in use: {prefix}")
            regular_file(marker)
            previous = json.loads(marker.read_bytes())
            if not isinstance(previous, dict) or any(previous.get(key) != value for key, value in identity.items()):
                raise ToolError(f"installation directory belongs to a different release: {prefix}")
            if tool["kind"] != "conda" or previous.get("complete"):
                return prefix
        artifact = temporary / "download"
        self.append(f"Downloading {tool['release']} and verifying SHA-256…\n".encode())
        download(tool["url"], artifact, tool["sha256"])
        if tool["kind"] == "archive":
            extraction = temporary / "extract"
            extraction.mkdir()
            extract_archive(artifact, extraction)
            source = extraction / tool["archive_root"]
            if not source.is_dir() or source.is_symlink():
                raise ToolError("tool archive has an unexpected layout")
            atomic_write(source / ".uvm-tools.json", json.dumps(identity).encode(), 0o600)
            os.replace(source, prefix)
        else:
            prefix.mkdir(exist_ok=True, mode=0o755)
            atomic_write(marker, json.dumps(identity).encode(), 0o600)
            if tool["kind"] == "conda":
                self.run(["/bin/bash", str(artifact), "-b", "-f", "-p", str(prefix)])
            else:
                artifact.chmod(0o755)
                os.replace(artifact, prefix / "micromamba")
        return prefix

    def configure(self, tool, prefix):
        for name, target in tool["links"].items():
            self.link(name, Path(target) if target.startswith("/") else prefix / target)
        if tool["id"] == "python":
            executable = prefix / "bin" / f"python{tool['version']}"
            self.write_managed(BIN_DIR / f"pip{tool['version']}",
                               f'#!/bin/sh\nexec {shlex.quote(str(executable))} -m pip "$@"\n'.encode(), 0o755)
        elif tool["id"] == "jdk":
            self.write_managed(PROFILE_DIR / "uvm-java.sh",
                               f"export JAVA_HOME={shlex.quote(tool['java_home'])}\n".encode())
        elif tool["id"] == "conda":
            for key, value in (("auto_activate_base", "false"),):
                self.run([str(prefix / "bin/conda"), "config", "--system", "--set", key, value])
            for key, value in (("envs_dirs", "~/.conda/envs"), ("pkgs_dirs", "~/.conda/pkgs")):
                self.run([str(prefix / "bin/conda"), "config", "--system", "--prepend", key, value])
            self.write_managed(PROFILE_DIR / "uvm-conda.sh",
                               f". {shlex.quote(str(prefix / 'etc/profile.d/conda.sh'))}\n".encode())
        elif tool["id"] == "micromamba":
            self.write_managed(PROFILE_DIR / "uvm-micromamba.sh",
                               b'export MAMBA_ROOT_PREFIX="${MAMBA_ROOT_PREFIX:-$HOME/micromamba}"\n'
                               b'eval "$(/usr/local/bin/micromamba shell hook --shell posix)"\n')

    def install(self, tool):
        safe_directory(LOCK.parent)
        self.lock = os.open(LOCK, os.O_RDWR | os.O_CREAT | os.O_NOFOLLOW, 0o600)
        try:
            metadata = os.fstat(self.lock)
            if not stat.S_ISREG(metadata.st_mode) or metadata.st_nlink != 1 or metadata.st_uid != os.geteuid():
                raise ToolError("unsafe guest tool installation lock")
            try:
                fcntl.flock(self.lock, fcntl.LOCK_EX | fcntl.LOCK_NB)
            except BlockingIOError:
                raise ToolError("another installation is still running inside this guest; wait before retrying")
            existing = probe(tool, self.username)
            if existing["installed"]:
                return {"status": "already_installed", "tool": existing}
            safe_directory(ROOT)
            manifest = ROOT / "managed.json"
            if manifest.exists():
                regular_file(manifest)
                self.managed = json.loads(manifest.read_bytes())
                if (not isinstance(self.managed, dict) or any(
                        not isinstance(values, list) or len(values) > 2 or any(not isinstance(value, str) for value in values)
                        for values in self.managed.values())):
                    raise ToolError("invalid managed tool file manifest")
            with tempfile.TemporaryDirectory(prefix=".install-", dir=ROOT) as name:
                temporary = Path(name)
                if tool["kind"] == "php":
                    self.php_repository(temporary)
                self.apt(tool["packages"], refresh=tool["kind"] == "php")
                prefix = self.payload(tool, temporary) if "url" in tool else None
                self.configure(tool, prefix)
                installed = probe(tool, self.username)
                if not installed["installed"]:
                    raise ToolError(installed["error"] or f"{tool['key']} did not pass its version check after installation")
                if prefix is not None:
                    atomic_write(prefix / ".uvm-tools.json", json.dumps({"release": tool["release"],
                                 "sha256": tool["sha256"], "complete": True}).encode(), 0o600)
                return {"status": "installed", "tool": installed}
        finally:
            os.close(self.lock)
            self.lock = None


def main(payload):
    installer = Installer(payload["username"])
    info = {}
    try:
        if os.geteuid() != 0:
            raise ToolError("tool installation and inspection require guest root")
        info = platform_info()
        info["supported"] = info["supported"] and info["architecture"] == payload["architecture"]
        if not info["supported"]:
            raise ToolError("one-click tools require an Ubuntu 22.04 guest on the selected x86_64/ARM64 architecture")
        if payload.get("boot_id") and payload["boot_id"] != info["boot_id"]:
            raise ToolError("the guest rebooted during installation; retry on the current boot")
        if payload["action"] == "probe":
            result = {**info, "busy": guest_busy(), "tools": [probe(tool, payload["username"]) for tool in payload["tools"]]}
        elif payload["action"] == "install":
            result = installer.install(payload["tools"][0])
        else:
            raise ToolError("unsupported tool operation")
        result["ok"] = True
    except (ToolError, OSError, ValueError, KeyError, tarfile.TarError, subprocess.SubprocessError) as error:
        result = {"ok": False, "error": str(error)[:2048], "supported": info.get("supported")}
    result.update(output=installer.log.decode("utf-8", errors="replace"), truncated=installer.truncated)
    print(json.dumps(result, ensure_ascii=True))


if __name__ == "__main__":
    main(json.loads(sys.argv[1]))