import hashlib
import json
import os
import shutil
import signal
import subprocess
import time
from pathlib import Path
from huggingface_hub import hf_hub_download

CONFIG = json.loads(Path(__file__).with_name("models.json").read_text())
BINARY = os.getenv("SD_CLI", "/usr/local/bin/sd-cli")


def download(spec):
    return hf_hub_download(repo_id=spec["repo"], filename=spec["file"],
                           revision=spec.get("revision", "main"),
                           token=os.getenv("HF_TOKEN") or None)


def diagnostics():
    pieces = ["CPU-only backend; model weights load on first request."]
    commit = Path("/engine-commit.txt")
    if commit.exists():
        pieces.append("Engine commit: " + commit.read_text().strip())
    for name in ("/sys/fs/cgroup/memory.max", "/sys/fs/cgroup/memory/memory.limit_in_bytes"):
        p = Path(name)
        if p.exists():
            pieces.append("Container memory limit (bytes): " + p.read_text().strip())
            break
    if shutil.which(BINARY):
        result = subprocess.run([BINARY, "--help"], capture_output=True, text=True, timeout=30)
        pieces.append((result.stdout + result.stderr)[-18000:])
    else:
        pieces.append("sd-cli missing: build the Docker image before running inference.")
    return "\n".join(pieces)


def run_one(image, output, prompt, negative, seed, steps, cfg, width, height,
            adapters, jobdir, status):
    help_result = subprocess.run([BINARY, "--help"], capture_output=True, text=True, timeout=30)
    help_text = help_result.stdout + help_result.stderr
    models = {}
    for name, spec in CONFIG.items():
        status(f"Downloading/checking {name}: {spec['file']}")
        models[name] = download(spec)
    lora_dir = jobdir / "loras"
    lora_dir.mkdir(exist_ok=True)
    tags = []
    for index, (spec, strength) in enumerate(adapters):
        path = Path(download(spec))
        name = f"adapter_{index}"
        target = lora_dir / (name + path.suffix)
        if not target.exists():
            target.symlink_to(path.resolve())
        tags.append(f"<lora:{name}:{strength:.4f}>")
    command = [BINARY, "--diffusion-model", models["diffusion"], "--vae", models["vae"],
               "--llm", models["llm"], "--llm_vision", models["vision"],
               "-r", str(image), "-p", prompt + " " + " ".join(tags),
               "-n", negative, "--seed", str(seed), "--steps", str(steps),
               "--cfg-scale", str(cfg), "--sampling-method", "euler", "--flow-shift", "3",
               "-W", str(width), "-H", str(height), "--threads", os.getenv("CPU_THREADS", "2"),
               "-o", str(output)]
    for flag in ("--mmap", "--vae-tiling", "--diffusion-fa"):
        if flag in help_text:
            command.append(flag)
    if tags:
        if "--lora-model-dir" not in help_text:
            raise RuntimeError("This engine version does not expose --lora-model-dir")
        command += ["--lora-model-dir", str(lora_dir)]
    log = output.with_suffix(".log")
    timeout = int(os.getenv("INFERENCE_TIMEOUT", "21600"))
    status("CPU inference running. This can take a long time; see container logs.")
    with log.open("w") as stream:
        process = subprocess.Popen(command, stdout=stream, stderr=subprocess.STDOUT, start_new_session=True)
        start = time.monotonic()
        try:
            while process.poll() is None:
                if time.monotonic() - start > timeout:
                    raise TimeoutError(f"CPU inference exceeded {timeout} seconds")
                time.sleep(2)
            if process.returncode or not output.exists():
                tail = log.read_text(errors="replace")[-7000:]
                if process.returncode == -9:
                    tail = "Process killed (possible container RAM exhaustion). " + tail
                raise RuntimeError(f"sd-cli exit={process.returncode}\n{tail}")
        finally:
            if process.poll() is None:
                os.killpg(process.pid, signal.SIGTERM)
                try:
                    process.wait(timeout=10)
                except subprocess.TimeoutExpired:
                    os.killpg(process.pid, signal.SIGKILL)
                    process.wait()
    return log