#!/usr/bin/env python3
"""
search_media standalone — substitui a MCP tool do app GPL por SigLIP Apache-2.0 local.

Implementa busca semântica de frames: dado um texto (query), acha os momentos nos vídeos
que melhor correspondem visualmente. Mesmo contrato de saída do search_media original:
  {"moments": [{"name": "<footage>", "startSeconds": float, "endSeconds": float, "score": float}]}

Arquitetura (clean-room, sem GPL):
  1. Carrega modelo SigLIP do HuggingFace (google/siglip-base-patch16-256, Apache-2.0)
     via transformers + torch. Alternativa: OpenCLIP com checkpoint LAION (MIT).
  2. Extrai frames dos vídeos em media_paths.json (ffmpeg, read-only, 1 fps default)
  3. Embed de imagem (SigLIP vision) + embed de texto (SigLIP text)
  4. Cosine similarity → top-K hits por query
  5. Cache de embeddings em disco (.search_cache/ no workdir) — re-rodar é instantâneo

Uso:
  # Construir índice de todos os vídeos
  python search_media.py --index

  # Buscar
  python search_media.py --query "the bride walking down the aisle" --limit 50

  # Como módulo Python (usado por director_moments.py)
  from search_media import search, ensure_index

Dependências (instalar no venv do VideoTask):
  pip install torch transformers pillow

Modelos testados:
  - google/siglip-base-patch16-256  (~400MB, bom balanço qualidade/velocidade)
  - google/siglip-so400m-patch14-384 (~1.6GB, melhor qualidade)
  Fallback: open_clip com ViT-B/32 (MIT, ~300MB) se transformers não disponível.
"""
import json, os, sys, struct, hashlib, time, subprocess
import numpy as np

FPS = 30
FRAME_INTERVAL = 1.0  # segundos entre frames indexados
FFMPEG = "/opt/homebrew/bin/ffmpeg"

SCRIPT_DIR = os.path.dirname(os.path.abspath(__file__))


# ══════════════════════════════════════════════════════════════════════════════
# Model loading (SigLIP via transformers, fallback OpenCLIP)
# ══════════════════════════════════════════════════════════════════════════════

_MODEL = None
_PROCESSOR = None
_DEVICE = None
_MODEL_NAME = os.environ.get("SIGLIP_MODEL", "google/siglip-base-patch16-256")


def _get_device():
    global _DEVICE
    if _DEVICE is None:
        try:
            import torch
            if torch.backends.mps.is_available():
                _DEVICE = "mps"
            elif torch.cuda.is_available():
                _DEVICE = "cuda"
            else:
                _DEVICE = "cpu"
        except Exception:
            _DEVICE = "cpu"
    return _DEVICE


def _load_siglip_transformers():
    """Carrega SigLIP via transformers (Apache-2.0)."""
    from transformers import AutoProcessor, AutoModel
    import torch

    device = _get_device()
    print(f"  Carregando {_MODEL_NAME} no device '{device}'...", flush=True)
    processor = AutoProcessor.from_pretrained(_MODEL_NAME)
    model = AutoModel.from_pretrained(_MODEL_NAME).to(device)
    model.eval()
    return model, processor, device


def _load_openclip():
    """Fallback: OpenCLIP (MIT) com checkpoint público."""
    import open_clip
    import torch

    device = _get_device()
    # Checkpoint público LAION-2B (MIT) — funciona bem pra busca visual genérica
    model_name = "ViT-B-32"
    pretrained = "laion2b_s34b_b79k"
    print(f"  Carregando OpenCLIP {model_name}/{pretrained} no device '{device}'...", flush=True)
    model, _, preprocess = open_clip.create_model_and_transforms(model_name, pretrained=pretrained)
    model = model.to(device)
    model.eval()
    tokenizer = open_clip.get_tokenizer(model_name)
    return model, preprocess, tokenizer, device


def load_model():
    """Carrega o modelo (SigLIP preferido, OpenCLIP fallback). Thread-safe após 1º load."""
    global _MODEL, _PROCESSOR, _DEVICE
    if _MODEL is not None:
        return _MODEL, _PROCESSOR, _DEVICE

    # Tenta SigLIP primeiro
    try:
        _MODEL, _PROCESSOR, _DEVICE = _load_siglip_transformers()
        return _MODEL, _PROCESSOR, _DEVICE
    except ImportError as e:
        msg = str(e)
        if "transformers" in msg.lower() or "torch" in msg.lower():
            print("  transformers/torch não instalado, tentando open_clip...", flush=True)
        else:
            # Ex: sentencepiece ausente — relança claro (melhor que cair em fallback errado)
            raise RuntimeError(
                f"SigLIP requer dependência adicional: {e}\n"
                f"  pip install sentencepiece"
            ) from e
    except Exception as e:
        print(f"  SigLIP falhou: {e}", flush=True)

    # Fallback OpenCLIP
    try:
        import open_clip
        _MODEL, preprocess, tokenizer, _DEVICE = _load_openclip()
        _PROCESSOR = (preprocess, tokenizer)
        return _MODEL, _PROCESSOR, _DEVICE
    except ImportError:
        raise RuntimeError(
            "Nenhum modelo disponível. Instale:\n"
            "  pip install transformers torch  # SigLIP Apache-2.0 (recomendado)\n"
            "  pip install open_clip_torch     # OpenCLIP MIT (fallback)"
        )


def embed_image_siglip(model, processor, device, image):
    """Embed de imagem SigLIP."""
    import torch
    inputs = processor(images=image, return_tensors="pt").to(device)
    with torch.no_grad():
        emb = model.get_image_features(**inputs)
    emb = emb / emb.norm(dim=-1, keepdim=True)
    return emb.cpu().numpy().astype("float32")[0]


def embed_text_siglip(model, processor, device, text):
    """Embed de texto SigLIP."""
    import torch
    inputs = processor(text=text, return_tensors="pt", padding=True, truncation=True).to(device)
    with torch.no_grad():
        emb = model.get_text_features(**inputs)
    emb = emb / emb.norm(dim=-1, keepdim=True)
    return emb.cpu().numpy().astype("float32")[0]


def embed_image_openclip(model, processor, device, image):
    """Embed de imagem OpenCLIP."""
    import torch
    preprocess, _ = processor
    img_tensor = preprocess(image).unsqueeze(0).to(device)
    with torch.no_grad():
        emb = model.encode_image(img_tensor)
    emb = emb / emb.norm(dim=-1, keepdim=True)
    return emb.cpu().numpy().astype("float32")[0]


def embed_text_openclip(model, processor, device, text):
    """Embed de texto OpenCLIP."""
    import torch
    _, tokenizer = processor
    text_tokens = tokenizer([text]).to(device)
    with torch.no_grad():
        emb = model.encode_text(text_tokens)
    emb = emb / emb.norm(dim=-1, keepdim=True)
    return emb.cpu().numpy().astype("float32")[0]


def _is_siglip():
    """True se o modelo carregado for SigLIP (via transformers)."""
    try:
        from transformers import PreTrainedModel
        return isinstance(_MODEL, PreTrainedModel)
    except ImportError:
        return False


# ══════════════════════════════════════════════════════════════════════════════
# Frame extraction & indexing
# ══════════════════════════════════════════════════════════════════════════════

def _frame_path(cache_dir, footage, t_sec):
    """Caminho do frame extraído (cache de imagem)."""
    return os.path.join(cache_dir, "frames", f"{footage}_{t_sec:.1f}.jpg")


def _extract_frame(video_path, output_path, t_sec):
    """Extrai 1 frame do vídeo em `t_sec` segundos. Thread-safe (tmp + os.replace)."""
    if os.path.exists(output_path):
        return True
    os.makedirs(os.path.dirname(output_path), exist_ok=True)
    tmp = f"{output_path}.{os.getpid()}.tmp"
    proc = subprocess.run(
        [FFMPEG, "-v", "error", "-ss", str(t_sec), "-i", video_path,
         "-frames:v", "1", "-f", "image2", "-y", tmp],
        capture_output=True, timeout=30
    )
    if proc.returncode == 0 and os.path.exists(tmp) and os.path.getsize(tmp) > 100:
        os.replace(tmp, output_path)
        return True
    if os.path.exists(tmp):
        os.remove(tmp)
    return False


def _video_duration(video_path):
    """Duração do vídeo em segundos (ffprobe)."""
    proc = subprocess.run(
        [FFMPEG.replace("ffmpeg", "ffprobe"), "-v", "error",
         "-show_entries", "format=duration", "-of", "json", video_path],
        capture_output=True, text=True, timeout=15
    )
    try:
        return float(json.loads(proc.stdout)["format"]["duration"])
    except Exception:
        return 0.0


def index_videos(workdir):
    """
    Constrói índice de embeddings para todos os vídeos em media_paths.json.
    Cacheia frames e embeddings em <workdir>/.search_cache/.

    Retorna: dict {footage_name: [(t_sec, embedding_npy_path), ...]}
    """
    media_paths_file = os.path.join(workdir, "media_paths.json")
    if not os.path.exists(media_paths_file):
        raise FileNotFoundError(f"{media_paths_file} não encontrado")

    media = json.load(open(media_paths_file))
    cache_dir = os.path.join(workdir, ".search_cache")
    emb_dir = os.path.join(cache_dir, "embeddings")
    os.makedirs(emb_dir, exist_ok=True)
    os.makedirs(os.path.join(cache_dir, "frames"), exist_ok=True)

    model, processor, device = load_model()
    siglip = _is_siglip()
    embed_fn = embed_image_siglip if siglip else embed_image_openclip
    from PIL import Image

    index = {}
    total_frames = 0
    t0 = time.time()

    for footage, video_path in sorted(media.items()):
        if not os.path.exists(video_path):
            print(f"  SKIP {footage}: arquivo não encontrado", flush=True)
            continue

        dur = _video_duration(video_path)
        if dur <= 0:
            print(f"  SKIP {footage}: duração zero", flush=True)
            continue

        frames = []
        for t_sec in np.arange(0, dur, FRAME_INTERVAL):
            t_sec = round(t_sec, 1)
            fpath = _frame_path(cache_dir, footage, t_sec)

            # Extrai frame se não cacheado
            if not os.path.exists(fpath):
                if not _extract_frame(video_path, fpath, t_sec):
                    continue

            # Embed se não cacheado
            emb_hash = hashlib.md5(f"{footage}_{t_sec}".encode()).hexdigest()[:16]
            epath = os.path.join(emb_dir, f"{emb_hash}.npy")
            if not os.path.exists(epath):
                try:
                    img = Image.open(fpath).convert("RGB")
                    emb = embed_fn(model, processor, device, img)
                    np.save(epath, emb)
                except Exception as e:
                    print(f"  ! embed falhou {footage}@{t_sec}s: {e}", flush=True)
                    continue

            frames.append((t_sec, epath))
            total_frames += 1

        if frames:
            index[footage] = frames
            print(f"  {footage}: {len(frames)} frames ({dur:.0f}s)", flush=True)

    dt = time.time() - t0
    print(f"  Índice: {total_frames} frames em {len(index)} vídeos ({dt:.0f}s)", flush=True)

    # Salva metadados do índice
    meta = {"videos": list(index.keys()), "total_frames": total_frames,
            "frame_interval": FRAME_INTERVAL, "model": _MODEL_NAME if siglip else "open_clip"}
    with open(os.path.join(cache_dir, "index_meta.json"), "w") as f:
        json.dump(meta, f)

    return index


def _load_index(workdir):
    """Carrega índice do cache em disco."""
    cache_dir = os.path.join(workdir, ".search_cache")
    emb_dir = os.path.join(cache_dir, "embeddings")
    meta_file = os.path.join(cache_dir, "index_meta.json")

    if not os.path.exists(meta_file):
        return None

    meta = json.load(open(meta_file))
    media = json.load(open(os.path.join(workdir, "media_paths.json")))

    index = {}
    for footage in meta["videos"]:
        if footage not in media:
            continue
        dur = _video_duration(media[footage])
        frames = []
        for t_sec in np.arange(0, dur, FRAME_INTERVAL):
            t_sec = round(t_sec, 1)
            emb_hash = hashlib.md5(f"{footage}_{t_sec}".encode()).hexdigest()[:16]
            epath = os.path.join(emb_dir, f"{emb_hash}.npy")
            if os.path.exists(epath):
                frames.append((t_sec, epath))
        if frames:
            index[footage] = frames

    return index if index else None


def _cosine_similarity(a, b):
    """Cosine similarity entre dois vetores numpy normalizados."""
    return float(np.dot(a, b))


# ══════════════════════════════════════════════════════════════════════════════
# Interface pública (mesmo contrato do search_media original)
# ══════════════════════════════════════════════════════════════════════════════

_QUERY_CACHE = {}  # cache de embeddings de texto (queries repetidas)


def search(workdir, query, limit=50):
    """
    Busca semântica: acha os top-K momentos nos vídeos indexados.

    Args:
        workdir: diretório com media_paths.json e .search_cache/
        query: texto da busca (ex: "the bride walking down the aisle")
        limit: número máximo de resultados

    Returns:
        dict com {"moments": [{"name": "<footage>", "startSeconds": float,
                                "endSeconds": float, "score": float}]}
    """
    # Carrega índice (constrói se não existir)
    index = _load_index(workdir)
    if index is None:
        index = index_videos(workdir)

    if not index:
        return {"moments": []}

    model, processor, device = load_model()
    siglip = _is_siglip()

    # Embed da query (com cache)
    cache_key = query.strip().lower()
    if cache_key in _QUERY_CACHE:
        query_emb = _QUERY_CACHE[cache_key]
    else:
        if siglip:
            query_emb = embed_text_siglip(model, processor, device, query)
        else:
            query_emb = embed_text_openclip(model, processor, device, query)
        _QUERY_CACHE[cache_key] = query_emb

    # Varre todos os frames, calcula similaridade
    hits = []
    for footage, frames in index.items():
        for t_sec, epath in frames:
            try:
                frame_emb = np.load(epath)
                score = _cosine_similarity(query_emb, frame_emb)
                # SigLIP cosine similarity após L2-norm é baixa (~0.01 relevante,~0.05 excelente).
                # Sem threshold: deixa o caller (director_moments) filtrar por score.
                hits.append({
                    "name": footage,
                    "startSeconds": t_sec,
                    "endSeconds": round(t_sec + FRAME_INTERVAL, 1),
                    "score": round(float(score), 4)
                })
            except Exception:
                continue

    # Ordena por score decrescente, top-K
    hits.sort(key=lambda h: -h["score"])
    return {"moments": hits[:limit]}


def ensure_index(workdir):
    """Garante que o índice existe (constrói se necessário). Thread-safe idempotente."""
    index = _load_index(workdir)
    if index is None:
        index = index_videos(workdir)
    return index


# ══════════════════════════════════════════════════════════════════════════════
# CLI
# ══════════════════════════════════════════════════════════════════════════════

if __name__ == "__main__":
    import argparse
    ap = argparse.ArgumentParser(description="search_media standalone (SigLIP Apache-2.0)")
    ap.add_argument("--workdir", default=os.getcwd(),
                    help="Diretório com media_paths.json (default: cwd)")
    ap.add_argument("--index", action="store_true", help="Construir índice de todos os vídeos")
    ap.add_argument("--query", type=str, help="Query de busca semântica")
    ap.add_argument("--limit", type=int, default=50, help="Limite de resultados (default: 50)")
    args = ap.parse_args()

    workdir = os.path.abspath(args.workdir)

    if args.index:
        index_videos(workdir)
    elif args.query:
        result = search(workdir, args.query, args.limit)
        print(json.dumps(result, indent=1))
    else:
        ap.print_help()
