#!/usr/bin/env python3
import argparse
import io
import subprocess
import sys
import tarfile
from pathlib import Path, PurePosixPath

import yaml

FLUX_KUSTOMIZATION = "kustomize.toolkit.fluxcd.io"
KUSTOMIZATION_FILES = ("kustomization.yaml", "kustomization.yml", "Kustomization")


class SliceError(Exception):
    pass


def is_remote(ref: str) -> bool:
    return "://" in ref or ref.startswith(("git@", "github.com/", "gitlab.com/"))


def kustomization_file(directory: Path) -> Path | None:
    for name in KUSTOMIZATION_FILES:
        candidate = directory / name
        if candidate.is_file():
            return candidate
    return None


def local_refs(doc: dict) -> list[str]:
    refs = []
    refs += doc.get("resources") or []
    refs += doc.get("components") or []
    refs += doc.get("patchesStrategicMerge") or []
    refs += [p["path"] for p in doc.get("patches") or [] if isinstance(p, dict) and "path" in p]
    for key in ("configMapGenerator", "secretGenerator"):
        for gen in doc.get(key) or []:
            for entry in (gen.get("files") or []) + (gen.get("envs") or []):
                refs.append(entry.split("=", 1)[-1])
            if gen.get("env"):
                refs.append(gen["env"])
    return [r for r in refs if isinstance(r, str)]


def flux_paths(path: Path) -> list[str]:
    try:
        docs = list(yaml.safe_load_all(path.read_text()))
    except yaml.YAMLError:
        return []
    return [
        d["spec"]["path"]
        for d in docs
        if isinstance(d, dict)
        and str(d.get("apiVersion", "")).startswith(FLUX_KUSTOMIZATION)
        and d.get("kind") == "Kustomization"
        and d.get("spec", {}).get("path")
    ]


def collect(repo: Path, start: str) -> tuple[set[PurePosixPath], list[str]]:
    repo = repo.resolve()
    files: set[PurePosixPath] = set()
    entries: list[str] = []
    seen_dirs: set[Path] = set()

    def rel(p: Path) -> PurePosixPath:
        try:
            return PurePosixPath(p.resolve().relative_to(repo).as_posix())
        except ValueError:
            raise SliceError(f"reference outside the repository: {p}")

    def add_file(p: Path):
        r = rel(p)
        if r in files:
            return
        files.add(r)
        if p.suffix in (".yaml", ".yml"):
            for fp in flux_paths(p):
                visit_dir(repo / fp, entry=True)

    def visit_dir(d: Path, entry: bool = False):
        d = d.resolve()
        if d in seen_dirs:
            return
        seen_dirs.add(d)
        kfile = kustomization_file(d)
        if kfile is None:
            raise SliceError(f"no kustomization.yaml in {rel(d)}")
        if entry:
            entries.append(str(rel(d)))
        add_file(kfile)
        doc = yaml.safe_load(kfile.read_text()) or {}
        for ref in local_refs(doc):
            if is_remote(ref):
                raise SliceError(f"remote reference {ref} in {rel(kfile)}")
            target = (d / ref)
            if target.is_dir():
                visit_dir(target)
            elif target.is_file():
                add_file(target)
            else:
                raise SliceError(f"reference {ref} in {rel(kfile)} does not exist")

    visit_dir(repo / start, entry=True)
    return files, sorted(set(entries))


def write_tar(repo: Path, files: set[PurePosixPath], out: Path):
    buf = io.BytesIO()
    with tarfile.open(fileobj=buf, mode="w", format=tarfile.PAX_FORMAT) as tar:
        for f in sorted(files):
            data = (repo / f).read_bytes()
            info = tarfile.TarInfo(str(f))
            info.size = len(data)
            info.mode = 0o644
            info.mtime = 0
            info.uid = info.gid = 0
            info.uname = info.gname = ""
            tar.addfile(info, io.BytesIO(data))
    out.write_bytes(buf.getvalue())


def kustomize_build(path: Path) -> bytes:
    p = subprocess.run(["kubectl", "kustomize", str(path)], capture_output=True)
    if p.returncode != 0:
        raise SliceError(f"kustomize build {path} failed: {p.stderr.decode().strip()}")
    return p.stdout


def verify(repo: Path, out: Path, entries: list[str]):
    import tempfile
    with tempfile.TemporaryDirectory() as tmp:
        with tarfile.open(out) as tar:
            tar.extractall(tmp, filter="data")
        for e in entries:
            if kustomize_build(repo / e) != kustomize_build(Path(tmp) / e):
                raise SliceError(f"kustomize build of {e} differs between slice and full repository")


def main() -> int:
    ap = argparse.ArgumentParser()
    ap.add_argument("contour")
    ap.add_argument("--repo", default=".")
    ap.add_argument("--out", required=True)
    ap.add_argument("--no-verify", action="store_true")
    args = ap.parse_args()
    repo = Path(args.repo).resolve()
    out = Path(args.out)
    try:
        files, entries = collect(repo, f"clusters/{args.contour}")
        write_tar(repo, files, out)
        if not args.no_verify:
            verify(repo, out, entries)
    except SliceError as e:
        print(f"slice {args.contour}: {e}", file=sys.stderr)
        return 1
    print(f"slice {args.contour}: {len(files)} files, entry points: {', '.join(entries)}")
    return 0


if __name__ == "__main__":
    sys.exit(main())