"""Checkpoint catalogue and cached downloads for the ZeroGPU Space."""

import hashlib
import os
import tempfile
from pathlib import Path

import requests


SDXL_BASE = "SDXL Base 1.0"
RAGNAROK = "Juggernaut XL - Ragnarok"
MODEL_ID = "stabilityai/stable-diffusion-xl-base-1.0"
RAGNAROK_FILENAME = "juggernautXL_ragnarokBy.safetensors"
RAGNAROK_URL = "https://civitai.com/api/download/models/1759168"
RAGNAROK_SHA256 = "dd08fa32f98d05a2443ca1419e46df1575a0811f6e3b246d9dd47ff20f5eb66a"
RAGNAROK_SIZE = 7105350162
MODEL_CHOICES = (RAGNAROK,)


def model_cache_dir():
    """Use an attached persistent volume when available, otherwise local disk."""
    configured = os.getenv("RF_MODEL_CACHE_DIR")
    if configured:
        return Path(configured)
    if Path("/data").is_dir():
        return Path("/data/ruinedfooocus/checkpoints")
    return Path(__file__).resolve().parent / "model_cache" / "checkpoints"


def file_sha256(path):
    digest = hashlib.sha256()
    with Path(path).open("rb") as stream:
        for chunk in iter(lambda: stream.read(4 * 1024 * 1024), b""):
            digest.update(chunk)
    return digest.hexdigest()


_verified_paths = set()


def download_verified_file(target, url, expected_sha256, progress=None, cached_size=None):
    """Cache one remote file, accepting it only when its checksum matches."""
    target = Path(target)
    target.parent.mkdir(parents=True, exist_ok=True)
    if target.is_file():
        if (str(target) in _verified_paths
                or (cached_size is not None and target.stat().st_size == cached_size)
                or file_sha256(target) == expected_sha256):
            _verified_paths.add(str(target))
            return str(target)
        target.unlink()

    partial = target.with_suffix(target.suffix + ".part")
    headers = {}
    if os.getenv("CIVITAI_TOKEN"):
        headers["Authorization"] = f"Bearer {os.environ['CIVITAI_TOKEN']}"

    try:
        with requests.get(url, headers=headers, stream=True, timeout=(15, 120)) as response:
            response.raise_for_status()
            total = int(response.headers.get("content-length", 0))
            digest = hashlib.sha256()
            downloaded = 0
            with partial.open("wb") as stream:
                for chunk in response.iter_content(chunk_size=4 * 1024 * 1024):
                    if not chunk:
                        continue
                    stream.write(chunk)
                    digest.update(chunk)
                    downloaded += len(chunk)
                    if progress and total:
                        progress(min(downloaded / total, 1), desc=f"Downloading {target.name}")
        if digest.hexdigest() != expected_sha256:
            raise ValueError(f"{target.name} download failed SHA-256 verification")
        partial.replace(target)
        _verified_paths.add(str(target))
        return str(target)
    finally:
        partial.unlink(missing_ok=True)


def stage_checkpoint(source, target, expected_sha256):
    """Copy a bucket checkpoint to local disk and verify it in one sequential read."""
    source = Path(source)
    target = Path(target)
    target.parent.mkdir(parents=True, exist_ok=True)
    if target.is_file() and target.stat().st_size == source.stat().st_size:
        if file_sha256(target) == expected_sha256:
            return str(target)
    partial = target.with_suffix(target.suffix + ".part")
    digest = hashlib.sha256()
    try:
        with source.open("rb") as input_file, partial.open("wb") as output_file:
            for chunk in iter(lambda: input_file.read(16 * 1024 * 1024), b""):
                output_file.write(chunk)
                digest.update(chunk)
        if digest.hexdigest() != expected_sha256:
            raise ValueError(f"{source.name} cache failed SHA-256 verification")
        partial.replace(target)
        return str(target)
    finally:
        partial.unlink(missing_ok=True)


def prepare_model(model_name, progress=None):
    """Return a local checkpoint path, downloading Ragnarok only if needed."""
    if model_name == SDXL_BASE:
        return ""
    if model_name != RAGNAROK:
        raise ValueError(f"Unknown model: {model_name}")
    cached_path = download_verified_file(
        model_cache_dir() / RAGNAROK_FILENAME,
        RAGNAROK_URL,
        RAGNAROK_SHA256,
        progress=progress,
        cached_size=RAGNAROK_SIZE if Path("/data").is_dir() else None,
    )
    if Path("/data").is_dir() and not os.getenv("RF_MODEL_CACHE_DIR"):
        local_path = Path(tempfile.gettempdir()) / "ruinedfooocus" / RAGNAROK_FILENAME
        try:
            return stage_checkpoint(cached_path, local_path, RAGNAROK_SHA256)
        except ValueError:
            Path(cached_path).unlink()
            _verified_paths.discard(cached_path)
            cached_path = download_verified_file(
                model_cache_dir() / RAGNAROK_FILENAME,
                RAGNAROK_URL,
                RAGNAROK_SHA256,
                progress=progress,
            )
            return stage_checkpoint(cached_path, local_path, RAGNAROK_SHA256)
    return cached_path