"""Command line interface for standalone diffusion generation."""

from __future__ import annotations

import argparse
import gc
from pathlib import Path
import sys

from diffusion_cli.checkpoint import inspectSourceTorchDtype
from diffusion_cli.api_profiles import API_PROFILES
from diffusion_cli.config import (
    UserConfig,
    buildGenerationConfig,
    loadUserConfig,
)
from diffusion_cli.errors import DiffusionCliError
from diffusion_cli.image_io import saveImages
from diffusion_cli.model_inspect import formatSummary, inspectModelSource
from diffusion_cli.paths import resolveModelSources
from diffusion_cli.sampling import sampleLatents
from diffusion_cli.server import serve, validateServerConfig
from diffusion_cli.text_encoder import ZImageTextEncoder
from diffusion_cli.vae import ZImageVae
from diffusion_cli.zimage_model import ZImageModel


def buildParser() -> argparse.ArgumentParser:
    """Build the argparse command line parser."""

    parser = argparse.ArgumentParser(
        prog="diffusion-cli",
        description="Standalone local Z-Image Turbo CLI.",
    )
    parser.add_argument("--prompt", help="Positive prompt.")
    parser.add_argument(
        "--negative-prompt",
        help="Negative prompt.",
    )
    parser.add_argument("--seed", type=int, help="Noise seed.")
    parser.add_argument("--width", type=int)
    parser.add_argument("--height", type=int)
    parser.add_argument("--batch-size", type=int)
    parser.add_argument("--steps", type=int)
    parser.add_argument("--cfg", type=float)
    parser.add_argument("--output", type=Path)
    parser.add_argument("--output-extension")
    parser.add_argument("--output-quality", type=int)
    parser.add_argument("--device")
    parser.add_argument(
        "--checkpoint",
        type=Path,
        help="Local all-in-one safetensors checkpoint file.",
    )
    parser.add_argument(
        "--diffusion-model",
        type=Path,
        help="Local Z-Image diffusion model safetensors file.",
    )
    parser.add_argument(
        "--text-encoder",
        type=Path,
        help="Local Qwen text encoder safetensors file.",
    )
    parser.add_argument(
        "--vae",
        type=Path,
        help="Local VAE safetensors file.",
    )
    parser.add_argument(
        "--tokenizer-path",
        type=Path,
        help="Local Qwen tokenizer directory.",
    )
    parser.add_argument(
        "--dtype",
        choices=("auto", "bf16", "fp16", "fp32"),
    )
    parser.add_argument(
        "--inspect-models",
        action="store_true",
        help="Inspect local safetensors metadata and exit.",
    )
    subparsers = parser.add_subparsers(dest="command")
    serve_parser = subparsers.add_parser(
        "serve",
        help="Start a long-running HTTP API server.",
    )
    serve_parser.add_argument(
        "--api-profile",
        required=True,
        choices=tuple(API_PROFILES),
        help="HTTP API profile to expose.",
    )
    serve_parser.add_argument(
        "--host",
        default="127.0.0.1",
        help="Host interface to bind.",
    )
    serve_parser.add_argument(
        "--port",
        default=7860,
        type=int,
        help="TCP port to bind.",
    )
    serve_parser.add_argument(
        "--model-residency",
        default="cpu-cache",
        choices=("staged", "cpu-cache"),
        help="How server mode keeps model components resident.",
    )
    return parser


def inspectModels(args, user_config: UserConfig) -> None:
    """Print metadata summaries for the required model files."""

    model_sources = resolveModelSources(args, user_config)
    summaries = [
        ("diffusion model", inspectModelSource(model_sources.diffusion_model)),
        ("text encoder", inspectModelSource(model_sources.text_encoder)),
        ("VAE", inspectModelSource(model_sources.vae)),
    ]
    output = "\n\n".join(
        formatSummary(name, summary) for name, summary in summaries
    )
    print(output)


def releaseMemory() -> None:
    """Release Python and CUDA caches between large model stages."""

    gc.collect()
    try:
        import torch

        if torch.cuda.is_available():
            torch.cuda.empty_cache()
    except ImportError:
        pass


def componentDtype(config, source):
    """Return the runtime dtype for one component source."""

    if config.dtype_name != "auto":
        return config.dtype
    return inspectSourceTorchDtype(source) or config.dtype


def generate(args, user_config: UserConfig) -> list[Path]:
    """Validate generation arguments and run the current milestone."""

    config = buildGenerationConfig(args, user_config)
    model_sources = resolveModelSources(args, user_config)
    text_dtype = componentDtype(config, model_sources.text_encoder)
    diffusion_dtype = componentDtype(config, model_sources.diffusion_model)
    vae_dtype = componentDtype(config, model_sources.vae)

    text_encoder = ZImageTextEncoder(
        model_sources.text_encoder,
        config.tokenizer_path,
        config.device,
        text_dtype,
    )
    conditioning = text_encoder.encodePrompts(
        config.prompt,
        config.negative_prompt,
    )
    del text_encoder
    releaseMemory()

    model = ZImageModel(
        model_sources.diffusion_model,
        config.device,
        diffusion_dtype,
    )
    latent = sampleLatents(
        model,
        conditioning,
        batch_size=config.batch_size,
        height=config.height,
        width=config.width,
        seed=config.seed,
        steps=config.steps,
        cfg=config.cfg,
        device=config.device,
        dtype=diffusion_dtype,
    )
    del model
    releaseMemory()

    vae = ZImageVae(model_sources.vae, config.device, vae_dtype)
    images = vae.decode(latent)
    return saveImages(
        images,
        config.output,
        config.output_extension,
        config.output_quality,
    )


def main(argv: list[str] | None = None) -> int:
    """Run the CLI entry point."""

    parser = buildParser()
    args = parser.parse_args(argv)

    try:
        user_config = loadUserConfig()
        if args.command == "serve":
            server_config = validateServerConfig(
                args.api_profile,
                args.host,
                args.port,
                args.model_residency,
            )
            serve(server_config, user_config)
            return 0

        if args.inspect_models:
            inspectModels(args, user_config)
            return 0

        paths = generate(args, user_config)
    except DiffusionCliError as exc:
        print(f"error: {exc}", file=sys.stderr)
        return 2

    for path in paths:
        print(path)
    return 0


if __name__ == "__main__":
    raise SystemExit(main())