import time
from contextlib import contextmanager
import torch
from PIL import Image
from modules import devices, shared, sd_models, timer, extra_networks
from modules.logger import log


def get_bucket(size: int):
    # LTX pipes validate width/height divisible by 32 across all families.
    ratio = getattr(shared.sd_model, 'vae_spatial_compression_ratio', None)
    if not isinstance(ratio, int) or ratio < 32:
        ratio = 32
    size = int(size)
    return size - (size % ratio)


def get_frames(frames: int):
    return int(8 * (int(frames) // 8)) + 1


def load_model(engine: str, model: str) -> str:
    if model is None or model == '' or model == 'None':
        shared.sd_model = None
        return 'Video model unloaded'
    from modules.video_models import models_def, video_load
    selected = models_def.find(engine, model)
    if selected is None: # the dropdown lists the separators it groups models under, and they name no model
        msg = f'Video model not loaded: engine="{engine}" model="{model}"'
        log.warning(msg)
        return msg
    t0 = time.time()
    # video_load owns the cache; pipe-class mismatch inside it invalidates the name-based hit
    # when Unload Models (or any external swap) silently replaced shared.sd_model.
    log.info(f'Load video: engine="{engine}" selected="{model}" {selected}')
    msg = video_load.load_model(selected)
    t1 = time.time()
    shared.sd_model = sd_models.apply_balanced_offload(shared.sd_model)
    t2 = time.time()
    timer.process.add('load', t1 - t0)
    timer.process.add('offload', t2 - t1)
    return msg or f'Video model loaded: {selected.name}'


def upsample_pipe_stale(upsample_pipe, upsample_repo_id) -> bool:
    # bound to one repo and borrowing the model's VAE; both change with the model, and a VAE from
    # an unloaded model holds meta tensors
    if upsample_pipe is None:
        return False
    if getattr(upsample_pipe, 'sdnext_upsample_repo', None) != upsample_repo_id:
        return True
    return upsample_pipe.vae is not getattr(shared.sd_model, 'vae', None)


def load_upsample(upsample_pipe, upsample_repo_id):
    if upsample_pipe_stale(upsample_pipe, upsample_repo_id):
        upsample_pipe = None
    if upsample_pipe is None:
        t0 = time.time()
        from diffusers.pipelines.ltx.pipeline_ltx_latent_upsample import LTXLatentUpsamplePipeline
        log.info(f'Load video: cls={LTXLatentUpsamplePipeline.__name__} repo="{upsample_repo_id}"')
        upsample_pipe = LTXLatentUpsamplePipeline.from_pretrained(
            upsample_repo_id,
            vae=shared.sd_model.vae,
            cache_dir=shared.opts.hfcache_dir,
            torch_dtype=devices.dtype,
        )
        # only the upsampler, since the pipe borrows the model's vae and moving the whole pipe
        # would drag that along into the meta tensors the caller's offload exclude avoids
        upsample_pipe.latent_upsampler.to(devices.device)
        upsample_pipe.sdnext_upsample_repo = upsample_repo_id
        t1 = time.time()
        timer.process.add('load', t1 - t0)
    return upsample_pipe


def load_upsample_2x(upsample_pipe, upsample_repo_id, variant: str = '2.x'):
    # 2.x ships the upsampler as a bare nn.Module in a subfolder; no from_pretrained on the
    # pipeline wrapper, so we load the model + construct the pipeline manually.
    if upsample_pipe_stale(upsample_pipe, upsample_repo_id):
        upsample_pipe = None
    if upsample_pipe is None:
        t0 = time.time()
        from diffusers.pipelines.ltx2.pipeline_ltx2_latent_upsample import LTX2LatentUpsamplePipeline
        from diffusers.pipelines.ltx2.latent_upsampler import LTX2LatentUpsamplerModel
        from modules import sd_checkpoint
        log.info(f'Load video: cls={LTX2LatentUpsamplePipeline.__name__} repo="{upsample_repo_id}"')
        latent_upsampler = LTX2LatentUpsamplerModel.from_pretrained(
            upsample_repo_id,
            subfolder='latent_upsampler',
            cache_dir=shared.opts.hfcache_dir,
            torch_dtype=devices.dtype,
        ).to(devices.device)
        upsample_pipe = LTX2LatentUpsamplePipeline(
            vae=shared.sd_model.vae,
            latent_upsampler=latent_upsampler,
        )
        # Synthetic checkpoint_info gives this pipe its own OffloadHook cache slot, so routing
        # it through apply_balanced_offload does not invalidate the main pipe's module map
        # (sd_offload keys on sd_checkpoint_info.name). Variant is in the name since each loads
        # different weights and the slot also names the disk offload folder.
        upsample_pipe.sdnext_upsample_repo = upsample_repo_id
        upsample_pipe.sd_checkpoint_info = sd_checkpoint.CheckpointInfo(f'ltx-upsampler-{variant}')
        t1 = time.time()
        timer.process.add('load', t1 - t0)
    return upsample_pipe


def scheduler_shift_key(scheduler) -> str | None:
    # UniPC and its relatives call the static shift flow_shift; flow-match schedulers call it shift.
    config = getattr(scheduler, 'config', None)
    if config is None:
        return None
    for key in ('flow_shift', 'shift'):
        if hasattr(config, key):
            return key
    return None


@contextmanager
def ltx_scheduler_opts(sd_model, *, dynamic_shift=None, sampler_shift=None, shift_terminal=None):
    # Run-scoped override of shared.opts scheduler settings and scheduler.config, restored on every
    # exit path. Keeps run settings out of config.json, protects default_scheduler from a deepcopy
    # of the mutated scheduler, and restores the scheduler object that Stage 2 refine swaps out.
    orig_dynamic_shift = shared.opts.schedulers_dynamic_shift
    orig_sampler_shift = shared.opts.schedulers_shift
    orig_scheduler = sd_model.scheduler
    orig_default_scheduler = getattr(sd_model, 'default_scheduler', None)
    restore = {}

    def write_config(values: dict):
        scheduler = getattr(sd_model, 'scheduler', None)
        if not values or scheduler is None or not hasattr(scheduler, 'config') or not hasattr(scheduler, 'register_to_config'):
            return
        for key, value in values.items():
            setattr(scheduler.config, key, value)
        scheduler.register_to_config(**values)

    def override_config(key, value):
        # only written keys are restored, so a key stored as None returns to None
        if key is None or value is None or not hasattr(getattr(orig_scheduler, 'config', None), key):
            return
        restore[key] = getattr(orig_scheduler.config, key)
        write_config({key: value})

    try:
        if dynamic_shift is not None:
            shared.opts.data['schedulers_dynamic_shift'] = dynamic_shift
        if sampler_shift is not None:
            shared.opts.data['schedulers_shift'] = sampler_shift
        override_config('use_dynamic_shifting', dynamic_shift)
        if sampler_shift is not None and sampler_shift >= 0:
            override_config(scheduler_shift_key(orig_scheduler), sampler_shift)
        override_config('shift_terminal', shift_terminal)
        yield
    finally:
        shared.opts.data['schedulers_dynamic_shift'] = orig_dynamic_shift
        shared.opts.data['schedulers_shift'] = orig_sampler_shift
        if sd_model.scheduler is not orig_scheduler:
            sd_model.scheduler = orig_scheduler
        if orig_default_scheduler is not None and sd_model.default_scheduler is not orig_default_scheduler:
            sd_model.default_scheduler = orig_default_scheduler
        write_config(restore)


def _condition_cls(family: str):
    if family == '2.x':
        try:
            from diffusers.pipelines.ltx2.pipeline_ltx2_condition import LTX2VideoCondition
            return LTX2VideoCondition
        except ImportError:
            log.warning('LTX conditions: LTX2VideoCondition not available in installed diffusers')
            return None
    from diffusers.pipelines.ltx.pipeline_ltx_condition import LTXVideoCondition
    return LTXVideoCondition


def make_condition(condition_cls, family: str, frames, strength: float, is_video: bool, index: int = 0):
    if family == '2.x':
        return condition_cls(frames=frames, index=index, strength=strength)
    if is_video:
        return condition_cls(video=frames, frame_index=index, strength=strength)
    return condition_cls(image=frames, frame_index=index, strength=strength)


def open_condition(src) -> Image.Image:
    """A conditioning source as a PIL image, from a gradio upload handle or from an api reference.

    A string goes to the api decoder, which reads base64 and upload refs, rather than being opened
    as a path: naming a file is the caller's way of reading one it never uploaded.
    """
    if hasattr(src, 'name'):
        return Image.open(src.name)
    if isinstance(src, str):
        from modules.api.api import decode_base64_to_image
        return decode_base64_to_image(src)
    return src


def get_conditions(width, height, condition_strength, condition_images, condition_files, condition_video, condition_video_frames, condition_video_skip, family: str = '0.9', num_frames=None, condition_last=None):
    condition_cls = _condition_cls(family)
    if condition_cls is None:
        return []
    conditions = []
    if condition_images is not None:
        for condition_image in condition_images:
            try:
                condition_image = open_condition(condition_image).convert('RGB').resize((width, height), resample=Image.Resampling.LANCZOS)
                conditions.append(make_condition(condition_cls, family, condition_image, condition_strength, is_video=False))
                log.debug(f'Video condition: family={family} image={condition_image.size} strength={condition_strength}')
            except Exception as e:
                log.error(f'LTX condition image: {e}')
    if condition_files is not None:
        batch_images = []
        for fn in condition_files:
            try:
                batch_images.append(open_condition(fn).convert('RGB').resize((width, height), resample=Image.Resampling.LANCZOS))
            except Exception as e:
                log.error(f'LTX condition files: {e}')
        if len(batch_images) > 0:
            conditions.append(make_condition(condition_cls, family, batch_images, condition_strength, is_video=True))
            log.debug(f'Video condition: family={family} files={len(batch_images)} size={batch_images[0].size} strength={condition_strength}')
    if condition_video is not None:
        from modules.video_models.video_utils import get_video_frames
        try:
            condition_frames = get_video_frames(condition_video, num_frames=condition_video_frames, skip_frames=condition_video_skip)
            condition_frames = [f.convert('RGB').resize((width, height), resample=Image.Resampling.LANCZOS) for f in condition_frames]
            if len(condition_frames) > 0:
                conditions.append(make_condition(condition_cls, family, condition_frames, condition_strength, is_video=True))
                log.debug(f'Video condition: family={family} frames={len(condition_frames)} size={condition_frames[0].size} strength={condition_strength}')
        except Exception as e:
            log.error(f'LTX condition video: {e}')
    if condition_last is not None:
        try:
            condition_last = open_condition(condition_last).convert('RGB').resize((width, height), resample=Image.Resampling.LANCZOS)
            # 2.x reads index as a latent index and accepts -1 for the final frame; 0.9 uses a pixel index.
            last_index = -1 if family == '2.x' else max((num_frames or 1) - 1, 0)
            conditions.append(make_condition(condition_cls, family, condition_last, condition_strength, is_video=False, index=last_index))
            log.debug(f'Video condition: family={family} last={condition_last.size} index={last_index} strength={condition_strength}')
        except Exception as e:
            log.error(f'LTX condition last image: {e}')
    return conditions


def get_prompts(p):
    prompt = shared.prompt_styles.apply_styles_to_prompt(p.prompt, p.styles)
    negative = shared.prompt_styles.apply_negative_styles_to_prompt(p.negative_prompt, p.styles)
    prompts, networks = extra_networks.parse_prompts([prompt])
    prompt = prompts[0] if len(prompts) > 0 else prompt
    return prompt, negative, networks


def get_generator(seed):
    import random
    if seed is None or seed < 0:
        random.seed()
        seed = int(random.randrange(4294967294))
    return torch.Generator().manual_seed(seed)


def vae_decode(latents, decode_timestep, seed, denormalize: bool = True):
    t0 = time.time()
    if latents.ndim == 4:
        latents = latents.unsqueeze(0)
    log.debug(f'Video: cls={shared.sd_model.vae.__class__.__name__} op=vae latents={latents.shape} timestep={decode_timestep} denormalize={denormalize}')
    from diffusers.utils.torch_utils import randn_tensor
    if denormalize:
        latents = shared.sd_model._denormalize_latents( # pylint: disable=protected-access
            latents,
            shared.sd_model.vae.latents_mean,
            shared.sd_model.vae.latents_std,
            shared.sd_model.vae.config.scaling_factor
        )
    latents = latents.to(device=devices.device, dtype=devices.dtype)
    if not shared.sd_model.vae.config.timestep_conditioning:
        timestep = None
    else:
        noise = randn_tensor(latents.shape, generator=get_generator(seed), device=devices.device, dtype=devices.dtype)
        timestep = torch.tensor([decode_timestep], device=devices.device, dtype=latents.dtype)
        noise_scale = torch.tensor([decode_timestep], device=devices.device, dtype=devices.dtype)[:, None, None, None, None]
        latents = (1 - noise_scale) * latents + noise_scale * noise
    frames = shared.sd_model.vae.decode(latents, timestep, return_dict=False)[0]
    t1 = time.time()
    timer.process.add('vae', t1 - t0)
    return frames