"""Reproducible synthetic-person dataset, SDXL LoRA training, and comparison."""

import argparse
import hashlib
import importlib.metadata
import json
import math
import os
import shutil
import signal
import subprocess
import sys
import time
from datetime import datetime, timezone
from pathlib import Path

try:
    import tomllib
except ModuleNotFoundError:
    import toml as tomllib


ROOT = Path(__file__).resolve().parent
TOKEN = "vprsn01"
PERSON = (
    "a fictional 29 year old East Asian woman, oval face, almond shaped dark brown eyes, "
    "straight eyebrows, long black hair, a few subtle freckles on her cheeks"
)
NEGATIVE = (
    "illustration, cartoon, painting, plastic skin, blurry face, deformed face, "
    "bad anatomy, extra limbs, duplicate person, multiple people, text, watermark"
)
SCENES = [
    "close-up portrait, white cotton shirt, neutral gray studio background, soft daylight",
    "three-quarter face portrait, green sweater, garden background, overcast daylight",
    "left side profile photograph, looking completely to the left, only one eye visible, navy jacket, pale studio background, side lighting",
    "right side profile photograph, looking completely to the right, only one eye visible, red blouse, brick wall background, open shade",
    "head and shoulders, charcoal turtleneck, window light, indoor background",
    "medium shot from head to hips, standing with both hands visible, blue denim shirt and jeans, park background, morning sunlight",
    "medium shot from head to hips, holding an open book with both hands, white blouse, bookstore interior, diffused window light",
    "medium shot from head to hips, arms loosely folded, gray blazer, plain white wall, natural daylight",
    "smiling head and shoulders portrait, green shirt, trees in the background",
    "serious expression, head and shoulders, blue sweater, soft studio lighting",
    "head slightly tilted, waist-up portrait, red cardigan, cafe interior",
    "seated waist-up portrait, black jacket, concrete wall, cloudy daylight",
]
EVAL_SCENES = [
    "head and shoulders, yellow raincoat, rainy city street, natural skin texture",
    "waist-up portrait, violet formal dress, theater lobby, evening light",
    "head and shoulders, white winter coat, snowy forest, cold daylight",
    "three-quarter face facing left, white shirt, cyan studio background",
    "waist-up portrait, striped shirt, seaside promenade, bright daylight",
    "seated portrait, black suit, modern library, window lighting",
]


def read_json(path):
    return json.loads(Path(path).read_text(encoding="utf-8"))


def save_json(path, data):
    path = Path(path)
    path.parent.mkdir(parents=True, exist_ok=True)
    path.write_text(json.dumps(data, ensure_ascii=False, indent=2) + "\n", encoding="utf-8")


def sha256(path):
    with Path(path).open("rb") as stream:
        return hashlib.file_digest(stream, "sha256").hexdigest() if sys.version_info >= (3, 11) else _hash_stream(stream)


def _hash_stream(stream):
    digest = hashlib.sha256()
    for chunk in iter(lambda: stream.read(1024 * 1024), b""):
        digest.update(chunk)
    return digest.hexdigest()


def event(stage, **fields):
    record = {"time_utc": datetime.now(timezone.utc).isoformat(), "stage": stage, **fields}
    (ROOT / "logs").mkdir(exist_ok=True)
    with (ROOT / "logs/events.jsonl").open("a", encoding="utf-8") as stream:
        stream.write(json.dumps(record, ensure_ascii=False) + "\n")
    print(json.dumps(record, ensure_ascii=False), flush=True)


def contained_file(parent, value):
    parent = Path(parent).resolve()
    path = (parent / value).resolve()
    if path.parent != parent or not path.is_file():
        raise ValueError(f"Expected a file directly in {parent}: {value}")
    return path


def preflight():
    import torch

    versions = {}
    for name in ["torch", "diffusers", "transformers", "accelerate", "peft", "bitsandbytes"]:
        try:
            versions[name] = importlib.metadata.version(name)
        except importlib.metadata.PackageNotFoundError:
            versions[name] = None
    info = {
        "python": sys.version,
        "packages": versions,
        "cuda_available": torch.cuda.is_available(),
        "cuda_runtime": torch.version.cuda,
        "free_disk_gb": round(shutil.disk_usage(ROOT).free / 1024**3, 2),
    }
    if info["cuda_available"]:
        info["gpu"] = torch.cuda.get_device_name(0)
        info["vram_gb"] = round(torch.cuda.get_device_properties(0).total_memory / 1024**3, 2)
    save_json(ROOT / "logs/environment.json", info)
    event("preflight", **info)
    if not info["cuda_available"]:
        raise RuntimeError("NVIDIA CUDA is unavailable; no generation or training started")
    if not torch.cuda.is_bf16_supported():
        raise RuntimeError("This configuration requires a BF16-capable GPU")
    if any(value is None for value in versions.values()):
        raise RuntimeError("Run setup.sh to install the missing packages")


def download():
    from huggingface_hub import snapshot_download

    lock = read_json(ROOT / "model-lock.json")
    if shutil.disk_usage(ROOT).free < 18 * 1024**3:
        raise RuntimeError("At least 18 GiB free space required before the first download")
    for key, patterns in [
        ("base", ["*.json", "tokenizer/*", "tokenizer_2/*", "README.md", lock["base"]["checkpoint"]]),
        ("adapter", ["models/image_encoder/config.json", "models/image_encoder/model.safetensors",
                     "sdxl_models/" + lock["adapter"]["weight"], "README.md"]),
    ]:
        event("download_start", model=key, revision=lock[key]["revision"])
        snapshot_download(lock[key]["repo"], revision=lock[key]["revision"],
                          allow_patterns=patterns, local_dir=ROOT / "models" / key, max_workers=2)
    tokenizer_names = ["openai_clip-vit-large-patch14", "laion_CLIP-ViT-bigG-14-laion2B-39B-b160k"]
    for source, name in zip(["tokenizer", "tokenizer_2"], tokenizer_names):
        shutil.copytree(ROOT / "models/base" / source, ROOT / "models/tokenizers" / name, dirs_exist_ok=True)
    weights = [ROOT / "models/base" / lock["base"]["checkpoint"],
               ROOT / "models/adapter/models/image_encoder/model.safetensors",
               ROOT / "models/adapter/sdxl_models" / lock["adapter"]["weight"]]
    save_json(ROOT / "logs/model-hashes.json", {str(p.relative_to(ROOT)): sha256(p) for p in weights})
    event("download_complete")


def pipeline(with_adapter=False):
    import torch
    from diffusers import EulerDiscreteScheduler, StableDiffusionXLPipeline

    if not torch.cuda.is_available():
        raise RuntimeError("CUDA is required; refusing accidental CPU generation")
    lock = read_json(ROOT / "model-lock.json")
    kwargs = {}
    if with_adapter:
        from transformers import CLIPVisionModelWithProjection
        kwargs["image_encoder"] = CLIPVisionModelWithProjection.from_pretrained(
            ROOT / "models/adapter/models/image_encoder", torch_dtype=torch.float16, local_files_only=True)
    pipe = StableDiffusionXLPipeline.from_single_file(
        str(ROOT / "models/base" / lock["base"]["checkpoint"]),
        config=str(ROOT / "models/base"), torch_dtype=torch.float16,
        local_files_only=True, **kwargs)
    pipe.scheduler = EulerDiscreteScheduler.from_config(pipe.scheduler.config)
    if with_adapter:
        pipe.load_ip_adapter(str(ROOT / "models/adapter"), subfolder="sdxl_models",
                             weight_name=lock["adapter"]["weight"], image_encoder_folder=None,
                             local_files_only=True)
    pipe.to("cuda")
    return pipe


def render(pipe, destination, prompt, seed, negative_prompt=NEGATIVE, **kwargs):
    import torch

    destination = Path(destination)
    settings = {"prompt": prompt, "negative_prompt": negative_prompt, "seed": seed,
                "steps": 30, "guidance_scale": 5.0, "width": 1024, "height": 1024,
                "scheduler": type(pipe.scheduler).__name__, "model_lock": read_json(ROOT / "model-lock.json")}
    if destination.exists():
        prior = read_json(destination.with_suffix(".json"))
        if any(prior.get(key) != value for key, value in settings.items()) or prior.get("sha256") != sha256(destination):
            raise ValueError(f"Existing image no longer matches its recipe: {destination}")
        return
    start = time.monotonic()
    image = pipe(prompt=prompt, negative_prompt=negative_prompt, num_inference_steps=30,
                 guidance_scale=5.0, width=1024, height=1024,
                 generator=torch.Generator(device="cpu").manual_seed(seed), **kwargs).images[0]
    destination.parent.mkdir(parents=True, exist_ok=True)
    image.save(destination)
    settings.update({"elapsed_seconds": round(time.monotonic() - start, 3), "sha256": sha256(destination)})
    save_json(destination.with_suffix(".json"), settings)
    event("image_saved", file=str(destination.relative_to(ROOT)), elapsed_seconds=settings["elapsed_seconds"])


def contact_sheet(paths, destination):
    from PIL import Image, ImageDraw, ImageOps

    paths = list(paths)
    if not paths:
        return
    width, cell, label = 4, 256, 30
    sheet = Image.new("RGB", (width * cell, math.ceil(len(paths) / width) * (cell + label)), "white")
    draw = ImageDraw.Draw(sheet)
    for index, path in enumerate(paths):
        x, y = index % width * cell, index // width * (cell + label)
        with Image.open(path) as image:
            sheet.paste(ImageOps.contain(image.convert("RGB"), (cell, cell)), (x, y))
        draw.text((x + 6, y + cell + 6), path.stem[:34], fill="black")
    sheet.save(destination)


def anchors(count):
    pipe = pipeline()
    for index in range(count):
        destination = ROOT / "anchors" / f"anchor_{index:02d}.png"
        render(pipe, destination, f"professional photograph of {PERSON}, {SCENES[0]}", 20260906 + index)
    contact_sheet(sorted((ROOT / "anchors").glob("anchor_*.png")), ROOT / "anchors/contact-sheet.jpg")


def expand(limit, scale):
    from PIL import Image

    review = read_json(ROOT / "review.json")
    if not review["anchor"] or not review["reference_crop"]:
        raise ValueError("Inspect anchors, then set anchor and reference_crop [left, top, right, bottom] in review.json")
    anchor = contained_file(ROOT / "anchors", review["anchor"])
    with Image.open(anchor) as image:
        left, top, right, bottom = review["reference_crop"]
        if not (0 <= left < right <= image.width and 0 <= top < bottom <= image.height):
            raise ValueError("Reference crop is outside the anchor image")
        reference = image.convert("RGB").crop((left, top, right, bottom))
    (ROOT / "candidates").mkdir(exist_ok=True)
    recipe = {"anchor": review["anchor"], "anchor_sha256": sha256(anchor),
              "reference_crop": review["reference_crop"], "adapter_scale": scale,
              "model_lock": read_json(ROOT / "model-lock.json")}
    recipe_path = ROOT / "candidates/recipe.json"
    if recipe_path.exists() and read_json(recipe_path) != recipe:
        raise ValueError("Candidate recipe changed; use a separate experiment directory")
    save_json(recipe_path, recipe)
    reference.save(ROOT / "candidates/reference-crop.png")
    pipe = pipeline(with_adapter=True)
    pipe.set_ip_adapter_scale(scale)
    for index in range(limit):
        destination = ROOT / "candidates" / f"candidate_{index:02d}.png"
        scene_index = index % len(SCENES)
        scene = SCENES[scene_index]
        negative = NEGATIVE
        if scene_index in {2, 3}:
            negative += ", frontal face, looking at camera, front view"
        elif scene_index in {5, 6, 7, 10, 11}:
            negative += ", close-up, headshot, cropped torso, cropped hands"
        render(pipe, destination, f"{scene}, professional photograph of {PERSON}",
               20300000 + index, negative_prompt=negative, ip_adapter_image=reference)
    contact_sheet(sorted((ROOT / "candidates").glob("candidate_*.png")), ROOT / "candidates/contact-sheet.jpg")
    event("candidates_ready", count=limit, next="Inspect images and write per-image acceptance and corrected captions")


def accepted_images(review, parent):
    accepted = []
    seen = set()
    for item in review["images"]:
        if item.get("accept") is not True:
            continue
        path = contained_file(parent, item["image"])
        if path.name in seen:
            raise ValueError(f"Duplicate accepted image: {path.name}")
        seen.add(path.name)
        if item.get("identity_ok") is not True or item.get("quality_ok") is not True:
            raise ValueError(f"Missing identity/quality review: {path.name}")
        caption = item["caption"].strip()
        if TOKEN not in caption.split() or len(caption.split()) < 6:
            raise ValueError(f"Caption needs trigger {TOKEN} and visible scene details: {path.name}")
        accepted.append((path, caption))
    if not 20 <= len(accepted) <= 30:
        raise ValueError(f"Need 20-30 reviewed images, found {len(accepted)}")
    if len({sha256(p) for p, _ in accepted}) != len(accepted):
        raise ValueError("Exact duplicate images in accepted dataset")
    return accepted


def curate():
    from PIL import Image

    accepted = accepted_images(read_json(ROOT / "review.json"), ROOT / "candidates")
    target = ROOT / "dataset"
    if target.exists():
        raise FileExistsError("Dataset already exists; preserve it and use a new experiment directory for revisions")
    target.mkdir()
    records = []
    for source, caption in accepted:
        with Image.open(source) as image:
            image.convert("RGB").save(target / source.name)
        (target / source.with_suffix(".txt").name).write_text(caption + "\n", encoding="utf-8")
        records.append({"image": source.name, "sha256": sha256(target / source.name), "caption": caption})
    save_json(target / "manifest.json", records)
    contact_sheet([target / p.name for p, _ in accepted], target / "contact-sheet.jpg")
    # Keep the review sheet outside the training directory; every image there is a sample.
    (target / "contact-sheet.jpg").replace(ROOT / "evidence/dataset-contact-sheet.jpg")
    event("dataset_curated", images=len(records))


def validate_dataset(directory):
    directory = Path(directory)
    records = read_json(directory / "manifest.json")
    if not 20 <= len(records) <= 30:
        raise ValueError("Dataset size is outside the reviewed range")
    expected = {row["image"] for row in records}
    actual = {p.name for p in directory.iterdir() if p.suffix.lower() in {".png", ".jpg", ".jpeg", ".webp"}}
    if len(expected) != len(records) or expected != actual:
        raise ValueError("Dataset images differ from the reviewed manifest")
    for row in records:
        path = contained_file(directory, row["image"])
        if sha256(path) != row["sha256"] or path.with_suffix(".txt").read_text(encoding="utf-8").strip() != row["caption"]:
            raise ValueError(f"Dataset changed after review: {path.name}")
    return records


def train(smoke):
    import toml

    preflight()
    validate_dataset(ROOT / "dataset")
    lock = read_json(ROOT / "model-lock.json")
    checkout = ROOT / "sd-scripts"
    commit = subprocess.check_output(["git", "-C", str(checkout), "rev-parse", "HEAD"], text=True).strip()
    if commit != lock["trainer"]["commit"]:
        raise ValueError("Trainer revision differs from model-lock.json")
    output = ROOT / "outputs" / ("smoke" if smoke else "run01")
    if output.exists():
        raise FileExistsError("Output run already exists; inspect it before resuming or creating another run")
    output.mkdir(parents=True)
    dataset_config = {"general": {"caption_extension": ".txt", "shuffle_caption": False},
                      "datasets": [{"resolution": 1024, "batch_size": 1, "enable_bucket": True,
                                    "bucket_no_upscale": True, "subsets": [{"image_dir": str(ROOT / "dataset"),
                                    "num_repeats": 1, "caption_dropout_rate": 0.0}]}]}
    dataset_path = output / "dataset.toml"
    dataset_path.write_text(toml.dumps(dataset_config), encoding="utf-8")
    config = tomllib.loads((ROOT / "train.toml").read_text(encoding="utf-8"))
    config.update({"pretrained_model_name_or_path": str(ROOT / "models/base" / lock["base"]["checkpoint"]),
                   "dataset_config": str(dataset_path), "output_dir": str(output), "output_name": TOKEN,
                   "logging_dir": str(output / "tensorboard"), "tokenizer_cache_dir": str(ROOT / "models/tokenizers")})
    if smoke:
        config.update(max_train_steps=20, save_every_n_steps=20)
    config_path = output / "train-resolved.toml"
    config_path.write_text(toml.dumps(config), encoding="utf-8")
    command = [sys.executable, "-m", "accelerate.commands.launch", "--num_processes=1", "--num_machines=1",
               "--num_cpu_threads_per_process=2", "--mixed_precision=bf16", "sdxl_train_network.py", "--config_file", str(config_path)]
    save_json(output / "command.json", {"argv": command, "cwd": str(checkout), "max_seconds": 600 if smoke else 5400})
    save_json(output / "dataset-manifest.json", validate_dataset(ROOT / "dataset"))
    event("training_start", output=str(output), smoke=smoke)
    start = time.monotonic()
    try:
        with (output / "train.log").open("w", encoding="utf-8") as stream:
            process = subprocess.Popen(command, cwd=checkout, stdout=stream,
                                       stderr=subprocess.STDOUT, start_new_session=True)
            save_json(output / "process.json", {"pid": process.pid, "started_at_utc": datetime.now(timezone.utc).isoformat()})
            try:
                returncode = process.wait(timeout=600 if smoke else 5400)
            except (subprocess.TimeoutExpired, KeyboardInterrupt):
                # Accelerate starts child workers; stop the whole group on timeout.
                os.killpg(process.pid, signal.SIGTERM)
                try:
                    process.wait(timeout=10)
                except subprocess.TimeoutExpired:
                    os.killpg(process.pid, signal.SIGKILL)
                    process.wait()
                raise
            if returncode:
                raise subprocess.CalledProcessError(returncode, command)
    except (subprocess.CalledProcessError, subprocess.TimeoutExpired) as error:
        event("training_failed", output=str(output), error=str(error), elapsed_seconds=time.monotonic() - start)
        raise
    weights = sorted(output.glob("*.safetensors"))
    if not weights:
        raise RuntimeError("Trainer exited successfully but produced no LoRA weights")
    save_json(output / "weight-hashes.json", {p.name: sha256(p) for p in weights})
    event("training_complete", output=str(output), elapsed_seconds=round(time.monotonic() - start, 3))


def evaluate(weight_paths):
    pipe = pipeline(with_adapter=False)
    conditions = [("base", None)]
    for weight in weight_paths:
        path = Path(weight).resolve(strict=True)
        conditions.append((path.stem, path))
    for name, weight in conditions:
        directory = ROOT / "evaluation" / name
        if directory.exists():
            raise FileExistsError(f"Evaluation already exists: {directory}")
        if weight:
            pipe.load_lora_weights(str(weight.parent), weight_name=weight.name, adapter_name="portrait")
            pipe.set_adapters("portrait", adapter_weights=0.8)
        paths = []
        for index, scene in enumerate(EVAL_SCENES):
            path = directory / f"eval_{index:02d}.png"
            render(pipe, path, f"photo of {TOKEN} woman, {scene}", 31415000 + index)
            paths.append(path)
        contact_sheet(paths, directory / "contact-sheet.jpg")
        save_json(directory / "condition.json", {"weight": str(weight) if weight else None,
                  "sha256": sha256(weight) if weight else None, "lora_scale": 0.8 if weight else 0,
                  "reference_image": None, "ip_adapter": False, "model_lock": read_json(ROOT / "model-lock.json")})
        if weight:
            pipe.unload_lora_weights()
    event("evaluation_complete", conditions=len(conditions), images_per_condition=len(EVAL_SCENES))


def generate(weight, prompt, seed):
    if TOKEN not in prompt.split():
        raise ValueError(f"Include the trained trigger {TOKEN} in the prompt")
    weight = Path(weight).resolve(strict=True)
    pipe = pipeline(with_adapter=False)
    pipe.load_lora_weights(str(weight.parent), weight_name=weight.name, adapter_name="portrait")
    pipe.set_adapters("portrait", adapter_weights=0.8)
    name = datetime.now(timezone.utc).strftime("%Y%m%dT%H%M%S%f")
    destination = ROOT / "generated" / f"{name}-{seed}.png"
    render(pipe, destination, prompt, seed)
    metadata = read_json(destination.with_suffix(".json"))
    metadata.update({"lora_sha256": sha256(weight), "lora_scale": 0.8, "ip_adapter": False})
    save_json(destination.with_suffix(".json"), metadata)


def main():
    parser = argparse.ArgumentParser(description=__doc__)
    sub = parser.add_subparsers(dest="stage", required=True)
    for stage in ["preflight", "download", "curate"]:
        sub.add_parser(stage)
    anchor = sub.add_parser("anchors")
    anchor.add_argument("--count", type=int, default=4, choices=range(1, 9))
    expansion = sub.add_parser("expand")
    expansion.add_argument("--limit", type=int, default=6, choices=range(1, 49))
    expansion.add_argument("--scale", type=float, default=0.75)
    training = sub.add_parser("train")
    training.add_argument("--smoke", action="store_true")
    evaluation = sub.add_parser("evaluate")
    evaluation.add_argument("weights", nargs="*")
    generation = sub.add_parser("generate")
    generation.add_argument("--weight", required=True)
    generation.add_argument("--prompt", required=True)
    generation.add_argument("--seed", type=int, default=42)
    args = parser.parse_args()
    (ROOT / "evidence").mkdir(exist_ok=True)
    if args.stage == "preflight":
        preflight()
    elif args.stage == "download":
        download()
    elif args.stage == "anchors":
        anchors(args.count)
    elif args.stage == "expand":
        if not 0 < args.scale <= 1:
            parser.error("adapter scale must be in (0, 1]")
        expand(args.limit, args.scale)
    elif args.stage == "curate":
        curate()
    elif args.stage == "train":
        train(args.smoke)
    elif args.stage == "evaluate":
        evaluate(args.weights)
    elif args.stage == "generate":
        generate(args.weight, args.prompt, args.seed)


if __name__ == "__main__":
    main()
