import os
import io
import json
import gc
import subprocess
from typing import Optional, List

import numpy as np
import torch
from PIL import Image
import ffmpeg
import gradio as gr
from transformers import CLIPProcessor, CLIPModel
import imagehash
from scipy.spatial.distance import cdist

os.environ.setdefault("OMP_NUM_THREADS", "2")
os.environ.setdefault("MKL_NUM_THREADS", "2")
torch.set_num_threads(2)
torch.set_grad_enabled(False)

CLIP_MODEL_ID = "openai/clip-vit-base-patch32"
MIN_HAMMING_THRESHOLD = 12
DTW_THRESHOLD = 0.80
MAX_CANDIDATES = 30
FRAME_TIMEOUT_SEC = 20
PROBE_TIMEOUT_SEC = 15
HTTP_RW_TIMEOUT_US = 10_000_000

_clip_model = None
_clip_processor = None
_device = "cpu"

def get_clip():
    global _clip_model, _clip_processor
    if _clip_model is None:
        _clip_processor = CLIPProcessor.from_pretrained(CLIP_MODEL_ID)
        _clip_model = CLIPModel.from_pretrained(CLIP_MODEL_ID).to(_device)
        _clip_model.eval()
    return _clip_model, _clip_processor

@torch.no_grad()
def embed_image(pil_img: Image.Image) -> np.ndarray:
    model, proc = get_clip()
    inputs = proc(images=pil_img, return_tensors="pt").to(_device)
    outputs = model.get_image_features(**inputs)
    if isinstance(outputs, torch.Tensor): embeddings = outputs
    elif hasattr(outputs, 'image_embeds'): embeddings = outputs.image_embeds
    elif hasattr(outputs, 'pooler_output'): embeddings = outputs.pooler_output
    else: embeddings = outputs[0]
    return embeddings[0].cpu().numpy().astype(np.float32)

def dtw_similarity(traj_a, traj_b):
    cost = cdist(traj_a, traj_b, metric="cosine")
    n, m = cost.shape
    D = np.full((n + 1, m + 1), np.inf)
    D[0, 0] = 0
    for i in range(1, n + 1):
        for j in range(1, m + 1):
            D[i, j] = cost[i-1, j-1] + min(D[i-1, j], D[i, j-1], D[i-1, j-1])
    return 1 - (D[n, m] / (n + m))

def probe_duration(url: str, timeout: int = PROBE_TIMEOUT_SEC) -> Optional[float]:
    try:
        cmd = ["ffprobe", "-v", "error", "-rw_timeout", str(HTTP_RW_TIMEOUT_US),
               "-show_entries", "format=duration", "-of", "default=noprint_wrappers=1:nokey=1", url]
        r = subprocess.run(cmd, capture_output=True, text=True, timeout=timeout)
        if r.returncode == 0 and r.stdout.strip(): return float(r.stdout.strip())
    except: pass
    return None

def extract_frame(url: str, seek_s: int, timeout: int = FRAME_TIMEOUT_SEC):
    """Returns (PIL.Image, error_string). If successful, error_string is None."""
    try:
        cmd = (ffmpeg.input(url, ss=seek_s, rw_timeout=HTTP_RW_TIMEOUT_US)
               .output("pipe:", vframes=1, format="image2pipe", vcodec="png")
               .overwrite_output().global_args("-loglevel", "error").compile(cmd="ffmpeg"))
        r = subprocess.run(cmd, capture_output=True, timeout=timeout)
        if r.stdout:
            return Image.open(io.BytesIO(r.stdout)).convert("RGB"), None
        else:
            err = r.stderr.decode('utf-8', errors='ignore').strip()
            return None, err
    except subprocess.TimeoutExpired:
        return None, "Timeout (20s)"
    except Exception as e:
        return None, str(e)

def match_video(payload_text: str):
    yield "🔄 Parsing JSON payload...", ""
    if not payload_text or not payload_text.strip():
        yield "❌ Empty payload.", ""
        return

    try: payload = json.loads(payload_text)
    except Exception as e:
        yield f"❌ Invalid JSON: {e}", ""
        return

    has_trajectory = "fingerprint_trajectory" in payload
    query_trajectory = np.array(payload.get("fingerprint_trajectory", []), dtype=np.float32) if has_trajectory else None
    query_mean = np.array(payload.get("visual_vector_mean", payload.get("visual_vector", [])), dtype=np.float32)
    phash_set = payload.get("phash_set", [])
    candidates: List[str] = payload.get("candidate_urls", []) or []

    if query_mean.size == 0 and query_trajectory is None:
        yield "❌ Payload must contain vector data.", ""
        return
    if len(candidates) > MAX_CANDIDATES: candidates = candidates[:MAX_CANDIDATES]

    yield f"✅ Payload parsed — {len(candidates)} candidate(s). Running pHash Pre-filter...", ""

    survivors = []
    prefilter_rejected = 0
    prefilter_errors = []
    
    for url in candidates:
        frame, err = extract_frame(url, seek_s=10, timeout=10)
        if frame is None:
            prefilter_rejected += 1
            prefilter_errors.append(f"- {url[:60]}... -> {err}")
            continue
        if phash_set:
            c_hash = imagehash.phash(frame)
            min_dist = min(c_hash - imagehash.hex_to_hash(h) for h in phash_set)
            if min_dist <= MIN_HAMMING_THRESHOLD: survivors.append(url)
            else: prefilter_rejected += 1
        else: survivors.append(url)
        del frame

    if not survivors and prefilter_errors:
        error_report = "❌ No candidates survived the pHash pre-filter.\n\n**FFmpeg Errors:**\n" + "\n".join(prefilter_errors[:5])
        yield error_report, ""
        return

    yield f"✅ Pre-filter complete: {len(survivors)} survivors, {prefilter_rejected} rejected. Loading CLIP...", ""

    if not survivors:
        yield "❌ No candidates survived.", ""
        return

    try: get_clip()
    except Exception as e:
        yield f"❌ Failed to load CLIP: {e}", ""
        return

    results = []
    for i, url in enumerate(survivors, 1):
        header = f"### Verifying Survivor {i}/{len(survivors)}\n\n**Candidate:**\n`{url[:140]}`\n"
        yield header + "⏳ Extracting 5-frame sequence...", ""

        dur = probe_duration(url, timeout=PROBE_TIMEOUT_SEC) or 10.0
        seek_points = [dur * 0.1, dur * 0.3, dur * 0.5, dur * 0.7, dur * 0.9] if dur > 20 else [0, 2, 4, 6, 8]
        
        candidate_frames = []
        for sp in seek_points:
            frame, _ = extract_frame(url, seek_s=int(sp), timeout=FRAME_TIMEOUT_SEC)
            if frame: candidate_frames.append(frame)
        
        if len(candidate_frames) < 2:
            yield header + "⚠️ Could not extract enough frames. Skipping.", ""
            continue

        try: candidate_traj = np.stack([embed_image(f) for f in candidate_frames], axis=0)
        except Exception as e:
            yield header + f"⚠️ CLIP failed: {e}", ""
            continue
        finally: del candidate_frames

        if has_trajectory:
            sim_score = dtw_similarity(query_trajectory, candidate_traj)
            score_type = "DTW"
        else:
            cand_mean = np.mean(candidate_traj, axis=0)
            sim_score = float(np.dot(query_mean, cand_mean) / (np.linalg.norm(query_mean) * np.linalg.norm(cand_mean) + 1e-8))
            score_type = "Cosine"

        duration_score = 1.0 if dur >= 30 else (dur / 30.0)
        combined_score = (0.7 * sim_score) + (0.2 * duration_score) + (0.1 * 0.5)
        
        results.append({"url": url, "score": combined_score, "sim": sim_score, "dur": dur, "type": score_type})
        yield header + f"🔎 Scored! {score_type} Sim: **{sim_score:.4f}** | Combined: **{combined_score:.4f}**", ""
        gc.collect()

    if not results:
        yield "❌ No candidates could be fully verified.", ""
        return

    results.sort(key=lambda x: x["score"], reverse=True)
    top_result = results[0]
    
    top_3_md = "### 🏆 Top 3 Ranked Results:\n"
    for idx, r in enumerate(results[:3], 1):
        top_3_md += f"{idx}. **Score:** {r['score']:.4f} | **Sim:** {r['sim']:.4f} | **Dur:** {r['dur']/60:.1f}m\n   `{r['url']}`\n"

    final_md = (
        f"## ✅ TOP MATCH CONFIRMED\n"
        f"- **URL:** `{top_result['url']}`\n"
        f"- **Combined Score:** {top_result['score']:.4f}\n"
        f"- **Visual Similarity ({top_result['type']}):** {top_result['sim']:.4f}\n"
        f"- **Duration:** {top_result['dur']/60:.1f} min\n\n"
    ) + top_3_md
    
    yield final_md, top_result["url"]

def match_from_file(file_obj):
    if file_obj is None:
        yield "❌ No file uploaded.", ""
        return
    try:
        with open(file_obj.name, 'r', encoding='utf-8') as f:
            payload_text = f.read()
        yield from match_video(payload_text)
    except Exception as e:
        yield f"❌ Error reading file: {e}", ""

with gr.Blocks(title="Space 2 — Streamer & Verifier", theme=gr.themes.Soft()) as demo:
    gr.Markdown("# 🎬 Space 2 — Streamer & Verifier")
    gr.Markdown("Upload the JSON Declaration Form from Space 1 to verify the video. Includes yt-dlp debug logging.")
    
    with gr.Row():
        with gr.Column(scale=3):
            file_in = gr.File(label="Upload JSON Declaration Form", file_types=[".json"])
            payload_in = gr.Textbox(label="Or paste JSON Payload manually", lines=6)
            with gr.Row():
                file_match_btn = gr.Button("📁 Verify Video (from File)", variant="primary")
                match_btn = gr.Button("🔍 Verify Video (from Text)", variant="secondary")
                clear_btn = gr.Button("🧹 Clear", variant="secondary")
        with gr.Column(scale=2):
            status_box = gr.Markdown("Awaiting payload...")
            result_box = gr.Textbox(label="Top Matched Video URL", interactive=False, lines=2)

    file_match_btn.click(fn=match_from_file, inputs=[file_in], outputs=[status_box, result_box], api_name="verify_video_from_file")
    match_btn.click(fn=match_video, inputs=[payload_in], outputs=[status_box, result_box], api_name="verify_video")
    clear_btn.click(fn=lambda: ("Awaiting payload...", ""), inputs=[], outputs=[status_box, result_box])

if __name__ == "__main__":
    demo.queue(default_concurrency_limit=1, max_size=10).launch(
        server_name="0.0.0.0",
        server_port=int(os.environ.get("PORT", 7860)),
        show_error=True,
    )