from fractions import Fraction
import os
import time
import cv2
import numpy as np
import torch
import einops
from PIL import Image
from modules import shared, errors ,timer, rife, processing
from modules.logger import log
from modules.video_models.video_utils import check_av
from modules.video_models.video_upscale import upscale_video


def get_audio_rate(p=None, default: int = 24000) -> int:
    # pipeline output wins when it reports a rate, else the loaded vocoder: LTX-2.0 runs at 24k,
    # 2.3 and 2.5 at 48k, and muxing at the wrong rate shifts the pitch
    rate = getattr(p, 'audio_sampling_rate', None) if p is not None else None
    if not rate:
        vocoder = getattr(shared.sd_model, 'vocoder', None)
        rate = getattr(getattr(vocoder, 'config', None), 'output_sampling_rate', None)
    return int(rate) if rate else default


def get_video_filename(p:processing.StableDiffusionProcessingVideo):
    from modules.image.namegen import FilenameGenerator
    from modules.paths import resolve_output_path
    namegen = FilenameGenerator(p, seed=p.seed if p is not None else 0, prompt=p.prompt if p is not None else '')
    filename = namegen.apply(shared.opts.samples_filename_pattern if shared.opts.samples_filename_pattern and len(shared.opts.samples_filename_pattern) > 0 else "[seq]-[prompt_words]")
    base_path = resolve_output_path(shared.opts.outdir_samples, shared.opts.outdir_video)
    if shared.opts.save_to_dirs:
        dirname = namegen.apply(shared.opts.directories_filename_pattern or "[prompt_words]")
        dirfile = os.path.dirname(filename)
        dirname = os.path.join(base_path, dirname, dirfile)
    else:
        dirname = base_path
    if not os.path.exists(dirname):
        os.makedirs(dirname, exist_ok=True)
    filename = os.path.join(dirname, filename)
    filename = namegen.sequence(filename)
    filename = namegen.sanitize(filename)
    return filename


def save_params(p, filename: str | None = None):
    from modules.paths import params_path
    if p is None:
        dct = {}
    else:
        # sampler_index, sampler_shift, dynamic_shift, guidance_scale, guidance_true, init_image, init_strength, last_image, vae_type, vae_tile_frames, mp4_fps, mp4_interpolate, mp4_codec, mp4_ext, mp4_opt, mp4_video, mp4_frames, mp4_sf, mp4_thumb, vlm_enhance, vlm_model, vlm_system_prompt, override_settings = args
        dct = {
            "Prompt": p.prompt,
            "Negative prompt": p.negative_prompt,
            "Steps": p.steps,
            "Sampler": p.sampler_name,
            "Seed": p.seed,
            "Engine": p.video_engine,
            "Model": p.video_model,
            "Frames": p.frames,
            "Size": f"{p.width}x{p.height}",
            "Styles": ','.join(p.styles) if isinstance(p.styles, list) else p.styles,
        }
    params = ', '.join([f'{k}: {v}' for k, v in dct.items() if v is not None and v != ''])
    fn = filename if filename is not None else params_path
    with open(fn, "w", encoding="utf8") as file:
        file.write(params)


def create_video_metadata(p: processing.StableDiffusionProcessingVideo | None, metadata: dict | None = None, filename: str | None = None):
    metadata = metadata.copy() if metadata is not None else {}
    if not shared.opts.image_metadata:
        return metadata
    if p is None:
        return metadata
    try:
        info = processing.create_infotext(p)
    except Exception as e:
        log.debug(f'Video metadata: infotext failed: {e}')
        info = ''
    if len(info) == 0:
        info = getattr(p, 'prompt', '') or ''
    if len(info) == 0:
        return metadata
    title = os.path.basename(filename) if filename else 'SD.Next video'
    metadata.setdefault('title', title)
    metadata.setdefault('encoder', 'SD.Next')
    metadata.setdefault('comment', info)
    metadata.setdefault('description', info)
    return metadata


def images_to_tensor(images):
    if images is None or len(images) == 0:
        return None
    array = [torch.from_numpy(np.array(image)) for image in images]
    tensor = torch.stack(array, dim=0) # n h w c
    tensor = tensor.unsqueeze(0) # 1, n, h, w, c
    tensor = tensor.permute(0, 4, 1, 2, 3).contiguous() # 1, c, n, h, w
    tensor = (tensor.float() / 127.5) - 1.0 # from [0,255] to [-1,1]
    # log.debug(f'Video output: images={len(images)} tensor={tensor.shape}')
    return tensor


def numpy_to_tensor(images):
    if images is None or len(images) == 0:
        return None
    images = (2.0 * images) - 1.0 # from [0,1] to [-1,1]
    array = [torch.from_numpy(images[i]) for i in range(images.shape[0])]
    tensor = torch.stack(array, dim=0) # n h w c
    tensor = tensor.unsqueeze(0) # 1, n, h, w, c
    tensor = tensor.permute(0, 4, 1, 2, 3).contiguous() # 1, c, n, h, w
    # tensor = (tensor.float() / 127.5) - 1.0 # from [0,255] to [-1,1]
    # log.debug(f'Video output: images={len(images)} tensor={tensor.shape}')
    return tensor


def add_audio_packets(container, audio_stream, audio: dict):
    if not audio or "frames" not in audio:
        return
    try:
        av = check_av()
        sr = audio.get("sr", 44100)
        layout = audio.get("layout", "stereo")
        resampler = av.AudioResampler(format="fltp", layout=layout, rate=sr)
        fifo = av.AudioFifo()
        for raw_frame in audio.get("frames", []):
            for resampled in resampler.resample(raw_frame):
                fifo.write(resampled)
        for resampled in resampler.resample(None):
            fifo.write(resampled)
        pts_counter = 0
        frame_size = audio_stream.codec_context.frame_size or 1024
        while fifo.samples >= frame_size:
            frame = fifo.read(frame_size)
            frame.pts = pts_counter
            pts_counter += frame.samples
            for packet in audio_stream.encode(frame):
                packet.stream = audio_stream
                container.mux_one(packet)
        if fifo.samples > 0:
            frame = fifo.read(fifo.samples)
            frame.pts = pts_counter
            pts_counter += frame.samples
            for packet in audio_stream.encode(frame):
                packet.stream = audio_stream
                container.mux_one(packet)
        for packet in audio_stream.encode():
            packet.stream = audio_stream
            container.mux_one(packet)
    except Exception as e:
        log.error(f"Video audio encoding: type=packets {e}")
        errors.display(e, "Audio")


def add_audio_tensor(container, audio_stream, audio: torch.Tensor, sample_rate: int):
    av = check_av()
    if torch.is_tensor(audio):
        audio = audio.detach().float().cpu().numpy()
    if audio.ndim > 2:
        audio = np.squeeze(audio)
    if audio.ndim == 1:
        audio = audio[None, :]
    elif audio.ndim == 2 and audio.shape[0] > audio.shape[1] and audio.shape[1] in (1, 2):
        audio = audio.T
    channels = audio.shape[0] if audio.shape[0] in (1, 2) else 1
    layout = "stereo" if channels == 2 else "mono"
    if audio.dtype != np.int16:
        audio = np.clip(audio, -1.0, 1.0)
        audio = (audio * 32767.0).astype(np.int16)
    audio_frame = av.AudioFrame.from_ndarray(audio, format="s16p", layout=layout)
    audio_frame.sample_rate = sample_rate
    add_audio_packets(container, audio_stream, {"sr": sample_rate, "layout": layout, "frames": [audio_frame]})


def parse_options(options):
    if isinstance(options, dict):
        return options
    if not isinstance(options, str) or not options.strip():
        return {}
    parsed_options = {}
    normalized = options.replace(',', ':') # Standardize delimiters by replacing commas with colons
    for item in normalized.split(':'):
        item = item.strip()
        if not item:
            continue
        if '=' in item:
            key, value = item.split('=', 1)
            parsed_options[key.strip()] = value.strip()
        else:
            parsed_options[item] = '1' # Handle flag options without explicit '=' (e.g., 'fastseek')
    return parsed_options


def atomic_save_video(
    filename: str,
    tensor: torch.Tensor,
    audio: torch.Tensor | None = None,
    fps: float = 24,
    codec: str = "libx264",
    pix_fmt: str = "yuv420p",
    options: str = "",
    sample_rate: int = 24000,
    metadata: dict | None = None,
    pbar=None,
):
    if metadata is None:
        metadata = {}
    av = check_av()
    if av is None:
        log.error('Video: ffmpeg/av not available')
        return
    savejob = shared.state.begin('Save video')
    frames, height, width, _channels = tensor.shape
    rate = round(fps)
    parsed_options = parse_options(options)
    log.info(f'Video: file="{filename}" codec={codec} frames={frames} width={width} height={height} fps={rate} audio={audio is not None} sample_rate={sample_rate} options={parsed_options}')
    video_array = torch.as_tensor(tensor, dtype=torch.uint8).numpy(force=True)
    task = pbar.add_task('encoding', total=frames) if pbar is not None else None
    if task is not None:
        pbar.update(task, description='video encoding')

    with av.open(filename, mode="w") as container:
        for k, v in metadata.items():
            container.metadata[k] = v
        stream: av.VideoStream = container.add_stream(codec, rate=rate, options=parsed_options)
        stream.width = video_array.shape[2]
        stream.height = video_array.shape[1]
        stream.pix_fmt = pix_fmt
        stream.time_base = Fraction(1, rate)
        audio_stream = None
        has_audio = (audio is not None) and ((torch.is_tensor(audio) or isinstance(audio, np.ndarray)) or (isinstance(audio, dict) and len(audio.get('frames', [])) > 0))
        if has_audio:
            sr = sample_rate if not isinstance(audio, dict) else audio.get("sr", sample_rate)
            layout = "stereo" if not isinstance(audio, dict) else audio.get("layout", "stereo")
            audio_stream = container.add_stream("aac", rate=sr)
            audio_stream.layout = layout
            audio_stream.time_base = Fraction(1, sr)
        for i, img in enumerate(video_array):
            frame = av.VideoFrame.from_ndarray(img, format="rgb24")
            frame.pts = i
            for packet in stream.encode_lazy(frame):
                container.mux_one(packet)
            if task is not None:
                pbar.update(task, advance=1)
        if (audio is not None) and (torch.is_tensor(audio) or isinstance(audio, np.ndarray)):
            add_audio_tensor(container, audio_stream, audio, sample_rate)
        elif (audio is not None) and isinstance(audio, dict) and len(audio.get('frames', [])) > 0:
            add_audio_packets(container, audio_stream, audio)
        for packet in stream.encode():
            container.mux_one(packet)

    shared.state.outputs(filename)
    shared.state.end(savejob)


def save_thumbnail(video_path, tensor=None):
    try:
        base = os.path.splitext(video_path)[0]
        thumb_path = f'{base}.thumb.jpg'
        if tensor is not None and len(tensor) > 0:
            frame = Image.fromarray(tensor[0].numpy())
        else:
            from modules.video import get_video_params
            _frames, _fps, _dur, _w, _h, _codec, frame = get_video_params(video_path, capture=True)
        if frame is not None:
            frame.thumbnail((512, 512), Image.Resampling.LANCZOS)
            frame.save(thumb_path, quality=80)
            return thumb_path
    except Exception as e:
        log.debug(f'Video thumbnail: {e}')
    return None


def save_video(
    p: processing.StableDiffusionProcessingVideo,
    pixels: torch.Tensor | None = None,
    audio: torch.Tensor | None = None,
    binary: bytes | None = None,
    mp4_fps: int = 24,
    mp4_codec: str = "libx264",
    mp4_opt: str = "",
    mp4_ext: str = "mp4",
    mp4_sf: bool = False,  # save safetensors
    mp4_video: bool = True,  # save video
    mp4_frames: bool = False,  # save frames
    mp4_thumb: bool = True,  # save thumbnail
    mp4_interpolate: int = 0,  # rife interpolation
    aac_sample_rate: int = 24000,  # audio sample rate
    upscale_scale: float = 1.0,  # upscale scale
    upscale_upscaler: str = "",  # upscale upscaler
    stream=None,  # async progress reporting stream
    metadata: dict | None = None,  # metadata for video
    pbar=None,  # progress bar for video
    reclamp: bool = True,  # reclamp pixels to [-1, 1] range
):
    if metadata is None:
        metadata = {}
    output_video = None

    if binary is not None:
        output_filename = get_video_filename(p)
        output_video = f'{output_filename}.{mp4_ext}'
        try:
            try:
                with open(output_video, 'wb') as f:
                    f.write(binary)
                log.info(f'Video output: file="{output_video}" size={len(binary)}')
                shared.state.outputs(output_video)
            except Exception as e:
                log.error(f'Video output: file="{output_video}" {e}')
        except Exception as e:
            log.error(f'Video output: file="{output_video}" write error {e}')
            errors.display(e, 'video')
        thumb = save_thumbnail(output_video) if mp4_thumb else None
        return 0, output_video, thumb

    if pixels is None:
        return 0, output_video, None
    if isinstance(pixels, np.ndarray):
        pixels = numpy_to_tensor(pixels)
    if isinstance(pixels, list) and isinstance(pixels[0], Image.Image):
        pixels = images_to_tensor(pixels)
    if not torch.is_tensor(pixels):
        log.error(f'Video: type={type(pixels)} not a tensor')
        return 0, output_video, None

    try:
        if mp4_interpolate > 0 and not getattr(p, 'video_interpolated', False):
            t_interpolate = time.time()
            x = pixels.squeeze(0).permute(1, 0, 2, 3)
            x = (x.clamp(-1., 1.) + 1.0) * 0.5  # RIFE expects [0, 1]; video pixels are [-1, 1]
            interpolated = rife.interpolate_nchw(x, count=mp4_interpolate+1)
            pixels = torch.stack(interpolated, dim=0)
            pixels = pixels.permute(1, 2, 0, 3, 4)
            pixels = pixels * 2.0 - 1.0
            timer.process.ts('interpolate', t_interpolate)
            p.video_interpolated = True
    except Exception as e:
        log.error(f'Video interpolate: {e}')
        errors.display(e, 'video')

    try:
        if upscale_upscaler is not None and len(upscale_upscaler) > 0 and not getattr(p, 'video_upscaled', False):
            t_upscale = time.time()
            pixels = upscale_video(pixels, scale=upscale_scale, upscaler_name=upscale_upscaler)
            timer.process.ts('upscale', t_upscale)
            p.video_upscaled = True
    except Exception as e:
        log.error(f'Video upscale: {e}')
        errors.display(e, 'video')

    t_save = time.time()
    if pixels.ndim == 4:
        pixels = pixels.unsqueeze(0)
    n, _c, t, h, w = pixels.shape
    size = pixels.element_size() * pixels.numel()
    log.debug(f'Video: video={mp4_video} export={mp4_frames} safetensors={mp4_sf} interpolate={mp4_interpolate}')
    if hasattr(audio, 'shape'):
        audio_txt = f'audio={audio.shape} aac={aac_sample_rate}' if audio is not None else 'no audio'
    elif isinstance(audio, dict):
        audio_txt = f'audio={audio.get("format", None)} packets={len(audio.get("frames", []))} '
    else:
        audio_txt = None
    log.debug(f'Video: encode={t} tensor={pixels.shape} bytes={size} {audio_txt} fps={mp4_fps} codec={mp4_codec} ext={mp4_ext} options="{mp4_opt}"')
    try:
        preparejob = shared.state.begin('Prepare video')
        if stream is not None:
            stream.output_queue.push(('progress', (None, 'Saving video...')))

        if reclamp:
            x = torch.clamp(pixels.float(), -1., 1.) * 127.5 + 127.5
        else:
            x = pixels.float() * 255.0
        x = x.detach().cpu().to(torch.uint8)
        x = einops.rearrange(x, '(m n) c t h w -> t (m h) (n w) c', n=n)
        x = x.contiguous()

        output_filename = get_video_filename(p)
        if shared.opts.save_txt:
            save_params(p, f'{output_filename}.txt')
        save_params(p)

        if mp4_sf:
            fn = f'{output_filename}.safetensors'
            log.info(f'Video export: file="{fn}" type=savetensors shape={x.shape}')
            from safetensors.torch import save_file
            shared.state.outputs(fn)
            save_file({ 'frames': x }, fn, metadata={'format': 'video', 'frames': str(t), 'width': str(w), 'height': str(h), 'fps': str(mp4_fps), 'codec': mp4_codec, 'options': mp4_opt, 'ext': mp4_ext, 'interpolate': str(mp4_interpolate)})

        if mp4_frames:
            log.info(f'Video frames: files="{output_filename}-00000.jpg" frames={t} width={w} height={h}')
            for i in range(t):
                image = cv2.cvtColor(x[i].numpy(), cv2.COLOR_RGB2BGR)
                fn = f'{output_filename}-{i:05d}.jpg'
                shared.state.outputs(fn)
                cv2.imwrite(fn, image)

        shared.state.end(preparejob)

        if mp4_video and (mp4_codec != 'none'):
            output_video = f'{output_filename}.{mp4_ext}'
            metadata = create_video_metadata(p, metadata, output_filename)
            atomic_save_video(output_video, tensor=x, audio=audio, fps=mp4_fps, codec=mp4_codec, options=mp4_opt, sample_rate=aac_sample_rate, metadata=metadata, pbar=pbar)
            if stream is not None:
                stream.output_queue.push(('progress', (None, f'Video {os.path.basename(output_video)} | Codec {mp4_codec} | Size {w}x{h}x{t} | FPS {mp4_fps}')))
                stream.output_queue.push(('file', output_video))
        else:
            if stream is not None:
                stream.output_queue.push(('progress', (None, '')))

    except Exception as e:
        log.error(f'Video save: raw={size} {e}')
        errors.display(e, 'video')
    timer.process.add('save', time.time()-t_save)
    thumb = save_thumbnail(output_video) if mp4_thumb and output_video is not None else None
    return t, output_video, thumb