"""Run every image in a folder through /v1/image/authenticity and write one CSV.

    export TRUESECURIX_API_KEY=tsx_live_xxx
    python batch_image_authenticity.py ./images results.csv
    python batch_image_authenticity.py ./images results.csv --workers 4 --labels labels.csv

Safe to stop and re-run: images already in the CSV are skipped, and each request uses
the image's SHA-256 as its Idempotency-Key, so an image is never charged twice even if
a run is interrupted mid-request. Requires only httpx (`pip install httpx`).

With --labels (a CSV of `file,label` where label is real / ai_generated / ai_edited /
digitally_edited, spelled either way, e.g. "AI Generated") a summary is printed at the end:
accuracy per class, the confusion matrix, manipulations caught, genuine images wrongly
flagged, and the share sent for human review.

Nothing is skipped silently: files that are not JPEG, PNG or WebP (e.g. iPhone HEIC) are
listed as not sent, and images the API refuses as not analysable are reported separately from
failures worth retrying.
"""

from __future__ import annotations

import argparse
import csv
import hashlib
import os
import sys
import threading
from concurrent.futures import ThreadPoolExecutor, as_completed
from pathlib import Path

sys.path.insert(0, str(Path(__file__).resolve().parent))
from truesecurix import DEFAULT_BASE_URL, TrueSecurix, TrueSecurixError  # noqa: E402

CLASSES = ("real", "ai_generated", "ai_edited", "digitally_edited")
EXTENSIONS = {".jpg", ".jpeg", ".jfif", ".jpe", ".png", ".webp"}
COLUMNS = (["file", "decision", "classification", "label", "confidence", "review_recommended",
            "suspected", "suspected_score"]
           + [f"score_{c}" for c in CLASSES]
           + ["top_indicators", "regions", "id", "model_version", "processing_ms", "sha256", "error"])
# The API refused the image itself (unreadable, unsupported, too small or too large): sending
# the same bytes again cannot succeed, so a re-run keeps these rows instead of retrying them.
REJECTED = ("[400]", "[413]", "[422]")
# The server was at capacity or unreachable even after the client's own retries.
BUSY = ("[429]", "[502]", "[503]", "[504]")


def find_images(folder: Path) -> tuple[list[Path], dict[str, int]]:
    """Images to send, and every other file counted by why it is not sent."""
    images, skipped = [], {}
    for p in sorted(folder.rglob("*")):
        if not p.is_file():
            continue
        if p.name.startswith("._"):
            reason = "macOS metadata (._ files)"
        elif p.suffix.lower() in EXTENSIONS:
            images.append(p)
            continue
        elif p.suffix.lower() in {".heic", ".heif"}:
            reason = "HEIC/HEIF (convert to JPEG to include)"
        else:
            reason = f"{p.suffix.lower() or 'no extension'} (not JPEG, PNG or WebP)"
        skipped[reason] = skipped.get(reason, 0) + 1
    return images, skipped


def resume_rows(output: Path, folder: Path) -> list[dict]:
    """Rows of an earlier run that stand: every result, and every refusal whose file is unchanged
    (a refused image replaced by a fixed file of the same name is sent again). Other failures
    are dropped so they are retried."""
    if not output.exists():
        return []
    with output.open(newline="", encoding="utf-8") as f:
        rows = list(csv.DictReader(f))
    def stands(r: dict) -> bool:
        if not r.get("error"):
            return True
        path = folder / r["file"]
        return r["error"].startswith(REJECTED) and path.is_file() and sha256_file(path) == r["sha256"]
    return [r for r in rows if stands(r)]


def sha256_file(path: Path) -> str:
    h = hashlib.sha256()
    with path.open("rb") as f:
        for chunk in iter(lambda: f.read(1 << 20), b""):
            h.update(chunk)
    return h.hexdigest()


def row_for(rel: str, digest: str, result: dict | None, error: str = "") -> dict:
    if result is None:
        return {"file": rel, "sha256": digest, "error": error}
    indicators = result.get("indicators") or []
    top = " | ".join(f"{i['code']} ({i['strength']})" for i in indicators[:3])
    suspected = result.get("suspected") or {}
    return {
        "file": rel,
        "decision": result.get("decision") or outcome(result["classification"], result["review_recommended"]),
        "suspected": suspected.get("classification", ""),
        "suspected_score": suspected.get("score", ""),
        "classification": result["classification"],
        "label": result["label"],
        "confidence": result["confidence"],
        "review_recommended": result["review_recommended"],
        **{f"score_{c}": result["scores"].get(c) for c in CLASSES},
        "top_indicators": top,
        "regions": len(result.get("regions") or []),
        "id": result.get("id"),
        "model_version": result.get("model_version"),
        "processing_ms": result.get("processing_ms"),
        "sha256": digest,
        "error": "",
    }


LABEL_ALIASES = {
    "real": "real", "genuine": "real", "authentic": "real", "original": "real",
    "ai_generated": "ai_generated", "generated": "ai_generated", "synthetic": "ai_generated",
    "ai_edited": "ai_edited", "ai_modified": "ai_edited",
    "digitally_edited": "digitally_edited", "edited": "digitally_edited",
    "photoshopped": "digitally_edited", "tampered": "digitally_edited",
}


def outcome(classification: str, review) -> str:
    """flag (acted on as manipulated), review (sent to a person) or pass (accepted as genuine).
    An image answered Real with review recommended is a review, not a pass."""
    if str(review).lower() == "true":
        return "review"
    return "pass" if classification == "real" else "flag"


def normalise_label(raw: str) -> str | None:
    """'AI Generated', 'ai-generated' and 'ai_generated' all mean the same class."""
    key = "_".join(raw.strip().lower().replace("-", " ").split())
    return LABEL_ALIASES.get(key)


def summarise(results_csv: Path, labels_csv: Path) -> dict:
    """Print and return per-class accuracy, the confusion matrix and the fraud-team view."""
    truth, unknown = {}, set()
    with labels_csv.open(newline="", encoding="utf-8-sig") as f:
        for r in csv.DictReader(f):
            label = normalise_label(r["label"])
            if label is None:
                unknown.add(r["label"].strip())
            else:
                truth[r["file"].replace("\\", "/")] = label
    if unknown:
        print(f"ignored labels not recognised as a class: {sorted(unknown)}")
    with results_csv.open(newline="", encoding="utf-8") as f:
        every = list(csv.DictReader(f))
    rows = [r for r in every if not r["error"] and r["file"] in truth]
    sent = {r["file"] for r in every}
    errored = sum(1 for r in every if r["error"] and r["file"] in truth)
    coverage = {"labelled": len(truth), "with_result": len(rows), "errored": errored,
                "never_sent": len(set(truth) - sent),
                "results_without_label": sum(1 for r in every if not r["error"] and r["file"] not in truth)}
    print(f"\nlabels: {coverage['labelled']}; scored: {coverage['with_result']}; refused or failed: {errored}; "
          f"labelled but not in the results: {coverage['never_sent']}; "
          f"results with no label: {coverage['results_without_label']}")
    if not rows:
        print("no labelled results to summarise (do the file names in the labels match the results CSV?)")
        return {"coverage": coverage}
    present = [c for c in CLASSES if any(truth[r["file"]] == c for r in rows)]
    matrix = {t: {p: 0 for p in CLASSES} for t in present}
    for r in rows:
        matrix[truth[r["file"]]][r["classification"]] += 1
    correct = sum(matrix[c][c] for c in present)
    out = {"n": len(rows), "accuracy": correct / len(rows), "confusion": matrix, "per_class": {},
           "coverage": coverage}
    print(f"\n{len(rows)} labelled images, accuracy {correct / len(rows):.1%}")
    for c in present:
        n = sum(matrix[c].values())
        out["per_class"][c] = {"n": n, "recall": matrix[c][c] / n}
        print(f"  {c:17s} {matrix[c][c]}/{n} = {matrix[c][c] / n:.1%}")
    print("\n  confusion (rows = your label, columns = our classification)")
    print("  " + " " * 17 + "".join(f"{p[:16]:>17s}" for p in CLASSES))
    for t in present:
        print(f"  {t:17s}" + "".join(f"{matrix[t][p]:>17d}" for p in CLASSES))
    fake = [r for r in rows if truth[r["file"]] != "real"]
    real = [r for r in rows if truth[r["file"]] == "real"]
    act = {r["file"]: outcome(r["classification"], r["review_recommended"]) for r in rows}
    out["outcomes"] = {t: {o: sum(1 for r in rows if truth[r["file"]] == t and act[r["file"]] == o)
                           for o in ("flag", "review", "pass")} for t in present}
    print("\n  outcome (flag = acted on as manipulated, review = sent to a person, pass = accepted)")
    print("  " + " " * 17 + "".join(f"{o:>10s}" for o in ("flag", "review", "pass")))
    for t in present:
        print(f"  {t:17s}" + "".join(f"{out['outcomes'][t][o]:>10d}" for o in ("flag", "review", "pass")))
    out["review_rate"] = sum(a == "review" for a in act.values()) / len(rows)
    if fake:
        out["manipulation_caught"] = sum(r["classification"] != "real" for r in fake) / len(fake)
        out["manipulation_flagged_or_review"] = sum(act[r["file"]] != "pass" for r in fake) / len(fake)
        print(f"\n  manipulated images classed as a manipulation: {out['manipulation_caught']:.1%} of {len(fake)}")
        print(f"  manipulated images flagged or sent to review: "
              f"{out['manipulation_flagged_or_review']:.1%} of {len(fake)}")
    if real:
        out["genuine_wrongly_flagged"] = sum(r["classification"] != "real" for r in real) / len(real)
        out["genuine_sent_to_review"] = sum(act[r["file"]] == "review" for r in real) / len(real)
        print(f"  genuine images wrongly flagged: {out['genuine_wrongly_flagged']:.1%} of {len(real)}")
        print(f"  genuine images sent to review: {out['genuine_sent_to_review']:.1%} of {len(real)}")
    print(f"  sent for human review overall: {out['review_rate']:.1%}")
    return out


def main() -> int:
    ap = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
    ap.add_argument("folder", type=Path)
    ap.add_argument("output", type=Path)
    ap.add_argument("--workers", type=int, default=4, help="parallel requests (default 4; more does not finish sooner, the server analyses a fixed number at a time)")
    ap.add_argument("--labels", type=Path, help="optional CSV file,label for an accuracy summary")
    ap.add_argument("--base-url", default=os.environ.get("TRUESECURIX_BASE_URL", DEFAULT_BASE_URL))
    args = ap.parse_args()

    api_key = os.environ.get("TRUESECURIX_API_KEY")
    if not api_key:
        print("Set TRUESECURIX_API_KEY to your API key.", file=sys.stderr)
        return 2
    if not args.folder.is_dir():
        print(f"{args.folder} is not a folder.", file=sys.stderr)
        return 2

    images, skipped = find_images(args.folder)
    if skipped:
        print(f"not sent ({sum(skipped.values())} files):")
        for reason, n in sorted(skipped.items()):
            print(f"  {n:6d}  {reason}")
    kept = resume_rows(args.output, args.folder)
    if args.output.exists():
        with args.output.open("w", newline="", encoding="utf-8") as f:
            w = csv.DictWriter(f, fieldnames=COLUMNS, extrasaction="ignore")
            w.writeheader()
            w.writerows(kept)
    done = {r["file"] for r in kept}
    todo = [p for p in images if p.relative_to(args.folder).as_posix() not in done]
    print(f"{len(images)} images found, {len(done)} already done, {len(todo)} to send")

    new_file = not args.output.exists()
    lock = threading.Lock()
    counts = {"ok": 0, "failed": 0, "refused": 0}
    with args.output.open("a", newline="", encoding="utf-8") as out, \
            TrueSecurix(api_key, base_url=args.base_url, timeout=180.0, max_retries=4) as client:
        writer = csv.DictWriter(out, fieldnames=COLUMNS)
        if new_file:
            writer.writeheader()

        def work(path: Path) -> dict:
            rel = path.relative_to(args.folder).as_posix()
            digest = sha256_file(path)
            try:
                result = client.image_authenticity(str(path), reference=rel[:120], idempotency_key=digest)
                return row_for(rel, digest, result)
            except (TrueSecurixError, OSError, ValueError) as exc:
                return row_for(rel, digest, None, str(exc)[:300])

        def record(row: dict) -> None:
            with lock:
                writer.writerow(row)
                out.flush()
                err = row["error"]
                counts["ok" if not err else "refused" if err.startswith(REJECTED) else "failed"] += 1

        busy = []
        with ThreadPoolExecutor(max_workers=max(1, args.workers)) as pool:
            for i, fut in enumerate(as_completed([pool.submit(work, p) for p in todo]), 1):
                row = fut.result()
                if row["error"].startswith(BUSY):
                    busy.append(row)
                else:
                    record(row)
                if i % 25 == 0 or i == len(todo):
                    print(f"  {i}/{len(todo)} done ({counts['refused']} refused, {counts['failed']} failed, "
                          f"{len(busy)} to retry)", flush=True)
        if busy:
            # The server was full while many requests were in flight. One at a time it has room,
            # so these finish in this run instead of being left for the user to re-run.
            print(f"  retrying {len(busy)} images one at a time (the server was busy)", flush=True)
            for row in busy:
                record(work(args.folder / row["file"]))

    print(f"finished: {counts['ok']} classified, {counts['refused']} refused, {counts['failed']} failed. "
          f"Results in {args.output}")
    if counts["refused"]:
        print("Refused images are not analysable as sent (the error column says why: unreadable, "
              "unsupported, too small or too large); they were not charged and a re-run will not change them.")
    if counts["failed"]:
        print("Re-run the same command to retry the failed images; nothing is charged twice.")
    if args.labels:
        summarise(args.output, args.labels)
    return 0 if not counts["failed"] else 1


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