import os
import random
import time
from dataclasses import dataclass
from pathlib import Path
from typing import Iterable, Optional

# -----------------------------------------------------------------------------
# Environment must be configured before importing DiffSynth/torch-heavy modules.
# DiffSynth defaults to ModelScope unless this is set; keep it pinned to HF.
# -----------------------------------------------------------------------------
os.environ.setdefault("DIFFSYNTH_DOWNLOAD_SOURCE", "huggingface")
os.environ.setdefault("DIFFSYNTH_SKIP_DOWNLOAD", "True")
os.environ.setdefault("HF_HUB_ENABLE_HF_TRANSFER", "1")
#os.environ.setdefault("PYTORCH_CUDA_ALLOC_CONF", "expandable_segments:True")


def _select_writable_dir(env_name: str, candidates: list[str]) -> str:
    existing = os.getenv(env_name)
    if existing:
        candidates = [existing] + [c for c in candidates if c != existing]
    for candidate in candidates:
        try:
            path = Path(candidate)
            path.mkdir(parents=True, exist_ok=True)
            test_file = path / ".write_test"
            test_file.write_text("ok")
            test_file.unlink(missing_ok=True)
            return str(path)
        except Exception:
            continue
    fallback = Path("/tmp") / env_name.lower()
    fallback.mkdir(parents=True, exist_ok=True)
    return str(fallback)


HF_HOME = _select_writable_dir("HF_HOME", ["/data/.cache/huggingface", "/tmp/.cache/huggingface"])
LOCAL_MODEL_DIR = _select_writable_dir("ANIMA_LOCAL_MODEL_DIR", ["/data/models/anima-v1", "/tmp/models/anima-v1"])
os.environ["HF_HOME"] = HF_HOME
os.environ.setdefault("HF_HUB_CACHE", str(Path(HF_HOME) / "hub"))
os.environ.setdefault("DIFFSYNTH_MODEL_BASE_PATH", str(Path(LOCAL_MODEL_DIR) / "diffsynth"))

# Import spaces before torch/gradio for ZeroGPU compatibility.
try:
    import spaces
except Exception:  # Allows local CPU/GPU testing outside Hugging Face Spaces.
    class _SpacesFallback:
        def GPU(self, *args, **kwargs):
            if args and callable(args[0]):
                return args[0]
            def decorator(fn):
                return fn
            return decorator
    spaces = _SpacesFallback()

import gradio as gr
import torch
from huggingface_hub import hf_hub_download, snapshot_download

from diffsynth.pipelines.anima_image import AnimaImagePipeline, ModelConfig

MODEL_ID = os.getenv("ANIMA_MODEL_ID", "circlestone-labs/Anima")
DIFFUSION_FILE = os.getenv(
    "ANIMA_DIFFUSION_FILE",
    "split_files/diffusion_models/anima-base-v1.0.safetensors",
)
TEXT_ENCODER_FILE = os.getenv(
    "ANIMA_TEXT_ENCODER_FILE",
    "split_files/text_encoders/qwen_3_06b_base.safetensors",
)
VAE_FILE = os.getenv(
    "ANIMA_VAE_FILE",
    "split_files/vae/qwen_image_vae.safetensors",
)
QWEN_TOKENIZER_ID = os.getenv("ANIMA_QWEN_TOKENIZER_ID", "Qwen/Qwen3-0.6B")
T5_TOKENIZER_ID = os.getenv("ANIMA_T5_TOKENIZER_ID", "google/t5-v1_1-xxl")
T5_TOKENIZER_SUBFOLDER = os.getenv("ANIMA_T5_TOKENIZER_SUBFOLDER", "")
T5_TOKENIZER_FALLBACK_ID = os.getenv("ANIMA_T5_TOKENIZER_FALLBACK_ID", "google/t5-v1_1-xxl")
VRAM_LIMIT_GB = os.getenv("ANIMA_VRAM_LIMIT_GB")
LOAD_AT_STARTUP = os.getenv("ANIMA_LOAD_AT_STARTUP", "1").strip().lower() not in {"0", "false", "no"}

DEFAULT_POSITIVE_PREFIX = "masterpiece, best quality, score_7, safe, "
DEFAULT_NEGATIVE = "worst quality, low quality, score_1, score_2, score_3, artist name, watermark, signature, bad anatomy"

# Keep tokenizer downloads small. These patterns intentionally exclude model weights.
TOKENIZER_ALLOW_PATTERNS = [
    "tokenizer.json",
    "tokenizer_config.json",
    "special_tokens_map.json",
    "added_tokens.json",
    "vocab.json",
    "merges.txt",
    "config.json",
    "*.model",
    "*.tiktoken",
    "*.jinja",
]


@dataclass
class AssetPaths:
    diffusion: str
    text_encoder: str
    vae: str
    qwen_tokenizer_dir: str
    t5_tokenizer_dir: str


_PIPE: Optional[AnimaImagePipeline] = None
_ASSETS: Optional[AssetPaths] = None
_STARTUP_STATUS = "Starting up."
_STARTUP_ERROR: Optional[str] = None
_HF_TOKEN = os.getenv("HF_TOKEN") or os.getenv("HUGGINGFACE_HUB_TOKEN")


def _repo_local_dir(repo_id: str, suffix: str = "") -> Path:
    safe = repo_id.replace("/", "--")
    if suffix:
        safe = f"{safe}--{suffix}"
    path = Path(LOCAL_MODEL_DIR) / safe
    path.mkdir(parents=True, exist_ok=True)
    return path


def _download_file_from_hf(repo_id: str, filename: str, local_dir: Path) -> str:
    """Download one Hub file into a predictable local directory and return its path."""
    return hf_hub_download(
        repo_id=repo_id,
        filename=filename,
        local_dir=str(local_dir),
        token=_HF_TOKEN,
    )


def _has_tokenizer_files(path: Path) -> bool:
    """Return True when a directory is usable by AutoTokenizer.from_pretrained."""
    if not path.is_dir():
        return False
    markers = (
        "tokenizer.json",
        "tokenizer.model",
        "spiece.model",
        "vocab.json",
        "merges.txt",
    )
    return any((path / marker).exists() for marker in markers)


def _download_tokenizer_repo(repo_id: str, *, subfolder: str = "", suffix: str = "tokenizer") -> str:
    """Download tokenizer-only files and return the local directory to pass to AutoTokenizer."""
    cleaned_subfolder = subfolder.strip("/")
    local_dir = _repo_local_dir(repo_id, suffix if not cleaned_subfolder else f"{suffix}--{cleaned_subfolder}")
    if cleaned_subfolder:
        allow_patterns = [f"{cleaned_subfolder}/{pattern}" for pattern in TOKENIZER_ALLOW_PATTERNS]
    else:
        allow_patterns = TOKENIZER_ALLOW_PATTERNS

    snapshot_dir = Path(
        snapshot_download(
            repo_id=repo_id,
            allow_patterns=allow_patterns,
            local_dir=str(local_dir),
            token=_HF_TOKEN,
        )
    )

    candidates = []
    if cleaned_subfolder:
        candidates.append(snapshot_dir / cleaned_subfolder)
    candidates.append(snapshot_dir)

    for candidate in candidates:
        if _has_tokenizer_files(candidate):
            return str(candidate)

    searched = ", ".join(str(candidate) for candidate in candidates)
    raise RuntimeError(
        f"Downloaded tokenizer files from {repo_id!r}, but no usable tokenizer files were found. "
        f"Checked: {searched}"
    )


def _download_t5_tokenizer() -> str:
    """Use the public T5 tokenizer first, with a fallback for custom mirrors."""
    try:
        return _download_tokenizer_repo(
            T5_TOKENIZER_ID,
            subfolder=T5_TOKENIZER_SUBFOLDER,
            suffix="t5-tokenizer",
        )
    except Exception as exc:
        if T5_TOKENIZER_ID == T5_TOKENIZER_FALLBACK_ID and not T5_TOKENIZER_SUBFOLDER:
            raise
        print(
            f"[startup] T5 tokenizer download from {T5_TOKENIZER_ID!r} "
            f"subfolder {T5_TOKENIZER_SUBFOLDER!r} failed: {exc}. "
            f"Falling back to {T5_TOKENIZER_FALLBACK_ID!r}.",
            flush=True,
        )
        return _download_tokenizer_repo(
            T5_TOKENIZER_FALLBACK_ID,
            subfolder="",
            suffix="t5-tokenizer-fallback",
        )


def _download_assets() -> AssetPaths:
    """Download all required files from Hugging Face Hub at application startup."""
    print(f"[startup] Downloading Anima assets from Hugging Face Hub into {LOCAL_MODEL_DIR}", flush=True)
    start = time.time()

    anima_dir = _repo_local_dir(MODEL_ID)
    diffusion = _download_file_from_hf(MODEL_ID, DIFFUSION_FILE, anima_dir)
    text_encoder = _download_file_from_hf(MODEL_ID, TEXT_ENCODER_FILE, anima_dir)
    vae = _download_file_from_hf(MODEL_ID, VAE_FILE, anima_dir)

    qwen_dir = _download_tokenizer_repo(QWEN_TOKENIZER_ID, suffix="qwen-tokenizer")
    t5_tokenizer_dir = _download_t5_tokenizer()

    elapsed = time.time() - start
    print(f"[startup] Hugging Face downloads ready in {elapsed:.1f}s", flush=True)
    return AssetPaths(
        diffusion=diffusion,
        text_encoder=text_encoder,
        vae=vae,
        qwen_tokenizer_dir=str(qwen_dir),
        t5_tokenizer_dir=t5_tokenizer_dir,
    )


def _ensure_assets() -> AssetPaths:
    global _ASSETS
    if _ASSETS is None:
        _ASSETS = _download_assets()
    return _ASSETS


def _set_progress(progress: Optional[gr.Progress], value=0, desc: str = "") -> None:
    if progress is None:
        return
    try:
        progress(value, desc=desc)
    except TypeError:
        try:
            progress(value)
        except Exception:
            pass


def _progress_wrapper(progress: gr.Progress):
    def _wrap(iterable: Iterable):
        try:
            total = len(iterable)
        except Exception:
            total = None
        for index, item in enumerate(iterable):
            if total:
                _set_progress(progress, (index, total), desc=f"Denoising {index + 1}/{total}")
            else:
                _set_progress(progress, 0, desc=f"Denoising step {index + 1}")
            yield item
        if total:
            _set_progress(progress, (total, total), desc="Decoding")
        else:
            _set_progress(progress, 0, desc="Decoding")
    return _wrap


def _normalize_dimension(value: int) -> int:
    value = int(value)
    value = max(256, min(1536, value))
    return int(round(value / 16) * 16)


def _normalize_seed(seed) -> int:
    try:
        seed = int(seed)
    except Exception:
        seed = -1
    if seed < 0:
        return random.randint(0, 2**31 - 1)
    return seed


def _prepare_prompt(prompt: str, add_prefix: bool) -> str:
    prompt = (prompt or "").strip()
    if not prompt:
        prompt = "1girl, solo, long silver hair, blue eyes, blue dress, underwater, floating hair, refraction, portrait"
    if add_prefix and not prompt.lower().startswith(DEFAULT_POSITIVE_PREFIX.strip().lower()):
        prompt = DEFAULT_POSITIVE_PREFIX + prompt
    return prompt


def _load_pipe(progress: Optional[gr.Progress] = None) -> AnimaImagePipeline:
    """Create the pipeline once, using local files that were downloaded from HF."""
    global _PIPE
    if _PIPE is not None:
        return _PIPE

    assets = _ensure_assets()

    _set_progress(progress, 0, desc="Loading Anima V1 from local files")
    torch.set_float32_matmul_precision("high")
    vram_limit = float(VRAM_LIMIT_GB) if VRAM_LIMIT_GB else None

    # Pass local paths so DiffSynth never tries to resolve/download from ModelScope.
    _PIPE = AnimaImagePipeline.from_pretrained(
        torch_dtype=torch.float32,
        device="cpu",
        model_configs=[
            ModelConfig(path=assets.diffusion),
            ModelConfig(path=assets.text_encoder),
            ModelConfig(path=assets.vae),
        ],
        tokenizer_config=ModelConfig(path=assets.qwen_tokenizer_dir),
        tokenizer_t5xxl_config=ModelConfig(path=assets.t5_tokenizer_dir),
        vram_limit=vram_limit,
    )
    print("[startup] Anima pipeline loaded", flush=True)
    return _PIPE


def _startup_download_and_load() -> str:
    """Best-effort eager startup. Download failures are shown in the UI; load failures retry on Generate."""
    global _ASSETS, _STARTUP_ERROR
    messages = []
    try:
        _ASSETS = _download_assets()
        messages.append("Downloaded required files from Hugging Face Hub at startup.")
    except Exception as exc:
        _STARTUP_ERROR = f"Startup download failed: {exc}"
        print(f"[startup] {_STARTUP_ERROR}", flush=True)
        return f"⚠️ {_STARTUP_ERROR} Generate will retry."

    if LOAD_AT_STARTUP:
        try:
            _load_pipe(progress=None)
            messages.append("Loaded the Anima pipeline at startup.")
        except Exception as exc:
            _STARTUP_ERROR = f"Startup GPU load was deferred/failed: {exc}"
            print(f"[startup] {_STARTUP_ERROR}", flush=True)
            messages.append("Startup load failed; Generate will retry only as a fallback.")
    else:
        messages.append("Startup model load disabled by ANIMA_LOAD_AT_STARTUP=0; Generate will load once.")

    return " ".join(messages)


_STARTUP_STATUS = _startup_download_and_load()


@spaces.GPU(duration=30)
def generate(
    prompt: str,
    negative_prompt: str,
    width: int,
    height: int,
    steps: int,
    cfg_scale: float,
    sigma_shift: float,
    seed: int,
    add_recommended_prefix: bool,
    progress: gr.Progress = gr.Progress(track_tqdm=False),
):
    prompt = _prepare_prompt(prompt, add_recommended_prefix)
    negative_prompt = (negative_prompt or DEFAULT_NEGATIVE).strip()
    width = _normalize_dimension(width)
    height = _normalize_dimension(height)
    steps = int(max(3, min(60, steps)))
    cfg_scale = float(max(1.0, min(8.0, cfg_scale)))
    sigma_shift_value = None if sigma_shift <= 0 else float(sigma_shift)
    seed = _normalize_seed(seed)

    pipe = _load_pipe(progress)

    _set_progress(progress, 0, desc="Generating")
    with torch.inference_mode():
        image = pipe(
            prompt=prompt,
            negative_prompt=negative_prompt,
            cfg_scale=cfg_scale,
            height=height,
            width=width,
            seed=seed,
            num_inference_steps=steps,
            sigma_shift=sigma_shift_value,
            progress_bar_cmd=_progress_wrapper(progress),
        )

    info = (
        f"**Seed:** `{seed}`  \n"
        f"**Size:** `{width}×{height}`  \n"
        f"**Steps / CFG:** `{steps}` / `{cfg_scale}`  \n"
        f"**Model file:** `{MODEL_ID}:{DIFFUSION_FILE}`  \n"
        f"**Startup:** {_STARTUP_STATUS}"
    )
    return image, info


CSS = """
#title {text-align: center;}
#info-box {font-size: 0.95rem;}
#startup-box {font-size: 0.9rem; opacity: 0.9;}
"""

with gr.Blocks(css=CSS, title="Anima V1 ZeroGPU Demo") as demo:
    gr.Markdown(
        "# Anima V1 CPU Demo\n"
        "This demo loads the `circlestone-labs/Anima` V1 single-file weights: "
        "`anima-base-v1.0.safetensors`, `qwen_3_06b_base.safetensors`, and `qwen_image_vae.safetensors`. Tokenizers are downloaded separately from public Hugging Face repos."
    )
    gr.Markdown(f"**Startup status:** {_STARTUP_STATUS}", elem_id="startup-box")
    with gr.Row():
        with gr.Column(scale=1):
            prompt = gr.Textbox(
                label="Prompt",
                lines=5,
                value=(
                    "1girl, solo, long silver hair, blue eyes, blue dress, underwater, "
                    "air bubbles, floating hair, refraction, portrait, looking at viewer"
                ),
            )
            negative_prompt = gr.Textbox(
                label="Negative prompt",
                lines=3,
                value=DEFAULT_NEGATIVE,
            )
            add_recommended_prefix = gr.Checkbox(
                label=f"Prepend recommended prefix: {DEFAULT_POSITIVE_PREFIX.strip()}",
                value=True,
            )
            with gr.Row():
                width = gr.Slider(256, 1536, value=256, step=16, label="Width")
                height = gr.Slider(256, 1536, value=256, step=16, label="Height")
            with gr.Row():
                steps = gr.Slider(3, 60, value=7, step=1, label="Steps")
                cfg_scale = gr.Slider(1.0, 8.0, value=4.5, step=0.1, label="CFG scale")
            with gr.Row():
                sigma_shift = gr.Slider(0.0, 8.0, value=0.0, step=0.1, label="Sigma shift (0 = DiffSynth default)")
                seed = gr.Number(value=-1, precision=0, label="Seed (-1 random)")
            generate_btn = gr.Button("Generate", variant="primary")
        with gr.Column(scale=1):
            image = gr.Image(label="Output", type="pil")
            info = gr.Markdown(elem_id="info-box")

    inputs = [
        prompt,
        negative_prompt,
        width,
        height,
        steps,
        cfg_scale,
        sigma_shift,
        seed,
        add_recommended_prefix,
    ]
    generate_btn.click(generate, inputs=inputs, outputs=[image, info], show_progress=True)
    gr.Examples(
        examples=[
            [
                "1girl, solo, long silver hair, blue eyes, blue dress, underwater, air bubbles, floating hair, refraction, portrait, looking at viewer",
                DEFAULT_NEGATIVE,
                1024,
                1024,
                35,
                4.5,
                0.0,
                0,
                True,
            ],
            [
                "year 2025, newest, highres, safe, 1girl, witch hat, black dress, glowing runes, moonlit forest, dynamic pose, dramatic lighting",
                DEFAULT_NEGATIVE,
                896,
                1152,
                40,
                4.5,
                0.0,
                12345,
                False,
            ],
            [
                "safe, digital painting of a small dragon sleeping on a stack of books in a cozy candlelit library, painterly, warm light, highly detailed background",
                DEFAULT_NEGATIVE,
                1216,
                832,
                35,
                4.0,
                0.0,
                -1,
                True,
            ],
        ],
        inputs=inputs,
        outputs=[image, info],
        fn=generate,
        cache_examples=False,
    )

if __name__ == "__main__":
    demo.queue(max_size=20).launch(ssr_mode=False)