"""
Shared frame extraction and CLIP retrieval utilities.

Frame extraction uses per-frame seeking to avoid loading entire videos
into memory (critical for hour-long movies on limited RAM).
"""

import os
import numpy as np
from PIL import Image
from pathlib import Path


def _save_frames_to_disk(images, cache_dir):
    """Save PIL images as JPEGs to cache_dir/000000.jpg, ..."""
    cache_dir = Path(cache_dir)
    cache_dir.mkdir(parents=True, exist_ok=True)
    for i, img in enumerate(images):
        img.save(cache_dir / f'{i:06d}.jpg', quality=85)


def _load_frames_from_disk(cache_dir):
    """Load JPEGs from cache_dir, sorted by filename. Eagerly loads to close file handles."""
    cache_dir = Path(cache_dir)
    files = sorted(cache_dir.glob('*.jpg'))
    images = []
    for f in files:
        img = Image.open(f)
        img.load()
        images.append(img)
    return images


MAX_FRAME_PIXELS = 720 * 720  # cap frame resolution: longest edge ~720px
# This ensures consistent token counts across movies regardless of source resolution.
# 1080p frames (1920x1080 = 2M pixels) get downscaled; 480p frames stay as-is.
# Only affects newly extracted frames — existing cached frames are not re-processed.


def _resize_if_needed(img, max_pixels=MAX_FRAME_PIXELS):
    """Resize image if it exceeds max_pixels, preserving aspect ratio."""
    w, h = img.size
    if w * h <= max_pixels:
        return img
    scale = (max_pixels / (w * h)) ** 0.5
    new_w, new_h = int(w * scale), int(h * scale)
    return img.resize((new_w, new_h), Image.LANCZOS)


def extract_frames_uniform(video_path, num_frames, cache_dir=None, low_memory=False):
    """
    Extract num_frames uniformly-spaced frames from a video.
    Frames are resized to MAX_FRAME_PIXELS to ensure consistent resolution.
    low_memory=True: per-frame seeking (slow but safe for limited RAM / GCP)
    low_memory=False: batch extraction (fast, needs more RAM / local GPU server)
    If cache_dir is provided and has the right count, loads from cache.
    Returns: (list[PIL.Image], list[float timestamps in seconds])
    """
    from decord import VideoReader, cpu

    # Check cache
    if cache_dir is not None:
        cache_dir = Path(cache_dir)
        existing = sorted(cache_dir.glob('*.jpg')) if cache_dir.exists() else []
        if len(existing) == num_frames:
            images = [Image.open(f) for f in existing]
            vr = VideoReader(str(video_path), ctx=cpu(0), num_threads=1)
            total = len(vr)
            indices = np.linspace(0, total - 1, num_frames, dtype=int)
            timestamps = (indices / vr.get_avg_fps()).tolist()
            del vr
            return images, timestamps

    vr = VideoReader(str(video_path), ctx=cpu(0), num_threads=1 if low_memory else 2)
    total = len(vr)
    avg_fps = vr.get_avg_fps()
    indices = np.linspace(0, total - 1, num_frames, dtype=int)
    timestamps = (indices / avg_fps).tolist()

    if cache_dir is not None:
        Path(cache_dir).mkdir(parents=True, exist_ok=True)

    if low_memory:
        images = []
        for j, idx in enumerate(indices):
            frame = vr[int(idx)].asnumpy()
            img = _resize_if_needed(Image.fromarray(frame))
            images.append(img)
            if cache_dir is not None:
                img.save(Path(cache_dir) / f'{j:06d}.jpg', quality=85)
    else:
        # Batch extraction — fast, needs RAM for all frames at once
        frames = vr.get_batch(indices).asnumpy()
        images = [_resize_if_needed(Image.fromarray(f)) for f in frames]
        if cache_dir is not None:
            _save_frames_to_disk(images, cache_dir)

    del vr
    return images, timestamps


def extract_frames_fps(video_path, fps=0.5, cache_dir=None, low_memory=False):
    """
    Extract frames at a fixed FPS rate, grouped into 1-minute chunks.
    low_memory=True: per-frame seeking (slow but safe for limited RAM / GCP)
    low_memory=False: batch extraction (fast, needs more RAM / local GPU server)
    If cache_dir is provided and has enough frames, loads from cache.
    Returns: (list of chunks [list[PIL.Image]], duration_sec)
    """
    from decord import VideoReader, cpu

    vr = VideoReader(str(video_path), ctx=cpu(0), num_threads=1 if low_memory else 2)
    native_fps = vr.get_avg_fps()
    total_frames = len(vr)
    duration_sec = total_frames / native_fps
    frames_per_minute = int(60 * fps)

    step = native_fps / fps
    indices = np.arange(0, total_frames, step).astype(int)

    # Check cache
    if cache_dir is not None:
        cache_dir = Path(cache_dir)
        existing = sorted(cache_dir.glob('*.jpg')) if cache_dir.exists() else []
        if len(existing) >= len(indices) * 0.96:
            all_images = [Image.open(f) for f in existing]
            chunks = []
            for start in range(0, len(all_images), frames_per_minute):
                end = min(start + frames_per_minute, len(all_images))
                chunks.append(all_images[start:end])
            del vr
            return chunks, duration_sec

    if cache_dir is not None:
        cache_dir.mkdir(parents=True, exist_ok=True)

    if low_memory:
        all_images = []
        for j, idx in enumerate(indices):
            frame = vr[int(idx)].asnumpy()
            img = _resize_if_needed(Image.fromarray(frame))
            all_images.append(img)
            if cache_dir is not None:
                img.save(cache_dir / f'{j:06d}.jpg', quality=85)
    else:
        # Batch extraction — fast
        frames_np = vr.get_batch(indices).asnumpy()
        all_images = [_resize_if_needed(Image.fromarray(f)) for f in frames_np]
        if cache_dir is not None:
            _save_frames_to_disk(all_images, cache_dir)

    del vr

    chunks = []
    for start in range(0, len(all_images), frames_per_minute):
        end = min(start + frames_per_minute, len(all_images))
        chunks.append(all_images[start:end])

    return chunks, duration_sec


def get_clip_embeddings(images, cache_path=None, model=None, processor=None):
    """
    Compute CLIP embeddings for a list of PIL images.
    If cache_path (.npy) exists and matches length, loads from cache.
    If model/processor not provided, loads openai/clip-vit-large-patch14.
    Returns numpy array of shape (N, D).
    """
    import torch

    # Check cache
    if cache_path is not None and os.path.exists(cache_path):
        cached = np.load(cache_path)
        if cached.shape[0] == len(images):
            return cached

    if model is None or processor is None:
        from transformers import CLIPModel, CLIPProcessor
        model = CLIPModel.from_pretrained('openai/clip-vit-large-patch14')
        processor = CLIPProcessor.from_pretrained('openai/clip-vit-large-patch14')
        model.eval()
        if torch.cuda.is_available():
            model = model.cuda()

    device = next(model.parameters()).device
    all_embs = []
    batch_size = 32
    for i in range(0, len(images), batch_size):
        batch = images[i:i + batch_size]
        inputs = processor(images=batch, return_tensors='pt', padding=True)
        inputs = {k: v.to(device) for k, v in inputs.items()}
        with torch.no_grad():
            embs = model.get_image_features(**inputs)
            if not isinstance(embs, torch.Tensor):
                embs = embs.pooler_output if hasattr(embs, 'pooler_output') else embs.last_hidden_state[:, 0]
            embs = embs / embs.norm(dim=-1, keepdim=True)
        all_embs.append(embs.cpu().numpy())

    result = np.concatenate(all_embs, axis=0)

    # Save to cache
    if cache_path is not None:
        os.makedirs(os.path.dirname(cache_path), exist_ok=True)
        np.save(cache_path, result)

    return result


def retrieve_frames_by_text(query, frame_embeddings, clip_model=None, clip_processor=None, top_k=5):
    """
    Retrieve top-k frames most similar to a text query using CLIP.
    Returns (indices, similarities) sorted by similarity descending.
    """
    import torch

    if clip_model is None or clip_processor is None:
        from transformers import CLIPModel, CLIPProcessor
        clip_model = CLIPModel.from_pretrained('openai/clip-vit-large-patch14')
        clip_processor = CLIPProcessor.from_pretrained('openai/clip-vit-large-patch14')
        clip_model.eval()
        if torch.cuda.is_available():
            clip_model = clip_model.cuda()

    device = next(clip_model.parameters()).device
    # CLIP's text encoder has a hard 77-token cap (max_position_embeddings).
    # Without truncation=True, any LLM-generated description longer than ~12
    # words crashes with `Sequence length must be less than 77`. We silently
    # truncate — the first 77 tokens carry enough signal for retrieval, and
    # the alternative is dropping the whole iteration.
    inputs = clip_processor(
        text=[query], return_tensors='pt',
        padding=True, truncation=True, max_length=77,
    )
    inputs = {k: v.to(device) for k, v in inputs.items()}
    with torch.no_grad():
        text_emb = clip_model.get_text_features(**inputs)
        if not isinstance(text_emb, torch.Tensor):
            text_emb = text_emb.pooler_output if hasattr(text_emb, 'pooler_output') else text_emb.last_hidden_state[:, 0]
        text_emb = text_emb / text_emb.norm(dim=-1, keepdim=True)

    text_emb = text_emb.cpu().numpy()
    sims = (frame_embeddings @ text_emb.T).squeeze()
    top_indices = np.argsort(sims)[::-1][:top_k]
    return top_indices.tolist(), sims[top_indices].tolist()