from dataclasses import dataclass, field
from typing import Optional

from modules.logger import log


@dataclass
class LTXCaps:
    name: str
    repo_cls_name: str
    family: str  # '0.9' or '2.x'
    variant: str  # '0.9', '2.0', '2.3', '2.5' (finer-grained sub-variant)
    is_distilled: bool
    is_i2v: bool
    supports_input_media: bool
    supports_multi_condition: bool
    supports_image_cond_noise_scale: bool
    supports_decode_timestep: bool
    supports_stg: bool
    supports_audio: bool
    supports_frame_rate_kwarg: bool
    # 2.3 transformer cross-attn reads the other modality's sigma; unset falls back to 2.0's
    # independent-sigma path, which is a joint-distribution mismatch for 2.3 weights.
    use_cross_timestep: bool
    default_cfg: float
    default_steps: int
    default_sampler_shift: float
    default_dynamic_shift: bool
    default_width: int
    default_height: int
    default_frames: int
    default_frame_rate: int
    stg_default_scale: float = 0.0
    stg_default_blocks: list = field(default_factory=list)
    # Dev 2.x trained under cfg + stg + modality + rescale four-way composition;
    # distilled bakes these into its sigma schedule and stays at pipeline identity.
    modality_default_scale: float = 1.0
    guidance_rescale_default: float = 0.0
    supports_canonical_stage2: bool = False
    # Stage 1 at target/2 -> 2x upsample -> Stage 2 at target. Distilled needs upsampled latents
    # for refine; same-res refine oversaturates. Condition variants rebuild conditions per stage.
    supports_two_stage_refine: bool = False
    stage2_dev_lora_repo: Optional[str] = None
    # 2.5 keeps the stage 2 LoRA in the model repo; 2.0 and 2.3 each have their own
    stage2_dev_lora_weight: Optional[str] = None
    # tied to the family VAE: the wrong one drifts per-channel latent statistics
    upsample_repo: Optional[str] = None
    # 2.5 ships the distilled scheduler config, so its Dev rows restore the terminal shift
    scheduler_shift_terminal: Optional[float] = None
    supports_auto_duration: bool = False


CONDITION_CLASSES = {'LTXConditionPipeline', 'LTX2ConditionPipeline'}
LTX2_CLASSES = {'LTX2Pipeline', 'LTX2ImageToVideoPipeline', 'LTX2ConditionPipeline'}
ALL_LTX_CLASSES = {
    'LTXPipeline',
    'LTXImageToVideoPipeline',
    'LTXConditionPipeline',
    'LTX2Pipeline',
    'LTX2ImageToVideoPipeline',
    'LTX2ConditionPipeline',
}


def _repo_cls_name(model_name: str) -> Optional[str]:
    from modules.video_models.models_def import models
    entries = models.get('LTX Video', [])
    for m in entries:
        if m.name == model_name:
            if m.repo_cls is None:
                return None
            elif isinstance(m.repo_cls, str):
                return m.repo_cls
            else:
                return m.repo_cls.__name__
    return None


def get_caps(model_name: str) -> Optional[LTXCaps]:
    if not model_name or model_name == 'None':
        return None
    cls_name = _repo_cls_name(model_name)
    if cls_name is None:
        log.warning(f'LTX caps: model="{model_name}" has no repo_cls registered')
        return None
    if cls_name not in ALL_LTX_CLASSES:
        log.warning(f'LTX caps: model="{model_name}" repo_cls="{cls_name}" is not an LTX pipeline')
        return None

    is_ltx2 = cls_name in LTX2_CLASSES
    family = '2.x' if is_ltx2 else '0.9'
    # 2.x sub-variant detection: unknown 2.x mirrors fall through to '2.0' (conservative default).
    if is_ltx2:
        if '2.5' in model_name:
            variant = '2.5'
        elif '2.3' in model_name:
            variant = '2.3'
        else:
            variant = '2.0'
    else:
        variant = '0.9'
    is_distilled = 'Distilled' in model_name
    is_i2v = 'I2V' in model_name or cls_name in ('LTXImageToVideoPipeline', 'LTX2ImageToVideoPipeline')
    is_condition_cls = cls_name in CONDITION_CLASSES
    supports_input_media = is_i2v or is_condition_cls

    caps = LTXCaps(
        name=model_name,
        repo_cls_name=cls_name,
        family=family,
        variant=variant,
        is_distilled=is_distilled,
        is_i2v=is_i2v,
        supports_input_media=supports_input_media,
        supports_multi_condition=is_condition_cls,
        supports_image_cond_noise_scale=(cls_name == 'LTXConditionPipeline'),
        supports_decode_timestep=(family == '0.9'),
        supports_stg=is_ltx2,
        supports_audio=is_ltx2,
        supports_frame_rate_kwarg=is_ltx2,
        use_cross_timestep=variant in ('2.3', '2.5'),
        default_cfg=3.0,
        default_steps=30 if is_ltx2 else 50,
        default_sampler_shift=-1.0,
        # distilled ships use_dynamic_shifting=False and runs explicit sigmas, which shifting remaps
        default_dynamic_shift=is_ltx2 and not is_distilled,
        default_width=768,
        default_height=512,
        default_frames=121 if is_ltx2 else 161,
        default_frame_rate=24 if is_ltx2 else 25,
    )

    if is_distilled:
        caps.default_cfg = 1.0
        caps.default_steps = 8

    if is_ltx2 and not is_distilled:
        if variant == '2.5':
            caps.stage2_dev_lora_repo = 'Lightricks/LTX-2.5-Diffusers'
            caps.stage2_dev_lora_weight = 'ltx-2.5-22b-distilled-lora-450-bf16.safetensors'
        elif variant == '2.3':
            caps.stage2_dev_lora_repo = 'CalamitousFelicitousness/LTX-2.3-distilled-lora-384-Diffusers'
        elif variant == '2.0':
            caps.stage2_dev_lora_repo = 'CalamitousFelicitousness/LTX-2.0-distilled-lora-384-Diffusers'
    caps.supports_canonical_stage2 = caps.stage2_dev_lora_repo is not None
    caps.supports_two_stage_refine = is_ltx2

    if is_ltx2:
        if variant == '2.5':
            caps.upsample_repo = 'Lightricks/LTX-2.5-Diffusers'
            caps.stg_default_blocks = [28]
            caps.supports_auto_duration = True
        elif variant == '2.3':
            caps.upsample_repo = 'CalamitousFelicitousness/LTX-2.3-Spatial-Upsampler-x2-1.1-Diffusers'
            caps.stg_default_blocks = [28]
        else:
            caps.upsample_repo = 'Lightricks/LTX-2'
            caps.stg_default_blocks = [29]
        if not is_distilled:
            # canonical T2V composition from huggingface/diffusers#13217
            caps.stg_default_scale = 1.0
            caps.modality_default_scale = 3.0
            caps.guidance_rescale_default = 0.7
            if variant == '2.5':
                caps.scheduler_shift_terminal = 0.1

    return caps