"""
Combine per-source manifests into a single master manifest.

The master manifest is the authoritative legal record for the dataset. It is:
  • Committed alongside the trained model (so the lineage is auditable).
  • Used by the training script (no image is loaded that isn't in the manifest).
  • The first thing a lawyer or auditor will ask to see.

Stage 3A manifest schema
------------------------
Required legal/training fields:
  path, class, source, license, license_url, sha256

Optional generator metadata fields:
  generator, model_family, model_id, prompt, seed, width, height,
  generation_params_json

Missing optional values are written as empty strings. This keeps Stage 2
manifest fragments readable while giving Stage 3A scripts stable columns for
per-generator evaluation.

Usage
-----
    python scripts/dataset/build_manifest.py \
        --inputs data/raw/real_manifest.csv data/raw/ai_generated_manifest.csv \
        --out data/manifest.csv

    python scripts/dataset/build_manifest.py \
        --input-dir data/raw \
        --out data/manifest.csv
"""
from __future__ import annotations

import argparse
import csv
from pathlib import Path

from generation_utils import APPROVED_GENERATORS

REQUIRED_FIELDS = ["path", "class", "source", "license", "license_url", "sha256"]
OPTIONAL_STAGE3A_FIELDS = [
    "generator",
    "model_family",
    "model_id",
    "prompt",
    "seed",
    "width",
    "height",
    "generation_params_json",
]
MANIFEST_FIELDS = REQUIRED_FIELDS + OPTIONAL_STAGE3A_FIELDS
MANIFEST_FRAGMENT_GLOB = "*_manifest.csv"
GENERATOR_SPECS_BY_SOURCE = {
    spec.source: spec for spec in APPROVED_GENERATORS.values()
}


def _normalise_row(row: dict[str, str]) -> dict[str, str]:
    """Return a manifest row with all Stage 3A optional fields present."""
    normalised = dict(row)
    source = normalised.get("source", "")
    spec = GENERATOR_SPECS_BY_SOURCE.get(source)
    if normalised.get("class") == "ai_generated" and spec is not None:
        if not normalised.get("generator"):
            normalised["generator"] = spec.generator
        if not normalised.get("model_family"):
            normalised["model_family"] = spec.model_family
        if not normalised.get("model_id"):
            normalised["model_id"] = spec.model_id

    for field in OPTIONAL_STAGE3A_FIELDS:
        if normalised.get(field) is None:
            normalised[field] = ""
        else:
            normalised[field] = str(normalised.get(field, ""))
    return normalised


def _discover_manifest_fragments(input_dirs: list[Path], out: Path) -> list[Path]:
    """Find manifest fragments below each input directory in stable order."""
    discovered: list[Path] = []
    out_resolved = out.resolve()
    for input_dir in input_dirs:
        if not input_dir.exists():
            raise FileNotFoundError(f"Input directory does not exist: {input_dir}")
        if not input_dir.is_dir():
            raise NotADirectoryError(f"Input path is not a directory: {input_dir}")
        for path in sorted(input_dir.rglob(MANIFEST_FRAGMENT_GLOB)):
            if path.resolve() == out_resolved:
                continue
            discovered.append(path)
    return discovered


def _unique_paths(paths: list[Path]) -> list[Path]:
    """Deduplicate paths while preserving caller/discovery order."""
    unique: list[Path] = []
    seen: set[Path] = set()
    for path in paths:
        resolved = path.resolve()
        if resolved in seen:
            continue
        seen.add(resolved)
        unique.append(path)
    return unique


def main() -> None:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument(
        "--inputs",
        type=Path,
        nargs="+",
        default=[],
        help="Explicit manifest fragments to merge",
    )
    parser.add_argument(
        "--input-dir",
        type=Path,
        action="append",
        default=[],
        help=(
            f"Directory to scan recursively for {MANIFEST_FRAGMENT_GLOB}; "
            "can be passed multiple times"
        ),
    )
    parser.add_argument("--out", type=Path, required=True)
    args = parser.parse_args()

    input_paths = _unique_paths(
        [*args.inputs, *_discover_manifest_fragments(args.input_dir, args.out)]
    )
    if not input_paths:
        raise ValueError("Provide at least one --inputs file or --input-dir")

    rows: list[dict] = []
    for src in input_paths:
        if not src.exists():
            raise FileNotFoundError(f"Input manifest does not exist: {src}")
        with src.open() as fh:
            reader = csv.DictReader(fh)
            for r in reader:
                missing = [f for f in REQUIRED_FIELDS if not r.get(f)]
                if missing:
                    raise ValueError(
                        f"{src}: row missing required fields {missing}: {r}"
                    )
                rows.append(_normalise_row(r))

    # Detect duplicates by sha256 — important for license cleanliness AND
    # to avoid train/test leakage.
    seen: dict[str, str] = {}
    deduped: list[dict] = []
    for r in rows:
        sha = r["sha256"]
        if sha in seen:
            print(f"  dropping duplicate {r['path']} (matches {seen[sha]})")
            continue
        seen[sha] = r["path"]
        deduped.append(r)

    args.out.parent.mkdir(parents=True, exist_ok=True)
    # Write a stable schema first, then any extra legacy/source-specific fields.
    fieldnames = sorted({k for r in deduped for k in r.keys()})
    ordered = MANIFEST_FIELDS + [f for f in fieldnames if f not in MANIFEST_FIELDS]

    with args.out.open("w", newline="") as fh:
        writer = csv.DictWriter(fh, fieldnames=ordered)
        writer.writeheader()
        writer.writerows(deduped)

    classes: dict[str, int] = {}
    for r in deduped:
        classes[r["class"]] = classes.get(r["class"], 0) + 1

    print(f"\nMaster manifest: {args.out}")
    print(f"  fragments: {len(input_paths)}")
    print(f"  total rows: {len(deduped)}")
    for cls, n in sorted(classes.items()):
        print(f"  {cls}: {n}")

    generators: dict[str, int] = {}
    for r in deduped:
        if r["class"] != "ai_generated":
            continue
        generator = r.get("generator") or r["source"]
        generators[generator] = generators.get(generator, 0) + 1
    for generator, n in sorted(generators.items()):
        print(f"  generator {generator}: {n}")


if __name__ == "__main__":
    main()