#!/usr/bin/env python3
"""
Standalone Offline CPU Image Generator for ModelsLab/indigo-furry-mix-v65.

Supports:
- 100% offline local loading from ./hf-model (and ./hf-model/lcm-lora for --lcm)
- Standard schedulers (Euler, Euler a, DDIM, DPM++ 2M, DPM++ 2M Karras, PNDM, Heun, LMS, LCM)
- Optional LCM-LoRA ultra-fast 4-8 step generation on CPU (--lcm)
- Single-run CLI mode (--prompt "...") or interactive loop mode (--interactive or omit --prompt)
"""

import argparse
import random
import shlex
from datetime import datetime
from pathlib import Path

import torch
from diffusers import (
    DDIMScheduler,
    DPMSolverMultistepScheduler,
    DPMSolverSinglestepScheduler,
    EulerAncestralDiscreteScheduler,
    EulerDiscreteScheduler,
    HeunDiscreteScheduler,
    LCMScheduler,
    LMSDiscreteScheduler,
    PNDMScheduler,
    StableDiffusionPipeline,
)

SCRIPT_DIR = Path(__file__).resolve().parent
DEFAULT_MODEL_PATH = SCRIPT_DIR / "hf-model"
DEFAULT_LCM_LORA_PATH = DEFAULT_MODEL_PATH / "lcm-lora"
HF_BASE_MODEL_ID = "ModelsLab/indigo-furry-mix-v65"
HF_LCM_LORA_ID = "latent-consistency/lcm-lora-sdv1-5"

DEFAULT_NEGATIVE_PROMPT = (
    "lowres, bad anatomy, bad hands, text, error, missing fingers, "
    "extra digit, fewer digits, cropped, worst quality, low quality, "
    "normal quality, jpeg artifacts, signature, watermark, username, blurry"
)


def configure_scheduler(pipe: StableDiffusionPipeline, scheduler_name: str) -> None:
    """Replaces the pipeline's scheduler based on the user's CLI choice."""
    name = scheduler_name.strip().lower().replace("-", "_").replace(" ", "_")
    config = pipe.scheduler.config

    if name in ("ddim",):
        pipe.scheduler = DDIMScheduler.from_config(config)
    elif name in ("euler",):
        pipe.scheduler = EulerDiscreteScheduler.from_config(config)
    elif name in ("euler_a", "euler_ancestral"):
        pipe.scheduler = EulerAncestralDiscreteScheduler.from_config(config)
    elif name in ("dpm++", "dpm++_2m", "dpmpp_2m"):
        pipe.scheduler = DPMSolverMultistepScheduler.from_config(config)
    elif name in ("dpm++_karras", "dpm++_2m_karras", "dpmpp_2m_karras", "karras"):
        pipe.scheduler = DPMSolverMultistepScheduler.from_config(
            config, use_karras_sigmas=True
        )
    elif name in ("dpm++_sde", "dpmpp_sde"):
        pipe.scheduler = DPMSolverSinglestepScheduler.from_config(config)
    elif name in ("dpm++_sde_karras", "dpmpp_sde_karras"):
        pipe.scheduler = DPMSolverSinglestepScheduler.from_config(
            config, use_karras_sigmas=True
        )
    elif name in ("pndm",):
        pipe.scheduler = PNDMScheduler.from_config(config)
    elif name in ("heun",):
        pipe.scheduler = HeunDiscreteScheduler.from_config(config)
    elif name in ("lms",):
        pipe.scheduler = LMSDiscreteScheduler.from_config(config)
    elif name in ("lcm",):
        pipe.scheduler = LCMScheduler.from_config(config)
    else:
        supported = (
            "ddim, euler, euler_a, dpm++, dpm++_karras, "
            "dpm++_sde, dpm++_sde_karras, pndm, heun, lms, lcm"
        )
        raise ValueError(
            f"Unknown scheduler '{scheduler_name}'. Supported values: {supported}"
        )


def resolve_seed(seed_arg: str) -> int:
    """Parses '--seed r' (or 'random' / '-1') into a random 32-bit integer, or converts numeric string."""
    val = str(seed_arg).strip().lower()
    if val in ("r", "random", "-1", "none", ""):
        return random.randint(0, 2**32 - 1)
    try:
        return int(val)
    except ValueError as err:
        raise ValueError(
            f"Invalid seed '{seed_arg}'. Use an integer (e.g. 12345) or 'r' for random."
        ) from err


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(
        description="Generate images offline on CPU using local hf-model",
        formatter_class=argparse.ArgumentDefaultsHelpFormatter,
    )
    parser.add_argument(
        "-p",
        "--prompt",
        type=str,
        default=None,
        help="Positive text prompt. If omitted, starts in interactive mode.",
    )
    parser.add_argument(
        "-i",
        "--interactive",
        action="store_true",
        help="Keep the model loaded in memory and enter an interactive prompt loop.",
    )
    parser.add_argument(
        "-n",
        "--negative-prompt",
        type=str,
        default=DEFAULT_NEGATIVE_PROMPT,
        help="Negative prompt (things to avoid in the image).",
    )
    parser.add_argument(
        "--steps",
        type=int,
        default=None,
        help="Number of inference steps (defaults to 4 when --lcm is used, otherwise 20).",
    )
    parser.add_argument(
        "--cfg",
        type=float,
        default=None,
        help="Guidance scale / CFG (defaults to 1.5 when --lcm is used, otherwise 7.0).",
    )
    parser.add_argument(
        "--seed",
        type=str,
        default="r",
        help="Seed number for reproducibility, or 'r' for a random seed.",
    )
    parser.add_argument(
        "--scheduler",
        type=str,
        default=None,
        help=(
            "Sampler/scheduler to use: ddim, euler, euler_a, dpm++, dpm++_karras, "
            "dpm++_sde, pndm, heun, lms, lcm (defaults to 'lcm' with --lcm, else 'dpm++_karras')."
        ),
    )
    parser.add_argument(
        "--lcm",
        action="store_true",
        help="Enable LCM-LoRA mode for ultra-fast 4-8 step generation on CPU.",
    )
    parser.add_argument(
        "--model-path",
        type=str,
        default=str(DEFAULT_MODEL_PATH),
        help="Path to the local model directory.",
    )
    parser.add_argument(
        "--lora-path",
        type=str,
        default=str(DEFAULT_LCM_LORA_PATH),
        help="Path to the local LCM-LoRA directory.",
    )
    parser.add_argument(
        "--width",
        type=int,
        default=512,
        help="Image width in pixels (must be a multiple of 8).",
    )
    parser.add_argument(
        "--height",
        type=int,
        default=512,
        help="Image height in pixels (must be a multiple of 8).",
    )
    parser.add_argument(
        "--count",
        "--num-images",
        "-c",
        dest="num_images",
        type=int,
        default=1,
        help="Number of images to generate per prompt.",
    )
    parser.add_argument(
        "-o",
        "--output",
        type=str,
        default=None,
        help="Custom output file path. If omitted, saves to outputs/<timestamp>_seed<seed>.png.",
    )
    parser.add_argument(
        "--threads",
        type=int,
        default=None,
        help="Optional number of CPU threads for PyTorch to use.",
    )
    return parser.parse_args()


def generate_one(
    pipe: StableDiffusionPipeline,
    prompt: str,
    negative_prompt: str,
    steps: int,
    cfg: float,
    seed_arg: str,
    scheduler_name: str,
    lcm_active: bool,
    width: int,
    height: int,
    num_images: int = 1,
    output_override: str | None = None,
) -> list[Path]:
    """Runs image generation (1 or more images) and saves the results to disk."""
    configure_scheduler(pipe, scheduler_name)
    base_seed = resolve_seed(seed_arg)
    seeds = [base_seed + i for i in range(num_images)]

    print("\n--- Generation Config ---")
    print(f"Prompt          : {prompt}")
    print(f"Negative Prompt : {negative_prompt}")
    print(f"Scheduler       : {scheduler_name} ({pipe.scheduler.__class__.__name__})")
    print(f"LCM-LoRA Active : {lcm_active}")
    print(f"Steps           : {steps}")
    print(f"CFG Scale       : {cfg}")
    print(f"Images (Count)  : {num_images}")
    print(f"Seed(s)         : {seeds if num_images > 1 else base_seed}")
    print(f"Resolution      : {width}x{height}")
    print("-------------------------\n")

    generators = [
        torch.Generator(device="cpu").manual_seed(s) for s in seeds
    ]

    result = pipe(
        prompt=prompt,
        negative_prompt=negative_prompt,
        num_inference_steps=steps,
        guidance_scale=cfg,
        width=width,
        height=height,
        num_images_per_prompt=num_images,
        generator=generators,
    )

    saved_paths: list[Path] = []
    timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")

    for idx, (image, s) in enumerate(zip(result.images, seeds), start=1):
        if output_override:
            orig = Path(output_override)
            orig.parent.mkdir(parents=True, exist_ok=True)
            if num_images == 1:
                output_path = orig
            else:
                output_path = orig.with_name(f"{orig.stem}_{idx}{orig.suffix}")
        else:
            output_dir = SCRIPT_DIR / "outputs"
            output_dir.mkdir(parents=True, exist_ok=True)
            if num_images == 1:
                output_path = output_dir / f"{timestamp}_seed{s}.png"
            else:
                output_path = output_dir / f"{timestamp}_{idx}_seed{s}.png"

        image.save(output_path)
        saved_paths.append(output_path)
        print(f"[+] Saved image ({idx}/{num_images}) to: {output_path.resolve()}")

    return saved_paths


def parse_inline_overrides(
    line: str,
    default_steps: int,
    default_cfg: float,
    default_seed: str,
    default_scheduler: str,
    default_width: int,
    default_height: int,
    default_num_images: int,
) -> tuple[str, int, float, str, str, int, int, int]:
    """
    Allows interactive prompts to optionally include inline CLI flags like:
    `a cool wolf --steps 6 --cfg 1.8 --seed 42 --scheduler euler --count 2`
    """
    if " --" not in line and " -c " not in line:
        return (
            line.strip(),
            default_steps,
            default_cfg,
            default_seed,
            default_scheduler,
            default_width,
            default_height,
            default_num_images,
        )

    tokens = shlex.split(line)
    prompt_parts = []
    steps = default_steps
    cfg = default_cfg
    seed = default_seed
    scheduler = default_scheduler
    width = default_width
    height = default_height
    num_images = default_num_images

    i = 0
    while i < len(tokens):
        tok = tokens[i]
        if tok == "--steps" and i + 1 < len(tokens):
            steps = int(tokens[i + 1])
            i += 2
        elif tok == "--cfg" and i + 1 < len(tokens):
            cfg = float(tokens[i + 1])
            i += 2
        elif tok == "--seed" and i + 1 < len(tokens):
            seed = tokens[i + 1]
            i += 2
        elif tok == "--scheduler" and i + 1 < len(tokens):
            scheduler = tokens[i + 1]
            i += 2
        elif tok == "--width" and i + 1 < len(tokens):
            width = int(tokens[i + 1])
            i += 2
        elif tok == "--height" and i + 1 < len(tokens):
            height = int(tokens[i + 1])
            i += 2
        elif tok in ("--count", "--num-images", "-c") and i + 1 < len(tokens):
            num_images = int(tokens[i + 1])
            i += 2
        else:
            prompt_parts.append(tok)
            i += 1

    return (
        " ".join(prompt_parts).strip(),
        steps,
        cfg,
        seed,
        scheduler,
        width,
        height,
        num_images,
    )


def main() -> None:
    args = parse_args()

    if args.threads is not None and args.threads > 0:
        torch.set_num_threads(args.threads)

    steps = args.steps if args.steps is not None else (4 if args.lcm else 20)
    cfg = args.cfg if args.cfg is not None else (1.5 if args.lcm else 7.0)
    scheduler_name = (
        args.scheduler
        if args.scheduler is not None
        else ("lcm" if args.lcm else "dpm++_karras")
    )

    model_path = Path(args.model_path).resolve()
    unet_file = model_path / "unet" / "diffusion_pytorch_model.fp16.safetensors"
    if model_path.exists() and unet_file.exists():
        model_source = str(model_path)
        local_model_only = True
        print(f"[*] Loading local model from '{model_path}' onto CPU (offline mode)...")
    else:
        model_source = HF_BASE_MODEL_ID
        local_model_only = False
        print(f"[*] Local model weights not found at '{model_path}'. Loading '{model_source}' from Hugging Face Hub...")

    # Load using variant="fp16" (matching *.fp16.safetensors)
    # and upcast in-memory to torch.float32 for CPU compatibility.
    pipe = StableDiffusionPipeline.from_pretrained(
        model_source,
        torch_dtype=torch.float32,
        variant="fp16",
        use_safetensors=True,
        safety_checker=None,
        requires_safety_checker=False,
        local_files_only=local_model_only,
    )

    if args.lcm:
        lora_path = Path(args.lora_path).resolve()
        lora_file = lora_path / "pytorch_lora_weights.safetensors"
        if lora_file.exists():
            lora_source = str(lora_path)
            local_lora_only = True
            print(f"[*] Loading and fusing local LCM-LoRA from '{lora_path}'...")
        else:
            lora_source = HF_LCM_LORA_ID
            local_lora_only = False
            print(f"[*] Local LCM-LoRA not found at '{lora_file}'. Loading '{lora_source}' from Hugging Face Hub...")
        pipe.load_lora_weights(
            lora_source,
            weight_name="pytorch_lora_weights.safetensors",
            local_files_only=local_lora_only,
        )
        pipe.fuse_lora()

    # CPU memory & speed optimizations
    pipe.to("cpu")
    pipe.enable_attention_slicing()
    if hasattr(pipe, "enable_vae_slicing"):
        pipe.enable_vae_slicing()

    # Run initial prompt if provided on command line
    if args.prompt:
        generate_one(
            pipe=pipe,
            prompt=args.prompt,
            negative_prompt=args.negative_prompt,
            steps=steps,
            cfg=cfg,
            seed_arg=args.seed,
            scheduler_name=scheduler_name,
            lcm_active=args.lcm,
            width=args.width,
            height=args.height,
            num_images=args.num_images,
            output_override=args.output,
        )

    # Enter interactive loop if --interactive was passed OR no --prompt was given
    if args.interactive or not args.prompt:
        print("\n========================================================")
        print(" Interactive Mode Ready! (Model stays loaded in RAM)")
        print(" Type a prompt and press Enter to generate.")
        print(" Tip: You can also append flags inline, e.g.:")
        print("   cyberpunk wolf --steps 4 --cfg 1.5 --seed 42 --count 2")
        print(" Type 'q' or 'exit' to quit.")
        print("========================================================\n")

        while True:
            try:
                user_input = input("Prompt > ").strip()
            except (EOFError, KeyboardInterrupt):
                print("\nExiting interactive mode.")
                break

            if not user_input:
                continue
            if user_input.lower() in ("q", "quit", "exit"):
                print("Exiting interactive mode.")
                break

            try:
                (
                    p_text,
                    p_steps,
                    p_cfg,
                    p_seed,
                    p_sched,
                    p_width,
                    p_height,
                    p_count,
                ) = parse_inline_overrides(
                    user_input,
                    default_steps=steps,
                    default_cfg=cfg,
                    default_seed=args.seed,
                    default_scheduler=scheduler_name,
                    default_width=args.width,
                    default_height=args.height,
                    default_num_images=args.num_images,
                )
                if not p_text:
                    print("[!] Please provide a non-empty prompt.")
                    continue

                generate_one(
                    pipe=pipe,
                    prompt=p_text,
                    negative_prompt=args.negative_prompt,
                    steps=p_steps,
                    cfg=p_cfg,
                    seed_arg=p_seed,
                    scheduler_name=p_sched,
                    lcm_active=args.lcm,
                    width=p_width,
                    height=p_height,
                    num_images=p_count,
                    output_override=None,
                )
            except Exception as err:
                print(f"[!] Error during generation: {err}")


if __name__ == "__main__":
    main()