#!/usr/bin/env python3
from __future__ import annotations

import argparse
import json
import os
import platform
import shutil
import subprocess
import sys
import tarfile
import zipfile
from pathlib import Path
from typing import Any


def run(cmd: list[str], *, capture: bool = False) -> str:
    if capture:
        return subprocess.check_output(cmd, text=True).strip()
    subprocess.check_call(cmd)
    return ""


def cargo_metadata() -> dict[str, Any]:
    out = run(["cargo", "metadata", "--no-deps", "--format-version", "1"], capture=True)
    return json.loads(out)


def rustc_host_triple() -> str:
    v = run(["rustc", "-vV"], capture=True)
    for line in v.splitlines():
        if line.startswith("host: "):
            return line.split("host: ", 1)[1].strip()
    raise RuntimeError("Could not determine host target triple from `rustc -vV`")


def is_windows_host() -> bool:
    # Works for normal Windows Python and most MSYS/Cygwin Pythons too.
    sp = sys.platform.lower()
    ps = platform.system().lower()
    return (
        os.name == "nt"
        or sp.startswith("win")
        or sp.startswith("cygwin")
        or sp.startswith("msys")
        or "windows" in ps
        or "cygwin" in ps
        or "msys" in ps
    )


def exe_suffix_for_target(target_triple: str) -> str:
    return ".exe" if "windows" in target_triple else ""


def find_bin_targets(meta: dict[str, Any]) -> list[tuple[str, str, str]]:
    bins: list[tuple[str, str, str]] = []
    for p in meta.get("packages", []):
        for t in p.get("targets", []):
            if "bin" in t.get("kind", []):
                bins.append((p["name"], p["version"], t["name"]))
    bins.sort(key=lambda x: (x[0], x[2], x[1]))  # stable deterministic choice
    return bins


def find_owner_package_for_bin(meta: dict[str, Any], bin_name: str) -> tuple[str, str]:
    for p in meta.get("packages", []):
        for t in p.get("targets", []):
            if t.get("name") == bin_name and "bin" in t.get("kind", []):
                return p["name"], p["version"]
    raise RuntimeError(f"Could not find a package providing bin '{bin_name}'")


def stage_and_archive(
    *,
    pkg_name: str,
    pkg_version: str,
    bin_path: Path,
    data_dir: Path,
    dist_dir: Path,
    stage_root: Path,
    target_triple_for_name: str,
) -> Path:
    pkg_base = f"{pkg_name}-v{pkg_version}-{target_triple_for_name}"
    stage_dir = stage_root / pkg_base
    stage_data_dir = stage_dir / "data"

    if stage_root.exists():
        shutil.rmtree(stage_root)
    stage_data_dir.mkdir(parents=True, exist_ok=True)
    dist_dir.mkdir(parents=True, exist_ok=True)

    shutil.copy2(bin_path, stage_dir / bin_path.name)

    mmdbs = sorted(data_dir.glob("*.mmdb")) if data_dir.exists() else []
    if mmdbs:
        for f in mmdbs:
            shutil.copy2(f, stage_data_dir / f.name)
    else:
        print("WARN: no ./data/*.mmdb found; packaging binary only.", file=sys.stderr)

    if is_windows_host():
        out = dist_dir / f"{pkg_base}.zip"
        with zipfile.ZipFile(out, "w", compression=zipfile.ZIP_DEFLATED) as z:
            for p in stage_dir.rglob("*"):
                if p.is_file():
                    z.write(p, arcname=str(Path(pkg_base) / p.relative_to(stage_dir)))
        return out
    else:
        out = dist_dir / f"{pkg_base}.tar.gz"
        with tarfile.open(out, "w:gz") as tf:
            tf.add(stage_dir, arcname=pkg_base)
        return out


def main() -> int:
    ap = argparse.ArgumentParser(description="Build and package Rust binary + data/*.mmdb")
    ap.add_argument("--bin", default="", help="Binary target name (optional)")
    ap.add_argument("--target", default="", help="Cargo target triple (optional)")
    ap.add_argument("--dist-dir", default="dist", help="Output directory for archives")
    ap.add_argument("--stage-root", default="target/release-package", help="Staging directory root")
    ap.add_argument("--data-dir", default="data", help="Directory containing .mmdb files")
    args = ap.parse_args()

    meta = cargo_metadata()
    bins = find_bin_targets(meta)
    if not bins:
        print("ERROR: no binary targets found in workspace.", file=sys.stderr)
        return 2

    bin_name = args.bin.strip()
    if not bin_name:
        _, _, bin_name = bins[0]
        print(f"INFO: --bin not provided; defaulting to '{bin_name}'", file=sys.stderr)

    pkg_name, pkg_version = find_owner_package_for_bin(meta, bin_name)

    host_triple = rustc_host_triple()
    target_triple_for_name = args.target.strip() or host_triple

    # Build only the owning package
    build_cmd = ["cargo", "build", "-p", pkg_name, "--release"]
    if args.target.strip():
        build_cmd += ["--target", args.target.strip()]
    run(build_cmd)

    # Locate binary
    exe_suffix = exe_suffix_for_target(target_triple_for_name)
    bin_dir = Path("target") / (args.target.strip() if args.target.strip() else "release") / "release" \
        if args.target.strip() else Path("target") / "release"
    if args.target.strip():
        bin_dir = Path("target") / args.target.strip() / "release"

    bin_path = bin_dir / f"{bin_name}{exe_suffix}"
    if not bin_path.exists():
        print(f"ERROR: built binary not found: {bin_path}", file=sys.stderr)
        print("Hint: pass the correct bin target name: just release bin=<name>", file=sys.stderr)
        return 3

    out = stage_and_archive(
        pkg_name=pkg_name,
        pkg_version=pkg_version,
        bin_path=bin_path,
        data_dir=Path(args.data_dir),
        dist_dir=Path(args.dist_dir),
        stage_root=Path(args.stage_root),
        target_triple_for_name=target_triple_for_name,
    )

    print(f"Created: {out}")
    return 0


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