#!/usr/bin/env python3
"""Shared helpers for the Overdeck K3s Phase 1 tooling.

The module is deliberately Python-stdlib-only so it can run on the workstation
and on clean Debian-family buildboxes without provisioning a virtualenv.
"""
from __future__ import annotations

import hashlib
import ipaddress
import json
import os
import re
import shutil
import stat
import subprocess
import tarfile
import tempfile
from contextlib import contextmanager
from dataclasses import dataclass
from datetime import datetime, timezone
from pathlib import Path, PurePosixPath
from typing import Any, BinaryIO, Iterable, Iterator, Mapping, Sequence

try:
    import fcntl
except ImportError:  # pragma: no cover - Overdeck's supported workstations are Unix
    fcntl = None  # type: ignore[assignment]

K3S_VERSION_RE = re.compile(r"\bk3s version (?P<version>v[^\s]+)")
SAFE_NAME_RE = re.compile(r"[^A-Za-z0-9._-]+")
SSH_USER_RE = re.compile(r"[A-Za-z_][A-Za-z0-9._-]*\Z")
DNS_LABEL_RE = re.compile(r"[A-Za-z0-9](?:[A-Za-z0-9-]{0,61}[A-Za-z0-9])?\Z")


class Phase1Error(RuntimeError):
    """Raised for a user-actionable Phase 1 failure."""


@dataclass(frozen=True)
class HostAccess:
    name: str
    host: str
    port: int
    user: str
    identity_file: str | None
    door: str
    magic_dns: str | None
    state: str

    @property
    def target(self) -> str:
        host = f"[{self.host}]" if ":" in self.host else self.host
        return f"{self.user}@{host}"


def utc_now() -> str:
    return datetime.now(timezone.utc).replace(microsecond=0).isoformat().replace("+00:00", "Z")


def utc_stamp() -> str:
    return datetime.now(timezone.utc).strftime("%Y%m%dT%H%M%SZ")


def safe_name(value: str) -> str:
    cleaned = SAFE_NAME_RE.sub("-", value).strip("-.")
    if not cleaned:
        raise Phase1Error(f"value cannot be converted to a safe name: {value!r}")
    return cleaned


def sha256_file(path: Path, *, chunk_size: int = 1024 * 1024) -> str:
    digest = hashlib.sha256()
    with path.open("rb") as handle:
        while True:
            chunk = handle.read(chunk_size)
            if not chunk:
                break
            digest.update(chunk)
    return digest.hexdigest()


def sha256_bytes(payload: bytes) -> str:
    return hashlib.sha256(payload).hexdigest()


def parse_k3s_version(text: str) -> str:
    match = K3S_VERSION_RE.search(text)
    if not match:
        raise Phase1Error(f"cannot parse K3s version from output: {text.strip()!r}")
    return match.group("version")


def read_json(path: Path) -> Any:
    try:
        with path.open(encoding="utf-8") as handle:
            return json.load(handle)
    except FileNotFoundError as exc:
        raise Phase1Error(f"required JSON file does not exist: {path}") from exc
    except json.JSONDecodeError as exc:
        raise Phase1Error(f"invalid JSON in {path}: {exc}") from exc


def atomic_write_json(path: Path, value: Any, *, mode: int = 0o600) -> None:
    path.parent.mkdir(parents=True, exist_ok=True)
    payload = json.dumps(value, indent=2, sort_keys=True) + "\n"
    fd, tmp_name = tempfile.mkstemp(prefix=f".{path.name}.", dir=path.parent)
    tmp = Path(tmp_name)
    try:
        os.fchmod(fd, mode)
        with os.fdopen(fd, "w", encoding="utf-8") as handle:
            handle.write(payload)
            handle.flush()
            os.fsync(handle.fileno())
        os.replace(tmp, path)
        dir_fd = os.open(path.parent, os.O_DIRECTORY)
        try:
            os.fsync(dir_fd)
        finally:
            os.close(dir_fd)
    finally:
        try:
            tmp.unlink()
        except FileNotFoundError:
            pass


@contextmanager
def exclusive_lock(path: Path) -> Iterator[BinaryIO]:
    """Hold a non-blocking process lock for the duration of a Phase 1 run.

    The workstation is the trusted controller.  A local lock prevents a manual
    invocation and the recurring systemd timer from mutating/backuping the same
    server concurrently.  The server helper independently rejects overlapping
    control-plane transactions from any controller.
    """

    if fcntl is None:  # pragma: no cover - defensive portability guard
        raise Phase1Error("Phase 1 requires Unix advisory file locking")
    path.parent.mkdir(parents=True, exist_ok=True)
    handle = path.open("a+b")
    os.chmod(path, 0o600)
    try:
        try:
            fcntl.flock(handle.fileno(), fcntl.LOCK_EX | fcntl.LOCK_NB)
        except BlockingIOError as exc:
            raise Phase1Error(f"another Phase 1 operation holds the lock: {path}") from exc
        handle.seek(0)
        handle.truncate()
        handle.write(f"pid={os.getpid()} started_utc={utc_now()}\n".encode())
        handle.flush()
        os.fsync(handle.fileno())
        yield handle
    finally:
        try:
            fcntl.flock(handle.fileno(), fcntl.LOCK_UN)
        finally:
            handle.close()


def normalize_path(path: str | Path) -> Path:
    return Path(os.path.expandvars(os.path.expanduser(str(path)))).resolve()


def absolute_path_without_symlink_resolution(path: str | Path) -> Path:
    """Expand a path without hiding whether its final component is a symlink.

    ``Path.resolve`` is appropriate for repository and inventory locations, but
    it defeats explicit symlink checks for private keys and restore targets.  A
    lexical absolute path preserves the final path component for those checks.
    """

    expanded = os.path.expandvars(os.path.expanduser(str(path)))
    return Path(os.path.abspath(expanded))


def load_host_registry(path: Path) -> dict[str, Any]:
    document = read_json(path)
    if not isinstance(document, dict) or document.get("schema_version") != 1:
        raise Phase1Error(f"unsupported host registry schema in {path}")
    hosts = document.get("hosts")
    if not isinstance(hosts, list):
        raise Phase1Error(f"host registry has no hosts array: {path}")
    return document


def validate_ssh_host(host: str) -> str:
    try:
        ipaddress.ip_address(host)
        return host
    except ValueError:
        pass
    if len(host) > 253 or not host or any(not DNS_LABEL_RE.fullmatch(label) for label in host.rstrip(".").split(".")):
        raise Phase1Error(f"invalid SSH host: {host!r}")
    return host


def validate_ssh_user(user: str) -> str:
    if not SSH_USER_RE.fullmatch(user):
        raise Phase1Error(f"invalid SSH user: {user!r}")
    return user


def resolve_host_access(
    registry: Mapping[str, Any],
    name: str,
    *,
    preferred_door: str = "tailscale_ip",
    require_reachable: bool = True,
) -> HostAccess:
    candidates = [item for item in registry.get("hosts", []) if isinstance(item, dict) and item.get("name") == name]
    if len(candidates) != 1:
        raise Phase1Error(f"expected exactly one host named {name!r}; found {len(candidates)}")
    entry = candidates[0]
    state = str(entry.get("state") or "unknown")
    if require_reachable and state != "reachable":
        raise Phase1Error(f"host {name} is declared {state!r}; Phase 1 requires registry state 'reachable'")
    access = entry.get("access")
    if not isinstance(access, dict):
        raise Phase1Error(f"host {name} has no access declaration")
    order = [preferred_door, "tailscale_ip", "tailscale_ssh", "lan"]
    seen: set[str] = set()
    door_name = ""
    door: Mapping[str, Any] | None = None
    for candidate in order:
        if candidate in seen:
            continue
        seen.add(candidate)
        raw = access.get(candidate)
        if isinstance(raw, dict) and raw.get("host") and raw.get("user"):
            door_name = candidate
            door = raw
            break
    if door is None:
        raise Phase1Error(f"host {name} has no usable access door")
    try:
        port = int(door.get("port") or 22)
    except (TypeError, ValueError) as exc:
        raise Phase1Error(f"host {name} has an invalid SSH port") from exc
    if not (1 <= port <= 65535):
        raise Phase1Error(f"host {name} has an out-of-range SSH port: {port}")
    identity = door.get("identity_file")
    if identity is not None:
        identity = str(normalize_path(str(identity)))
    magic_dns_raw = access.get("tailscale_ssh")
    magic_dns = str(magic_dns_raw.get("host")) if isinstance(magic_dns_raw, dict) and magic_dns_raw.get("host") else None
    return HostAccess(
        name=name,
        host=validate_ssh_host(str(door["host"])),
        port=port,
        user=validate_ssh_user(str(door["user"])),
        identity_file=identity,
        door=door_name,
        magic_dns=magic_dns,
        state=state,
    )


def eligible_reachable_hosts(registry: Mapping[str, Any]) -> list[str]:
    names: list[str] = []
    for item in registry.get("hosts", []):
        if not isinstance(item, dict):
            continue
        name = item.get("name")
        if item.get("state") == "reachable" and isinstance(name, str) and name:
            names.append(name)
    return sorted(set(names))


def build_ssh_command(access: HostAccess, *, timeout: int = 12) -> list[str]:
    if timeout < 1:
        raise Phase1Error("SSH timeout must be positive")
    command = [
        "ssh",
        "-F",
        "/dev/null",
        "-p",
        str(access.port),
        "-o",
        "BatchMode=yes",
        "-o",
        f"ConnectTimeout={timeout}",
        "-o",
        "ServerAliveInterval=15",
        "-o",
        "ServerAliveCountMax=2",
        "-o",
        "StrictHostKeyChecking=yes",
    ]
    if access.identity_file:
        command.extend(["-o", "IdentitiesOnly=yes", "-i", access.identity_file])
    command.extend(["--", access.target])
    return command


def run_command(
    argv: Sequence[str],
    *,
    input_bytes: bytes | None = None,
    input_text: str | None = None,
    check: bool = True,
    timeout: int | None = None,
    env: Mapping[str, str] | None = None,
    cwd: Path | None = None,
) -> subprocess.CompletedProcess[Any]:
    if input_bytes is not None and input_text is not None:
        raise ValueError("provide only one of input_bytes or input_text")
    text = input_bytes is None
    input_value: bytes | str | None = input_text if input_text is not None else input_bytes
    completed = subprocess.run(
        list(argv),
        input=input_value,
        stdout=subprocess.PIPE,
        stderr=subprocess.PIPE,
        text=text,
        timeout=timeout,
        env=dict(env) if env is not None else None,
        cwd=str(cwd) if cwd else None,
        check=False,
    )
    if check and completed.returncode != 0:
        stdout = completed.stdout.decode(errors="replace") if isinstance(completed.stdout, bytes) else completed.stdout
        stderr = completed.stderr.decode(errors="replace") if isinstance(completed.stderr, bytes) else completed.stderr
        rendered = " ".join(str(part) for part in argv)
        raise Phase1Error(
            f"command failed with exit {completed.returncode}: {rendered}\n"
            f"stdout:\n{stdout[-4000:]}\n"
            f"stderr:\n{stderr[-4000:]}"
        )
    return completed


def require_commands(commands: Iterable[str]) -> None:
    missing = [command for command in commands if shutil.which(command) is None]
    if missing:
        raise Phase1Error(f"missing required command(s): {', '.join(missing)}")


def validate_tar_member_name(name: str) -> None:
    pure = PurePosixPath(name)
    if not name or name in {".", "./"} or pure.is_absolute() or ".." in pure.parts:
        raise Phase1Error(f"archive contains unsafe path: {name!r}")


def _extract_regular_member(handle: tarfile.TarFile, member: tarfile.TarInfo, target: Path) -> None:
    source = handle.extractfile(member)
    if source is None:
        raise Phase1Error(f"archive member has no readable payload: {member.name!r}")
    target.parent.mkdir(parents=True, exist_ok=True)
    if target.exists() or target.is_symlink():
        raise Phase1Error(f"archive contains duplicate member path: {member.name!r}")
    fd = os.open(target, os.O_WRONLY | os.O_CREAT | os.O_EXCL | os.O_NOFOLLOW, 0o600)
    try:
        with os.fdopen(fd, "wb") as destination:
            shutil.copyfileobj(source, destination, length=1024 * 1024)
            destination.flush()
            os.fsync(destination.fileno())
    finally:
        source.close()
    # Strip set-id/sticky bits.  Backup payload files are data, never executables.
    os.chmod(target, int(member.mode) & 0o0777)


def safe_extract_tar(archive: Path, destination: Path) -> None:
    """Extract only ordinary files/directories without trusting tarfile defaults.

    This intentionally avoids ``TarFile.extractall`` so the verifier behaves
    consistently on all supported Python versions and never follows archive
    links, applies ownership, or creates device nodes.
    """

    destination.mkdir(parents=True, exist_ok=True)
    root = destination.resolve()
    with tarfile.open(archive, "r:*") as handle:
        members = handle.getmembers()
        seen: set[str] = set()
        for member in members:
            validate_tar_member_name(member.name)
            normalized = PurePosixPath(member.name).as_posix().rstrip("/")
            if normalized in seen:
                raise Phase1Error(f"archive contains duplicate member path: {member.name!r}")
            seen.add(normalized)
            if member.issym() or member.islnk() or member.isdev() or member.isfifo():
                raise Phase1Error(f"archive contains unsupported special member: {member.name!r}")
            target = (root / member.name).resolve()
            if target != root and root not in target.parents:
                raise Phase1Error(f"archive member escapes extraction root: {member.name!r}")
            if not member.isdir() and not member.isfile():
                raise Phase1Error(f"archive contains unsupported member type: {member.name!r}")
        directory_modes: list[tuple[Path, int]] = []
        for member in members:
            target = root / member.name
            if member.isdir():
                target.mkdir(parents=True, exist_ok=True)
                os.chmod(target, 0o700)
                directory_modes.append((target, int(member.mode) & 0o0777))
            else:
                _extract_regular_member(handle, member, target)
        for target, mode in sorted(directory_modes, key=lambda item: len(item[0].parts), reverse=True):
            os.chmod(target, mode)


def verify_manifest(extracted_root: Path, manifest_path: Path) -> dict[str, Any]:
    manifest = read_json(manifest_path)
    if not isinstance(manifest, dict) or manifest.get("schema_version") not in {1, 2}:
        raise Phase1Error("unsupported backup manifest schema")
    schema = int(manifest["schema_version"])
    files = manifest.get("files")
    if not isinstance(files, list) or not files:
        raise Phase1Error("backup manifest has no files")
    seen: set[str] = set()
    for item in files:
        if not isinstance(item, dict):
            raise Phase1Error("backup manifest contains an invalid file entry")
        rel = item.get("path")
        expected_sha = item.get("sha256")
        expected_size = item.get("size")
        if not isinstance(rel, str) or not rel or not isinstance(expected_sha, str) or not isinstance(expected_size, int):
            raise Phase1Error("backup manifest file entry is incomplete")
        validate_tar_member_name(rel)
        if rel in seen:
            raise Phase1Error(f"backup manifest contains duplicate path: {rel}")
        seen.add(rel)
        path = (extracted_root / rel).resolve()
        if extracted_root.resolve() not in path.parents:
            raise Phase1Error(f"manifest path escapes extraction root: {rel}")
        if not path.is_file() or path.is_symlink():
            raise Phase1Error(f"manifest file is missing or unsafe: {rel}")
        actual_size = path.stat().st_size
        if actual_size != expected_size:
            raise Phase1Error(f"size mismatch for {rel}: expected {expected_size}, got {actual_size}")
        expected_mode = item.get("mode")
        if expected_mode is not None:
            if not isinstance(expected_mode, str) or not re.fullmatch(r"[0-7]{4}", expected_mode):
                raise Phase1Error(f"backup manifest has an invalid mode for {rel}")
            actual_mode = f"{stat.S_IMODE(path.stat().st_mode):04o}"
            if actual_mode != expected_mode:
                raise Phase1Error(
                    f"mode mismatch for {rel}: expected {expected_mode}, got {actual_mode}"
                )
        actual_sha = sha256_file(path)
        if actual_sha != expected_sha:
            raise Phase1Error(f"SHA-256 mismatch for {rel}")
    token_rel = manifest.get("server_token_path")
    if not isinstance(token_rel, str) or token_rel not in seen:
        raise Phase1Error("backup manifest does not pair a server token with the datastore")
    datastore = manifest.get("datastore")
    if not isinstance(datastore, dict) or datastore.get("type") not in {"sqlite", "etcd"}:
        raise Phase1Error("backup manifest datastore type is unsupported")
    if schema == 2:
        purpose = manifest.get("purpose")
        if purpose not in {"prechange", "postchange", "scheduled"}:
            raise Phase1Error("backup manifest has an unsupported purpose")
        agent = manifest.get("agent_token")
        if not isinstance(agent, dict):
            raise Phase1Error("backup manifest is missing the agent-token restoration descriptor")
        classification = agent.get("classification")
        payload_path = agent.get("payload_path")
        if classification == "regular":
            if payload_path != "payload/server/agent-token" or payload_path not in seen:
                raise Phase1Error("regular agent token is not paired with its payload file")
        elif classification == "symlink-to-server-token":
            if agent.get("restore") != "symlink-to-token" or agent.get("link_target") != "token":
                raise Phase1Error("agent-token symlink restoration descriptor is unsafe")
            if payload_path is not None or "payload/server/agent-token" in seen:
                raise Phase1Error("agent-token symlink must not be archived as a file or link")
        elif classification == "absent":
            if payload_path is not None or "payload/server/agent-token" in seen:
                raise Phase1Error("absent agent token unexpectedly has payload material")
        else:
            raise Phase1Error("backup manifest has an unsupported agent-token classification")
    actual_files = {
        path.relative_to(extracted_root).as_posix()
        for path in extracted_root.rglob("*")
        if path.is_file() and not path.is_symlink() and path != manifest_path
    }
    if actual_files != seen:
        missing = sorted(seen - actual_files)
        unexpected = sorted(actual_files - seen)
        raise Phase1Error(
            "backup archive and manifest file sets differ: "
            f"missing={missing[:10]} unexpected={unexpected[:10]}"
        )
    return manifest


def prune_backup_pairs(backup_dir: Path, *, prefix: str, retain: int) -> list[str]:
    if retain < 1:
        raise Phase1Error("backup retention count must be at least one")
    backup_dir.mkdir(parents=True, exist_ok=True)
    candidates: list[Path] = []
    for path in backup_dir.glob(f"{prefix}-*.tar.age"):
        if not path.is_file():
            continue
        metadata = path.with_suffix("").with_suffix(".json")
        try:
            document = read_json(metadata)
        except Phase1Error:
            # Unpaired or malformed archives are retained for operator review;
            # automated retention never destroys unverified evidence.
            continue
        if not isinstance(document, dict) or document.get("status") != "verified":
            continue
        expected_sha = document.get("encrypted_sha256")
        if not isinstance(expected_sha, str) or expected_sha != sha256_file(path):
            continue
        candidates.append(path)
    candidates = sorted(
        candidates,
        key=lambda path: (path.stat().st_mtime_ns, path.name),
        reverse=True,
    )
    removed: list[str] = []
    for archive in candidates[retain:]:
        metadata = archive.with_suffix("").with_suffix(".json")
        archive.unlink(missing_ok=True)
        metadata.unlink(missing_ok=True)
        removed.append(archive.name)
    return removed


def mode_string(mode: int) -> str:
    return stat.filemode(mode)


def redact_mapping(value: Any, *, sensitive_keys: set[str] | None = None) -> Any:
    sensitive = sensitive_keys or {
        "token",
        "server_token",
        "agent_token",
        "password",
        "secret",
        "credential",
        "private_key",
        "identity",
    }
    if isinstance(value, dict):
        redacted: dict[str, Any] = {}
        for key, child in value.items():
            normalized = str(key).lower().replace("-", "_")
            if any(term in normalized for term in sensitive):
                redacted[key] = "<redacted>"
            else:
                redacted[key] = redact_mapping(child, sensitive_keys=sensitive)
        return redacted
    if isinstance(value, list):
        return [redact_mapping(item, sensitive_keys=sensitive) for item in value]
    return value