#!/usr/bin/env python3
"""Create a metadata-free, consent-traceable Timmy training record."""
import argparse
import hashlib
import hmac
import io
import json
import os
import sys
from datetime import datetime, timezone
from pathlib import Path

from PIL import Image, ImageOps

COLORS = ("brown", "green", "yellow", "pale", "red", "black")
QUALITIES = ("good", "fair", "poor")
REVIEWERS = ("user-confirmed", "clinician-reviewed")


def parser() -> argparse.ArgumentParser:
    p = argparse.ArgumentParser()
    p.add_argument("--input", required=True)
    p.add_argument("--output-dir", required=True)
    p.add_argument("--subject-id", required=True)
    p.add_argument("--consent-version", required=True)
    p.add_argument("--bristol-type", required=True, type=int, choices=range(1, 8))
    p.add_argument("--color", required=True, choices=COLORS)
    p.add_argument("--quality", required=True, choices=QUALITIES)
    p.add_argument("--reviewer", required=True, choices=REVIEWERS)
    return p


def main() -> int:
    args = parser().parse_args()
    salt = os.environ.get("TIMMY_DATASET_SALT", "")
    if len(salt) < 8:
        raise SystemExit("TIMMY_DATASET_SALT must contain at least 8 characters")
    if not args.consent_version.strip():
        raise SystemExit("consent version is required")

    try:
        with Image.open(args.input) as source:
            image = ImageOps.exif_transpose(source).convert("RGB")
            image.thumbnail((1200, 1200))
            buffer = io.BytesIO()
            image.save(buffer, "JPEG", quality=85, optimize=True, exif=b"")
    except Exception as exc:
        raise SystemExit(f"input is not a decodable image: {exc}") from exc

    data = buffer.getvalue()
    image_id = hashlib.sha256(data).hexdigest()
    subject_key = hmac.new(salt.encode(), args.subject_id.encode(), hashlib.sha256).hexdigest()[:16]
    bucket = int(subject_key[:2], 16)
    split = "train" if bucket < 205 else "validation" if bucket < 230 else "test"

    output = Path(args.output_dir)
    images = output / "images"
    images.mkdir(parents=True, exist_ok=True)
    relative_path = f"images/{image_id}.jpg"
    target = output / relative_path
    if not target.exists():
        target.write_bytes(data)

    record = {
        "schemaVersion": 1,
        "imageId": image_id,
        "relativePath": relative_path,
        "subjectKey": subject_key,
        "split": split,
        "consentVersion": args.consent_version.strip(),
        "ingestedAt": datetime.now(timezone.utc).isoformat(),
        "labels": {
            "bristolType": args.bristol_type,
            "color": args.color,
            "quality": args.quality,
        },
        "review": {"status": args.reviewer},
        "derivation": {
            "format": "jpeg",
            "maxDimension": 1200,
            "metadataStripped": True,
        },
    }
    with (output / "manifest.jsonl").open("a", encoding="utf-8") as manifest:
        manifest.write(json.dumps(record, separators=(",", ":")) + "\n")
    print(json.dumps(record))
    return 0


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