#!/usr/bin/env python3
"""Build the canonical source manifest and archive for remote image builds."""

from __future__ import annotations

import argparse
import hashlib
import json
import os
import pathlib
import re
import stat
import subprocess
import sys
import tarfile
import tempfile
from typing import Iterable


CONTRACT = "cp9.source-manifest/v1"
LOCAL_REPLACEMENT = re.compile(r"=>\s+((?:\.\.?/|/)[^\s]+)")
class ManifestError(RuntimeError):
    pass


def run_git(root: pathlib.Path, *args: str) -> bytes:
    command = ["git", "-C", str(root), *args]
    try:
        return subprocess.run(
            command,
            check=True,
            stdout=subprocess.PIPE,
            stderr=subprocess.PIPE,
        ).stdout
    except (OSError, subprocess.CalledProcessError) as exc:
        detail = ""
        if isinstance(exc, subprocess.CalledProcessError):
            detail = exc.stderr.decode(errors="replace").strip()
        suffix = f": {detail}" if detail else ""
        raise ManifestError(
            f"Git source discovery failed for {root} while running "
            f"{' '.join(command)}{suffix}"
        ) from exc


def git_root(path: pathlib.Path) -> pathlib.Path:
    output = run_git(path, "rev-parse", "--show-toplevel")
    value = output.decode().strip()
    if not value:
        raise ManifestError(f"Git source discovery failed for {path}: empty worktree root")
    return pathlib.Path(value).resolve()


def module_file(path: pathlib.Path, boundary: pathlib.Path) -> pathlib.Path | None:
    current = path if path.is_dir() else path.parent
    while True:
        candidate = current / "go.mod"
        if candidate.is_file():
            return candidate
        if current == boundary or current.parent == current:
            return None
        current = current.parent


def local_replacements(go_mod: pathlib.Path) -> list[pathlib.Path]:
    try:
        text = go_mod.read_text(encoding="utf-8")
    except OSError as exc:
        raise ManifestError(f"cannot read local replacement manifest {go_mod}: {exc}") from exc
    paths: list[pathlib.Path] = []
    for match in LOCAL_REPLACEMENT.finditer(text):
        raw = match.group(1).strip('"')
        candidate = pathlib.Path(raw)
        if not candidate.is_absolute():
            candidate = go_mod.parent / candidate
        candidate = candidate.resolve()
        if not candidate.exists():
            raise ManifestError(
                f"intentional local replacement {raw!r} from {go_mod} does not exist at {candidate}"
            )
        paths.append(candidate)
    return paths


def discover_roots(repo: pathlib.Path) -> tuple[list[pathlib.Path], pathlib.Path]:
    if not repo.is_dir():
        raise ManifestError(f"source repository does not exist or is not a directory: {repo}")
    primary_root = git_root(repo)
    roots: set[pathlib.Path] = {primary_root}
    visited_modules: set[pathlib.Path] = set()
    queue: list[pathlib.Path] = [repo]
    while queue:
        source = queue.pop(0)
        source_root = git_root(source)
        roots.add(source_root)
        go_mod = module_file(source, source_root)
        if go_mod is None or go_mod in visited_modules:
            continue
        visited_modules.add(go_mod)
        for replacement in local_replacements(go_mod):
            replacement_root = git_root(replacement)
            roots.add(replacement_root)
            queue.append(replacement)
    return sorted(roots, key=lambda item: str(item)), primary_root


def nul_paths(payload: bytes) -> list[pathlib.PurePosixPath]:
    return [
        pathlib.PurePosixPath(os.fsdecode(value))
        for value in payload.split(b"\0")
        if value
    ]


def is_below(path: pathlib.Path, parent: pathlib.Path) -> bool:
    try:
        return os.path.commonpath((str(path), str(parent))) == str(parent)
    except ValueError:
        return False


def enumerate_entries(roots: list[pathlib.Path], sync_root: pathlib.Path) -> list[str]:
    entries: dict[str, pathlib.Path] = {}
    for root in roots:
        tracked = nul_paths(run_git(root, "ls-files", "-z", "--cached"))
        untracked = nul_paths(run_git(root, "ls-files", "-z", "--others", "--exclude-standard"))
        for relative in tracked + untracked:
            if relative.is_absolute() or ".." in relative.parts:
                raise ManifestError(f"Git returned unsafe source path {relative!s} for {root}")
            source = root.joinpath(*relative.parts)
            if any(other != root and is_below(source, other) for other in roots):
                continue
            try:
                info = source.lstat()
            except FileNotFoundError:
                # A tracked deletion is intentionally absent from the new remote tree.
                continue
            except OSError as exc:
                raise ManifestError(f"cannot inspect source path {source}: {exc}") from exc
            if not (stat.S_ISREG(info.st_mode) or stat.S_ISLNK(info.st_mode)):
                raise ManifestError(
                    f"unsupported tracked source type at {source}; only files and symlinks are allowed"
                )
            archive_name = source.relative_to(sync_root).as_posix()
            previous = entries.get(archive_name)
            if previous is not None and previous != source:
                raise ManifestError(
                    f"source roots map two files to archive path {archive_name}: {previous} and {source}"
                )
            entries[archive_name] = source
    return sorted(entries)


def validate_symlinks(entries: Iterable[str], sync_root: pathlib.Path) -> None:
    for entry in entries:
        source = sync_root / pathlib.PurePosixPath(entry)
        if not source.is_symlink():
            continue
        try:
            target = pathlib.Path(os.readlink(source))
        except OSError as exc:
            raise ManifestError(f"cannot read source symlink {source}: {exc}") from exc
        if target.is_absolute():
            raise ManifestError(f"source symlink {source} has an absolute target {target}")
        resolved = (source.parent / target).resolve(strict=False)
        if not is_below(resolved, sync_root):
            raise ManifestError(
                f"source symlink {source} escapes the synchronized roots via target {target}"
            )


def add_hash_field(digest: "hashlib._Hash", value: bytes) -> None:
    digest.update(len(value).to_bytes(8, "big"))
    digest.update(value)


def fingerprint(entries: Iterable[str], sync_root: pathlib.Path) -> str:
    digest = hashlib.sha256()
    add_hash_field(digest, CONTRACT.encode())
    for entry in entries:
        source = sync_root / pathlib.PurePosixPath(entry)
        info = source.lstat()
        add_hash_field(digest, os.fsencode(entry))
        add_hash_field(digest, oct(stat.S_IMODE(info.st_mode)).encode())
        if stat.S_ISLNK(info.st_mode):
            add_hash_field(digest, b"symlink")
            add_hash_field(digest, os.fsencode(os.readlink(source)))
            continue
        add_hash_field(digest, b"file")
        with source.open("rb") as handle:
            while True:
                chunk = handle.read(1024 * 1024)
                if not chunk:
                    break
                add_hash_field(digest, chunk)
    return digest.hexdigest()


def write_archive(path: pathlib.Path, entries: Iterable[str], sync_root: pathlib.Path) -> None:
    path.parent.mkdir(parents=True, exist_ok=True)
    fd, temporary_name = tempfile.mkstemp(prefix=f".{path.name}.", dir=path.parent)
    os.close(fd)
    temporary = pathlib.Path(temporary_name)
    try:
        with tarfile.open(temporary, mode="w", format=tarfile.PAX_FORMAT) as archive:
            for entry in entries:
                source = sync_root / pathlib.PurePosixPath(entry)
                tar_info = archive.gettarinfo(str(source), arcname=entry)
                tar_info.uid = 0
                tar_info.gid = 0
                tar_info.uname = ""
                tar_info.gname = ""
                tar_info.mtime = 0
                if tar_info.isreg():
                    with source.open("rb") as handle:
                        archive.addfile(tar_info, handle)
                else:
                    archive.addfile(tar_info)
        os.replace(temporary, path)
    except Exception:
        temporary.unlink(missing_ok=True)
        raise


def write_metadata(path: pathlib.Path, value: dict[str, object]) -> None:
    path.parent.mkdir(parents=True, exist_ok=True)
    fd, temporary_name = tempfile.mkstemp(prefix=f".{path.name}.", dir=path.parent)
    temporary = pathlib.Path(temporary_name)
    try:
        with os.fdopen(fd, "w", encoding="utf-8") as handle:
            json.dump(value, handle, indent=2, sort_keys=True)
            handle.write("\n")
        os.replace(temporary, path)
    except Exception:
        temporary.unlink(missing_ok=True)
        raise


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--repo", required=True, type=pathlib.Path)
    parser.add_argument("--metadata", required=True, type=pathlib.Path)
    parser.add_argument("--archive", required=True, type=pathlib.Path)
    parser.add_argument("--git-root-only", action="store_true")
    return parser.parse_args()


def main() -> int:
    args = parse_args()
    try:
        repo = args.repo.resolve()
        if args.git_root_only:
            primary_root = git_root(repo)
            roots = [primary_root]
        else:
            roots, primary_root = discover_roots(repo)
        sync_root = pathlib.Path(os.path.commonpath([str(root) for root in roots])).resolve()
        entries = enumerate_entries(roots, sync_root)
        validate_symlinks(entries, sync_root)
        source_fingerprint = fingerprint(entries, sync_root)
        source_key = hashlib.sha256(
            b"\0".join(os.fsencode(str(root)) for root in roots)
        ).hexdigest()[:16]
        metadata = {
            "version": CONTRACT,
            "sync_root": str(sync_root),
            "repo_relative": repo.relative_to(sync_root).as_posix(),
            "primary_root": str(primary_root),
            "roots": [str(root) for root in roots],
            "source_key": source_key,
            "fingerprint": source_fingerprint,
            "entries": entries,
        }
        write_archive(args.archive, entries, sync_root)
        write_metadata(args.metadata, metadata)
        return 0
    except (ManifestError, OSError, ValueError) as exc:
        print(f"source manifest failed: {exc}", file=sys.stderr)
        return 1


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