#!/usr/bin/env python3
"""Render one Pepe the Frog illustration with Stable Diffusion 1.5 + three public LoRAs on CPU.

Run from the repository root:

    python render.py

Downloads any missing weights into ./models, verifies every file against its pinned sha256
(stopping on a mismatch), renders with fixed settings and writes into ./artifacts:

    pepe.png     raw 512x512 output, settings in the PNG text chunk "parameters"
    render.py    a copy of this exact script
    params.json  settings, file hashes, library versions, hardware and timings
    run.log      everything printed to the console during the run

If the run fails at any point, run.log and params.json are still written with the stage
where it stopped.
"""

import hashlib
import json
import os
import platform
import shutil
import sys
import time
import traceback
import urllib.request

T_START = time.time()

MODELS_DIR = os.environ.get("MODELS_DIR", "models")
CONFIG_DIR = os.environ.get("SD15_CONFIG_DIR", "sd15-config")
OUT_DIR = os.environ.get("OUT_DIR", "artifacts")

BASE_MODEL = {
    "name": "v1-5-pruned-emaonly.safetensors",
    "repo": "stable-diffusion-v1-5/stable-diffusion-v1-5",
    "url": "https://huggingface.co/stable-diffusion-v1-5/stable-diffusion-v1-5/resolve/main/v1-5-pruned-emaonly.safetensors",
    "sha256": "6ce0161689b3853acaa03779ec93eafe75a02f4ced659bee03f50797806fa2fa",
}

LORAS = [
    {
        "name": "pepe-frog",
        "adapter": "pepe_frog",
        "file": "pepe-frog.safetensors",
        "url": "https://huggingface.co/stablediffusionapi/load_lora_embeddings/resolve/main/pepe-frog.safetensors",
        "sha256": "b777ec0c7b40ea6f5c7fc324c497042df41b89581223ca23354d3f217622cd04",
        "weight": 0.6,
    },
    {
        "name": "Line Art + Flat Colors v2",
        "adapter": "lineart_flat_colors",
        "file": "lineart_flat_colors_v2.safetensors",
        "url": "https://civitai.com/api/download/models/147989?fileId=109723",
        "sha256": "9a51a3d5f74740cec30599fdc695e7dd7ee96e43940b1e0637d69c02165fd2be",
        "weight": 0.7,
    },
    {
        "name": "Retro - Style",
        "adapter": "retro_style",
        "file": "retro_style.safetensors",
        "url": "https://civitai.com/api/download/models/69897?fileId=48789",
        "sha256": "c07d26d0a564b6214783dbe78ee34063ece5d8d95159a01b18e90097364ed201",
        "weight": 0.6,
    },
]

PROMPT = (
    "pepe_frog, pepe the frog, cartoon character, 1970s cartoon, old printed comic, faded colors, "
    "flat poster colors, heavy sleepy half closed eyelids, small black pupils, wide red lips, "
    "two eyes with white eyeballs and small round black pupils, lineart, flat colors"
)
NEGATIVE = (
    "photo, realistic, 3d render, plush, knitted, felt, toy, huge pupils, big black eyes, "
    "wide open eyes, extra eyes, many eyes, deformed, melting, blurry, text, watermark, signature, "
    "blank white eyes, empty eyes, pupilless eyes, solid black eyes, closed eyes, one eye, sunglasses"
)
SEED = 314
STEPS = 25
CFG = 7.0
SCHEDULER = "EulerAncestralDiscreteScheduler"
WIDTH = 512
HEIGHT = 512


class Tee:
    """Mirror a console stream into the run log."""

    def __init__(self, stream, log):
        self.stream = stream
        self.log = log

    def write(self, data):
        self.stream.write(data)
        self.log.write(data)
        self.log.flush()

    def flush(self):
        self.stream.flush()
        self.log.flush()

    def isatty(self):
        return False


def log(msg):
    print(f"[{time.time() - T_START:8.1f}s] {msg}", flush=True)


def cpu_model():
    try:
        with open("/proc/cpuinfo") as f:
            for line in f:
                if line.startswith("model name"):
                    return line.split(":", 1)[1].strip()
    except OSError:
        pass
    return platform.processor() or platform.machine()


def ram_gb():
    try:
        with open("/proc/meminfo") as f:
            for line in f:
                if line.startswith("MemTotal:"):
                    return round(int(line.split()[1]) / 1024 / 1024, 2)
    except OSError:
        pass
    return None


def sha256_of(path):
    h = hashlib.sha256()
    with open(path, "rb") as f:
        for chunk in iter(lambda: f.read(16 * 1024 * 1024), b""):
            h.update(chunk)
    return h.hexdigest()


def download(url, path):
    tmp = path + ".part"
    log(f"downloading {url} -> {path}")
    req = urllib.request.Request(url, headers={"User-Agent": "sd15-pepe-render/1.0"})
    with urllib.request.urlopen(req) as resp, open(tmp, "wb") as out:
        shutil.copyfileobj(resp, out, 16 * 1024 * 1024)
    os.replace(tmp, path)
    log(f"downloaded {os.path.getsize(path)} bytes")


def fetch_and_verify(name, url, path, expected):
    if not os.path.exists(path):
        download(url, path)
    else:
        log(f"found {path} ({os.path.getsize(path)} bytes), not downloading")
    t = time.time()
    actual = sha256_of(path)
    ok = actual == expected
    log(f"sha256 {name}: expected {expected}")
    log(f"sha256 {name}: actual   {actual}  -> {'OK' if ok else 'MISMATCH'} ({time.time() - t:.1f}s)")
    if not ok:
        raise RuntimeError(f"sha256 mismatch for {name} ({path}): expected {expected}, got {actual}")


def main():
    os.makedirs(OUT_DIR, exist_ok=True)
    os.makedirs(MODELS_DIR, exist_ok=True)
    log_file = open(os.path.join(OUT_DIR, "run.log"), "w", encoding="utf-8")
    sys.stdout = Tee(sys.__stdout__, log_file)
    sys.stderr = Tee(sys.__stderr__, log_file)

    params = {
        "model": BASE_MODEL["name"],
        "model_repo": BASE_MODEL["repo"],
        "model_sha256": BASE_MODEL["sha256"],
        "loras": [{k: l[k] for k in ("name", "url", "sha256", "weight")} for l in LORAS],
        "prompt": PROMPT,
        "negative": NEGATIVE,
        "seed": SEED,
        "steps": STEPS,
        "cfg": CFG,
        "scheduler": SCHEDULER,
        "width": WIDTH,
        "height": HEIGHT,
        "dtype": "float32",
        "device": "cpu",
        "torch_version": None,
        "diffusers_version": None,
        "cpu_model": cpu_model(),
        "cpu_count": os.cpu_count(),
        "ram_gb": ram_gb(),
        "seconds_total": None,
        "seconds_sampling": None,
        "status": "running",
        "stage": "start",
    }

    def write_params():
        params["seconds_total"] = round(time.time() - T_START, 2)
        with open(os.path.join(OUT_DIR, "params.json"), "w") as f:
            json.dump(params, f, indent=2)
            f.write("\n")

    try:
        log(f"python {sys.version.split()[0]} on {platform.platform()}")
        log(f"cpu: {params['cpu_model']} x{params['cpu_count']}, ram: {params['ram_gb']} GB")

        script_copy = os.path.join(OUT_DIR, "render.py")
        if os.path.abspath(__file__) != os.path.abspath(script_copy):
            shutil.copyfile(__file__, script_copy)
        log(f"script sha256 {sha256_of(__file__)}")

        params["stage"] = "download+verify"
        base_path = os.path.join(MODELS_DIR, BASE_MODEL["name"])
        fetch_and_verify(BASE_MODEL["name"], BASE_MODEL["url"], base_path, BASE_MODEL["sha256"])
        for l in LORAS:
            fetch_and_verify(l["name"], l["url"], os.path.join(MODELS_DIR, l["file"]), l["sha256"])

        params["stage"] = "import"
        import torch
        import diffusers
        from diffusers import EulerAncestralDiscreteScheduler, StableDiffusionPipeline
        from PIL.PngImagePlugin import PngInfo

        diffusers.utils.logging.disable_progress_bar()
        params["torch_version"] = torch.__version__
        params["diffusers_version"] = diffusers.__version__
        for mod in ("transformers", "peft", "safetensors", "accelerate"):
            try:
                params[f"{mod}_version"] = __import__(mod).__version__
            except ImportError:
                params[f"{mod}_version"] = None
        log(f"torch {torch.__version__}, diffusers {diffusers.__version__}, "
            f"transformers {params['transformers_version']}, peft {params['peft_version']}, "
            f"threads {torch.get_num_threads()}")

        params["stage"] = "load base model"
        t = time.time()
        # A local copy of the diffusers config for SD 1.5 keeps the run offline; without it,
        # diffusers fetches the same config from the Hugging Face repo.
        config = CONFIG_DIR if os.path.isdir(CONFIG_DIR) else BASE_MODEL["repo"]
        log(f"loading {base_path} with config {config}")
        pipe = StableDiffusionPipeline.from_single_file(
            base_path,
            config=config,
            torch_dtype=torch.float32,
            safety_checker=None,
            feature_extractor=None,
            requires_safety_checker=False,
        )
        pipe.scheduler = EulerAncestralDiscreteScheduler.from_config(pipe.scheduler.config)
        pipe.to("cpu")
        pipe.set_progress_bar_config(disable=True)
        log(f"base model loaded in {time.time() - t:.1f}s")
        log(f"scheduler config: {dict(pipe.scheduler.config)}")

        params["stage"] = "load loras"
        t = time.time()
        for l in LORAS:
            pipe.load_lora_weights(MODELS_DIR, weight_name=l["file"], adapter_name=l["adapter"])
            log(f"loaded LoRA {l['name']!r} as adapter {l['adapter']!r}")
        pipe.set_adapters([l["adapter"] for l in LORAS], adapter_weights=[l["weight"] for l in LORAS])
        log(f"active adapters: {pipe.get_active_adapters()} "
            f"weights {[l['weight'] for l in LORAS]} ({time.time() - t:.1f}s)")

        params["stage"] = "sampling"
        log(f"prompt: {PROMPT}")
        log(f"negative: {NEGATIVE}")
        log(f"seed {SEED}, steps {STEPS}, cfg {CFG}, {SCHEDULER}, {WIDTH}x{HEIGHT}, float32, cpu")
        generator = torch.Generator("cpu").manual_seed(SEED)
        t_step = [time.time()]

        def on_step_end(pipeline, step, timestep, callback_kwargs):
            now = time.time()
            log(f"step {step + 1}/{STEPS} (t={int(timestep)}) {now - t_step[0]:.1f}s")
            t_step[0] = now
            return callback_kwargs

        t_sample = time.time()
        with torch.inference_mode():
            latents = pipe(
                prompt=PROMPT,
                negative_prompt=NEGATIVE,
                width=WIDTH,
                height=HEIGHT,
                num_inference_steps=STEPS,
                guidance_scale=CFG,
                generator=generator,
                output_type="latent",
                callback_on_step_end=on_step_end,
            ).images
            params["seconds_sampling"] = round(time.time() - t_sample, 2)
            log(f"sampling done in {params['seconds_sampling']}s")

            # Same decode the pipeline does for output_type="pil" (no safety checker).
            params["stage"] = "vae decode"
            t = time.time()
            decoded = pipe.vae.decode(latents / pipe.vae.config.scaling_factor, return_dict=False)[0]
            image = pipe.image_processor.postprocess(decoded, output_type="pil")[0]
            log(f"vae decode in {time.time() - t:.1f}s")

        params["stage"] = "save"
        lora_text = "; ".join(f"{l['name']} weight {l['weight']} sha256 {l['sha256']}" for l in LORAS)
        parameters = (
            f"{PROMPT}\n"
            f"Negative prompt: {NEGATIVE}\n"
            f"Steps: {STEPS}, Sampler: Euler a ({SCHEDULER}), CFG scale: {CFG}, Seed: {SEED}, "
            f"Size: {WIDTH}x{HEIGHT}, Model: {BASE_MODEL['name']}, Model hash (sha256): {BASE_MODEL['sha256']}, "
            f"Precision: float32, Device: cpu, LoRAs: {lora_text}"
        )
        info = PngInfo()
        info.add_text("parameters", parameters)
        png_path = os.path.join(OUT_DIR, "pepe.png")
        image.save(png_path, pnginfo=info)
        params["png_sha256"] = sha256_of(png_path)
        log(f"saved {png_path} {image.size[0]}x{image.size[1]} sha256 {params['png_sha256']}")

        params["status"] = "ok"
        params["stage"] = "done"
        write_params()
        log(f"total {params['seconds_total']}s, sampling {params['seconds_sampling']}s")
        return 0
    except BaseException as e:
        params["status"] = "failed"
        params["error"] = f"{type(e).__name__}: {e}"
        traceback.print_exc()
        log(f"STOPPED at stage {params['stage']!r}: {params['error']}")
        write_params()
        return 1


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