"""Model path resolution for local Z-Image Turbo files."""

from __future__ import annotations

from pathlib import Path

from diffusion_cli.config import (
    CHECKPOINT_DIFFUSION_PREFIX,
    CHECKPOINT_TEXT_ENCODER_PREFIX,
    CHECKPOINT_VAE_PREFIX,
    DIFFUSION_ROLE,
    TEXT_ENCODER_ROLE,
    VAE_ROLE,
    ModelPathConfig,
    ModelFiles,
    ModelSource,
    ModelSources,
    UserConfig,
)
from diffusion_cli.errors import DiffusionCliError


def requireFile(path: Path | None, role: str, config_key: str) -> Path:
    """Return an existing regular file or raise a direct CLI error."""

    if path is None:
        raise DiffusionCliError(f"Missing {config_key}")

    resolved_path = path.expanduser().resolve()
    if not resolved_path.is_file():
        raise DiffusionCliError(f"Missing {role}: {resolved_path}")
    return resolved_path


def _pathFromArgs(args, name: str) -> Path | None:
    return getattr(args, name, None)


def _requireSelectedFile(path: Path, role: str) -> Path:
    resolved_path = path.expanduser().resolve()
    if not resolved_path.is_file():
        raise DiffusionCliError(f"Missing {role}: {resolved_path}")
    return resolved_path


def resolveComponentSource(
    cli_path: Path | None,
    config_path: Path | None,
    checkpoint_path: Path | None,
    role: str,
    checkpoint_prefix: str,
    config_key: str,
) -> ModelSource:
    """Resolve one component path using component-before-checkpoint order."""

    path = cli_path or config_path
    if path is not None:
        return ModelSource(
            path=_requireSelectedFile(path, role),
            role=role,
            checkpoint_prefix=None,
        )

    if checkpoint_path is None:
        raise DiffusionCliError(f"Missing {config_key} or models.checkpoint")

    return ModelSource(
        path=_requireSelectedFile(checkpoint_path, "checkpoint"),
        role=role,
        checkpoint_prefix=checkpoint_prefix,
    )


def resolveModelSources(args, user_config: UserConfig) -> ModelSources:
    """Resolve and validate model sources for the active workflow."""

    models = user_config.models
    checkpoint = _pathFromArgs(args, "checkpoint") or models.checkpoint

    return ModelSources(
        diffusion_model=resolveComponentSource(
            _pathFromArgs(args, "diffusion_model"),
            models.diffusion_model,
            checkpoint,
            DIFFUSION_ROLE,
            CHECKPOINT_DIFFUSION_PREFIX,
            "models.diffusion_model",
        ),
        text_encoder=resolveComponentSource(
            _pathFromArgs(args, "text_encoder"),
            models.text_encoder,
            checkpoint,
            TEXT_ENCODER_ROLE,
            CHECKPOINT_TEXT_ENCODER_PREFIX,
            "models.text_encoder",
        ),
        vae=resolveComponentSource(
            _pathFromArgs(args, "vae"),
            models.vae,
            checkpoint,
            VAE_ROLE,
            CHECKPOINT_VAE_PREFIX,
            "models.vae",
        ),
    )


def resolveModelSourcesFromConfig(models: ModelPathConfig) -> ModelSources:
    """Resolve model sources using only user configuration values."""

    checkpoint = models.checkpoint
    return ModelSources(
        diffusion_model=resolveComponentSource(
            None,
            models.diffusion_model,
            checkpoint,
            DIFFUSION_ROLE,
            CHECKPOINT_DIFFUSION_PREFIX,
            "models.diffusion_model",
        ),
        text_encoder=resolveComponentSource(
            None,
            models.text_encoder,
            checkpoint,
            TEXT_ENCODER_ROLE,
            CHECKPOINT_TEXT_ENCODER_PREFIX,
            "models.text_encoder",
        ),
        vae=resolveComponentSource(
            None,
            models.vae,
            checkpoint,
            VAE_ROLE,
            CHECKPOINT_VAE_PREFIX,
            "models.vae",
        ),
    )


def resolveModelFiles(args, user_config: UserConfig) -> ModelFiles:
    """Resolve standalone model paths for compatibility with old callers."""

    sources = resolveModelSources(args, user_config)
    if any(
        source.checkpoint_prefix is not None
        for source in (
            sources.diffusion_model,
            sources.text_encoder,
            sources.vae,
        )
    ):
        raise DiffusionCliError(
            "resolveModelFiles cannot return checkpoint-backed sources"
        )
    return ModelFiles(
        diffusion_model=sources.diffusion_model.path,
        text_encoder=sources.text_encoder.path,
        vae=sources.vae.path,
    )