"""
TRIBE v2 Neural Engagement Analyzer — HuggingFace Space
Runs on T4 GPU with full WhisperX transcription pipeline.
Falls back to audio-only mode if WhisperX fails.
"""

import json
import os
import subprocess
import sys
import tempfile
import threading
import time
import traceback
from pathlib import Path
from urllib.parse import urlparse

# Disable hf_transfer — uvx sandboxes don't have it
os.environ["HF_HUB_ENABLE_HF_TRANSFER"] = "0"

# Persistent storage, when the Space has it: the WhisperX environment (uvx),
# the whisper large-v3 weights and every HF model then survive a restart
# instead of being downloaded again on each cold start.
if os.path.isdir("/data") and os.access("/data", os.W_OK):
    os.environ.setdefault("UV_CACHE_DIR", "/data/uv-cache")
    os.environ.setdefault("HF_HOME", "/data/hf-home")
    for _d in (os.environ["UV_CACHE_DIR"], os.environ["HF_HOME"]):
        os.makedirs(_d, exist_ok=True)
    print("Persistent /data found: uv and HF caches live there")

import gradio as gr
import numpy as np
import pandas as pd
import requests

# Set HF_TOKEN
hf_token = os.environ.get("HF_TOKEN", "")
if hf_token:
    os.environ["HUGGING_FACE_HUB_TOKEN"] = hf_token
    print("HF token configured")

import torch

CUDA = torch.cuda.is_available()
print(f"torch {torch.__version__} | cuda={CUDA}" + (f" | {torch.cuda.get_device_name(0)} | {torch.cuda.get_device_properties(0).total_memory / 2**30:.1f} GiB" if CUDA else ""))
if CUDA:
    torch.backends.cudnn.benchmark = True


# Frames are cached per video as uint8 arrays; at 16 fps a 60 s ad is ~960
# frames of 320x568, about half a GiB, well within the Space's RAM.
MAX_CACHED_FRAMES = 2400
# Clips per VJEPA2 forward. Each clip's fp16 hidden states are ~0.9 GiB (40
# layers x 8192 tokens x 1408); two fit comfortably next to the fp32 weights
# and the fp16 Llama on a 16 GiB T4. Raise on a bigger GPU.
CLIP_BATCH = int(os.environ.get("TRIBE_CLIP_BATCH", "2"))


def _patch_video_pipeline() -> None:
    """Two fixes to neuralset's HuggingFaceVideo path, measured on the T4.

    1. Frame decoding. For every 4 s clip the extractor asked moviepy for 64
       frames by timestamp, and each get_frame() seeks and decodes through
       ffmpeg on the CPU: ~13 s per clip, the whole cost of "Encoding video".
       We decode each video once, sequentially, into memory and serve frames
       by index. Same frames, same PIL output, no per-frame seeks.
    2. The VJEPA2 (ViT-Giant) forward ran in fp32. Under fp16 autocast the T4's
       tensor cores take it; hidden states are cast back to fp32 so nothing
       downstream sees a dtype change.
    """
    try:
        from neuralset.extractors import video as ns_video
    except Exception as e:  # pragma: no cover
        print(f"[perf] could not import neuralset.extractors.video: {e}")
        return

    import PIL.Image
    from tqdm import tqdm

    frame_cache: dict[int, tuple[np.ndarray, float]] = {}
    original_read = ns_video._VideoImage._read

    def _read_cached(self):
        video = self.video
        key = id(video)
        entry = frame_cache.get(key)
        if entry is None:
            n_est = int(round(float(video.duration) * float(video.fps))) + 1
            if n_est > MAX_CACHED_FRAMES:
                return original_read(self)
            t0 = time.time()
            frames = np.stack([f.astype("uint8", copy=False) for f in video.iter_frames(dtype="uint8")])
            frame_cache.clear()  # one video at a time under the analysis lock
            entry = (frames, float(video.fps))
            frame_cache[key] = entry
            print(f"[perf] decoded {len(frames)} frames @ {video.fps:.1f} fps once in {time.time() - t0:.1f}s ({frames.nbytes / 2**20:.0f} MiB)")
        frames, fps = entry
        idx = min(len(frames) - 1, max(0, int(round(float(self.time) * fps))))
        return PIL.Image.fromarray(frames[idx])

    ns_video._VideoImage._read = _read_cached

    if CUDA:
        original_predict = ns_video._HFVideoModel.predict

        def _predict_fp16(self, images, audio=None):
            if getattr(self.model, "device", torch.device("cpu")).type != "cuda":
                return original_predict(self, images, audio)
            with torch.autocast("cuda", dtype=torch.float16):
                pred = original_predict(self, images, audio)
            hs = getattr(pred, "hidden_states", None)
            if hs is not None:
                pred.hidden_states = tuple(h.float() for h in hs)
            return pred

        ns_video._HFVideoModel.predict = _predict_fp16

    # Image-model path (DINOv2 for image_feature), same idea.
    try:
        from neuralset.extractors import image as ns_image
        original_hidden = ns_image.HuggingFaceImage._get_hidden_states

        def _fp16_hidden_states(self, images):
            if CUDA and str(self.device).startswith("cuda"):
                with torch.autocast("cuda", dtype=torch.float16):
                    states = original_hidden(self, images)
                return [st.float() for st in states]
            return original_hidden(self, images)

        ns_image.HuggingFaceImage._get_hidden_states = _fp16_hidden_states
    except Exception as e:  # pragma: no cover
        print(f"[perf] image extractor patch skipped: {e}")
    # 3. Several clips per GPU forward. neuralset encodes one 4 s clip per
    #    forward, so the T4 idles while the next clip's 64 frames are gathered
    #    and preprocessed on the CPU. We gather CLIP_BATCH clips, run one
    #    forward, and aggregate per clip exactly as the original did (token
    #    aggregation in fp32 on one clip at a time, so peak memory stays that
    #    of a single clip's hidden states).
    def _forward_clips(model, clips):
        """clips: list of (num_frames, H, W, 3) uint8 arrays -> list of (L, tokens, F) fp16 tensors."""
        kwargs = {"return_tensors": "pt", "videos": [list(c) for c in clips]}
        inputs = model.processor(**kwargs)
        ns_video._fix_pixel_values(inputs)
        inputs = inputs.to(model.model.device)
        with torch.inference_mode():
            if CUDA and model.model.device.type == "cuda":
                with torch.autocast("cuda", dtype=torch.float16):
                    pred = model.model(**inputs)
            else:
                pred = model.model(**inputs)
        states = pred.hidden_states  # tuple of L tensors, each (B, tokens, F)
        out = []
        for b in range(len(clips)):
            out.append(torch.stack([st[b] for st in states]))  # (L, tokens, F)
        return out

    # The loader asks for the same video's features more than once per job
    # (one call per segment it builds). exca's disk cache used to make the
    # repeats free; this memo does the same in memory, for the current video.
    feature_memo: dict[tuple, "ns_video.nsbase.TimedArray"] = {}

    def _get_data_batched(self, events):
        if "vjepa2" not in self.image.model_name:
            yield from _get_data_original(self, events)
            return
        pending = []
        for event in events:
            key = (str(event.filepath), round(float(getattr(event, "start", 0.0) or 0.0), 3), round(float(event.duration or 0.0), 3))
            hit = feature_memo.get(key)
            if hit is not None:
                print(f"[perf] video features served from memo for {key[0].rsplit('/', 1)[-1]} ({key[1]}s+{key[2]}s)")
                yield hit
            else:
                pending.append((key, event))
        if not pending:
            return
        for key, ta in zip([k for k, _ in pending], _encode_events(self, [e for _, e in pending])):
            if any(k[0] != key[0] for k in feature_memo):
                feature_memo.clear()  # a new video: drop the previous one's features
            feature_memo[key] = ta
            yield ta

    def _encode_events(self, events):
        model = ns_video._HFVideoModel(
            model_name=self.image.model_name, pretrained=self.image.pretrained,
            layer_type=self.layer_type, num_frames=self.num_frames,
        )
        if model.model.device.type == "cpu":
            model.model.to(self.image.device)
        freq = events[0].frequency if self.frequency == "native" else self.frequency
        T = 1 / freq if self.clip_duration is None else self.clip_duration
        subtimes = [k / model.num_frames * T for k in reversed(range(model.num_frames))]
        for event in events:
            video = event.read()
            freq = self.frequency if self.frequency != "native" else event.frequency
            expect_frames = ns_video.nsbase.Frequency(freq).to_ind(event.duration)
            times = np.linspace(0, video.duration, expect_frames + 1)[1:]
            output = None
            t_start = time.time()
            for start in tqdm(range(0, len(times), CLIP_BATCH), desc=f"Encoding video (x{CLIP_BATCH})"):
                chunk = times[start:start + CLIP_BATCH]
                clips = []
                for t in chunk:
                    ims = [ns_video._VideoImage(video=video, time=max(0, t - t2)) for t2 in subtimes]
                    pil_imgs = [i.read() for i in ims]
                    if pil_imgs and self.max_imsize is not None:
                        factor = max(pil_imgs[0].size) / self.max_imsize
                        if factor > 1:
                            size = tuple(int(sz / factor) for sz in pil_imgs[0].size)
                            pil_imgs = [pi.resize(size) for pi in pil_imgs]
                    clips.append(np.array([np.array(pi) for pi in pil_imgs]))
                for k, t_embd in zip(range(start, start + len(chunk)), _forward_clips(model, clips)):
                    embd = self.image._aggregate_tokens(t_embd.float()).cpu().numpy()
                    if not self.image.cache_all_layers and self.image.cache_n_layers is None:
                        embd = self.image._aggregate_layers(embd)
                    if output is None:
                        output = np.zeros((len(times),) + embd.shape)
                    output[k] = embd
                del clips
            video.close()
            print(f"[perf] encoded {len(times)} clips in {time.time() - t_start:.0f}s ({(time.time() - t_start) / max(1, len(times)):.2f}s/clip, batch {CLIP_BATCH})")
            output = output.transpose(list(range(1, output.ndim)) + [0])
            yield ns_video.nsbase.TimedArray(
                data=output.astype(np.float32), frequency=freq,
                start=ns_video.nsbase._UNSET_START, duration=event.duration,
            )

    # exca's infra.apply wraps _get_data with a per-event feature cache. That
    # cache is built for re-running studies over the same stimuli; every ad here
    # is analysed once, so the wrapper buys nothing and the method is replaced
    # outright. The wrapped original is kept for any non-VJEPA2 model.
    _get_data_original = ns_video.HuggingFaceVideo._get_data
    ns_video.HuggingFaceVideo._get_data = _get_data_batched
    patched_loop = True

    print(f"[perf] video pipeline patched: cached frame decoding, fp16 autocast={'on' if CUDA else 'off'}, clip batching={'x%d' % CLIP_BATCH if patched_loop else 'off'}")


def _patch_whisperx_language() -> None:
    """Let WhisperX detect the language instead of the library's hard-coded English.

    Ads are short and speech starts early, so Whisper's detection on the first
    30 s is reliable here; forcing English turns Italian or Hebrew copy into
    garbage words, and garbage words become garbage text features.
    Set TRIBE_WHISPER_LANGUAGE=english (or another supported name) to force one.
    """
    try:
        from tribev2 import eventstransforms as et
    except Exception as e:  # pragma: no cover
        print(f"[whisperx] could not import tribev2.eventstransforms: {e}")
        return
    forced = os.environ.get("TRIBE_WHISPER_LANGUAGE", "auto").strip().lower()
    original = et.ExtractWordsFromAudio._get_transcript_from_audio

    def _transcribe(wav_filename, language):
        import json as _json
        import subprocess as _sp
        import tempfile as _tf

        if forced != "auto":
            return original(wav_filename, forced)
        device = "cuda" if torch.cuda.is_available() else "cpu"
        with _tf.TemporaryDirectory() as output_dir:
            cmd = [
                "uvx", "whisperx", str(wav_filename),
                "--model", "large-v3",
                "--device", device,
                "--compute_type", "float16" if device == "cuda" else "int8",
                "--batch_size", "16",
                "--output_dir", output_dir,
                "--output_format", "json",
            ]
            env = {k: v for k, v in os.environ.items() if k != "MPLBACKEND"}
            t0 = time.time()
            result = _sp.run(cmd, capture_output=True, text=True, env=env)
            if result.returncode != 0:
                raise RuntimeError(f"whisperx failed:\n{result.stderr[-3000:]}")
            json_path = Path(output_dir) / f"{Path(wav_filename).stem}.json"
            transcript = _json.loads(json_path.read_text())
        detected = transcript.get("language", "?")
        n_words = sum(len(seg.get("words", [])) for seg in transcript.get("segments", []))
        print(f"[whisperx] language={detected} words={n_words} in {time.time() - t0:.0f}s")
        words = []
        for i, segment in enumerate(transcript.get("segments", [])):
            sentence = segment.get("text", "").replace('"', "")
            for word in segment.get("words", []):
                if "start" not in word or "end" not in word:
                    continue
                words.append({
                    "text": word["word"].replace('"', ""),
                    "start": word["start"],
                    "duration": word["end"] - word["start"],
                    "sequence_id": i,
                    "sentence": sentence,
                })
        return pd.DataFrame(words)

    et.ExtractWordsFromAudio._get_transcript_from_audio = staticmethod(_transcribe)
    print(f"[whisperx] language mode: {forced}")


def _prewarm_whisperx() -> None:
    """Build the uvx whisperx environment once at startup so the first job does not pay for it."""
    def run():
        try:
            t0 = time.time()
            r = subprocess.run(["uvx", "whisperx", "--help"], capture_output=True, text=True, timeout=1800)
            print(f"[whisperx] prewarm {'ok' if r.returncode == 0 else 'FAILED'} in {time.time() - t0:.0f}s" + ("" if r.returncode == 0 else f": {r.stderr[-800:]}"))
        except Exception as e:
            print(f"[whisperx] prewarm error: {e}")
    threading.Thread(target=run, daemon=True, name="whisperx-prewarm").start()


_patch_video_pipeline()
_patch_whisperx_language()
_prewarm_whisperx()

FUNCTIONAL_REGIONS = {
    "visual_processing": {
        "description": "Visual cortex",
        "labels": ["G_cuneus", "S_calcarine", "G_occipital_sup", "G_occipital_middle", "Pole_occipital", "G_and_S_occipital_inf", "G_oc-temp_lat-fusifor"],
    },
    "attention": {
        "description": "Prefrontal and parietal",
        "labels": ["G_front_sup", "G_front_middle", "S_front_sup", "G_parietal_sup", "S_intrapariet_and_P_trans", "G_precuneus"],
    },
    "emotional_response": {
        "description": "Limbic and temporal",
        "labels": ["G_cingul-Post-dorsal", "G_temp_sup-Lateral", "S_temporal_sup", "G_and_S_cingul-Ant"],
    },
    "language_processing": {
        "description": "Broca and Wernicke",
        "labels": ["G_front_inf-Opercular", "G_front_inf-Triangul", "G_temp_sup-G_T_transv", "G_temporal_middle", "G_pariet_inf-Angular"],
    },
    "motor_planning": {
        "description": "Motor and premotor",
        "labels": ["G_precentral", "S_precentral-sup-part", "S_central"],
    },
}

WEIGHTS = {"attention": 0.30, "visual_processing": 0.25, "emotional_response": 0.25, "language_processing": 0.15, "motor_planning": 0.05}

model = None
label_map = None


def build_events_without_whisperx(video_path: str) -> pd.DataFrame:
    """
    Build events dataframe from video WITHOUT using WhisperX.
    Extracts audio and creates Video+Audio events only.
    TRIBE v2 can still predict brain responses from video and audio
    features even without word-level transcription.
    """
    from neuralset.events.transforms import ExtractAudioFromVideo, ChunkEvents
    from neuralset.events.utils import standardize_events

    # Create initial video event
    event = {
        "type": "Video",
        "filepath": str(video_path),
        "start": 0,
        "timeline": "default",
        "subject": "default",
    }
    events = pd.DataFrame([event])
    events = standardize_events(events)

    # Extract audio from video
    transform_audio = ExtractAudioFromVideo()
    events = transform_audio(events)

    # Chunk long clips — keep short for faster processing
    transform_chunk_audio = ChunkEvents(event_type_to_chunk="Audio", max_duration=15, min_duration=5)
    transform_chunk_video = ChunkEvents(event_type_to_chunk="Video", max_duration=15, min_duration=5)
    events = transform_chunk_audio(events)
    events = transform_chunk_video(events)

    events = standardize_events(events)
    print(f"Built events: {len(events)} rows (audio_only mode, no WhisperX)")
    return events


def load_model():
    global model, label_map
    if model is not None:
        return

    from tribev2 import TribeModel
    from nilearn.datasets import fetch_atlas_surf_destrieux

    print("Loading TRIBE v2 model...")
    # Text features come from Llama-3.2-3B. In fp32 that is ~12.8 GiB and cannot
    # share a 16 GiB T4 with VJEPA2 (~4.4 GiB); "accelerate" loads it in fp16
    # with device_map=auto (~6.4 GiB, spilling to CPU rather than OOM-ing).
    # data.num_workers: the training config forks 20 DataLoader workers. Here
    # one video is analysed at a time and its features need the GPU, so the
    # extractors must run in this process: a forked worker cannot initialise
    # CUDA and the job dies inside the DataLoader. Zero workers, no fork.
    config_update = {"data.num_workers": 0}
    if CUDA:
        config_update["data.text_feature.device"] = "accelerate"
    model = TribeModel.from_pretrained("facebook/tribev2", cache_folder="./cache", config_update=config_update)

    print("Loading Destrieux atlas...")
    atlas = fetch_atlas_surf_destrieux()
    label_map = {}
    labels = atlas.labels

    for i, label_entry in enumerate(labels):
        if isinstance(label_entry, (tuple, list)) and len(label_entry) == 2:
            idx, name = label_entry
        elif isinstance(label_entry, bytes):
            idx, name = i, label_entry.decode("utf-8")
        elif isinstance(label_entry, str):
            idx, name = i, label_entry
        else:
            idx, name = i, str(label_entry)

        if name == "Unknown" or idx == 0:
            continue
        left_verts = np.where(np.array(atlas.map_left) == idx)[0]
        right_verts = np.where(np.array(atlas.map_right) == idx)[0] + len(atlas.map_left)
        if len(left_verts) > 0 or len(right_verts) > 0:
            label_map[name] = np.concatenate([left_verts, right_verts])

    print(f"Loaded {len(label_map)} brain regions. Model ready.")


def compute_regions(preds):
    region_activations = {}
    for rname, rinfo in FUNCTIONAL_REGIONS.items():
        verts = []
        for lbl in rinfo["labels"]:
            if lbl in label_map:
                verts.append(label_map[lbl])
        if not verts:
            region_activations[rname] = np.zeros(preds.shape[0])
            continue
        all_v = np.concatenate(verts)
        all_v = all_v[all_v < preds.shape[1]]
        region_activations[rname] = np.mean(preds[:, all_v], axis=1)
    return region_activations


def compute_engagement(region_activations):
    """
    Compute per-second engagement from region activations.
    Uses z-score normalization instead of min-max to avoid the
    "everything peaks at the last second" artifact.
    """
    n = len(next(iter(region_activations.values())))
    composite = np.zeros(n)
    for region, weight in WEIGHTS.items():
        if region in region_activations:
            vals = region_activations[region]
            mean, std = vals.mean(), vals.std()
            if std > 0:
                # Z-score: how many SDs above/below mean activation
                z = (vals - mean) / std
                # Sigmoid to [0, 1] range — z=0 maps to 0.5
                normalized = 1 / (1 + np.exp(-z))
            else:
                normalized = np.full(n, 0.5)
            composite += weight * normalized
    return composite


def compute_score(engagement):
    """
    Score 1-10 based on variability and peak strength.
    Good ads have high peaks and dynamic variation (not flat).
    """
    if len(engagement) == 0:
        return 5
    peak = float(np.max(engagement))
    valley = float(np.min(engagement))
    mean = float(np.mean(engagement))
    dynamic_range = peak - valley  # 0 to ~1
    # Weight: 60% peak strength, 40% dynamic range
    raw = 0.6 * peak + 0.4 * dynamic_range
    # Scale to 1-10 (0.3 → 1, 0.8 → 10)
    score = 1 + 9 * max(0, min(1, (raw - 0.3) / 0.5))
    return max(1, min(10, round(score)))


def find_peak_region_seconds(region_activations, hrf_offset):
    """
    Find per-region peak seconds, excluding the very last second
    (which is an artifact of the model's attention accumulation).
    """
    summary = {}
    for rname, rinfo in FUNCTIONAL_REGIONS.items():
        if rname not in region_activations:
            continue
        vals = region_activations[rname]
        # Exclude last 2 time points (attention accumulation artifact)
        search_range = vals[:-2] if len(vals) > 3 else vals
        peak_idx = int(np.argmax(search_range))
        mean_val = float(np.mean(vals))
        peak_val = float(np.max(search_range))
        # Relative strength: how much above mean is the peak
        rel_strength = (peak_val - mean_val) / (mean_val + 1e-8)
        summary[rname] = {
            "description": rinfo["description"],
            "mean": round(mean_val, 4),
            "peak": round(peak_val, 4),
            "peak_second": peak_idx + hrf_offset,
            "relative_strength": round(rel_strength, 2),
        }
    return summary


def generate_insights(engagement, region_activations, hrf_offset=0, ad_context=None):
    """Generate data-specific creative insights from brain activation patterns."""
    # Ad context for tailored language
    ctx = ad_context or {}
    offer = ctx.get("offer", "key message")  # "the app", "the service", "the shoe", etc.
    cta_label = ctx.get("cta", "CTA")
    brand = ctx.get("brand", "the brand")

    insights = []
    n = len(engagement)
    if n < 3:
        return insights

    mean_eng = float(np.mean(engagement))
    std_eng = float(np.std(engagement))
    peak_idx = int(np.argmax(engagement))
    low_idx = int(np.argmin(engagement))
    peak_val = float(engagement[peak_idx])
    low_val = float(engagement[low_idx])

    # ── Compute per-region stats ──
    region_stats = {}
    for rname in ["visual_processing", "attention", "emotional_response", "language_processing", "motor_planning"]:
        if rname in region_activations:
            vals = region_activations[rname]
            safe = vals[:-2] if len(vals) > 3 else vals
            region_stats[rname] = {
                "mean": float(np.mean(safe)),
                "std": float(np.std(safe)),
                "peak_idx": int(np.argmax(safe)),
                "peak_val": float(np.max(safe)),
            }

    sorted_regions = sorted(region_stats.items(), key=lambda x: x[1]["mean"], reverse=True) if region_stats else []
    top_region = sorted_regions[0][0] if sorted_regions else None
    bottom_region = sorted_regions[-1][0] if sorted_regions else None

    # ── 1. Engagement Shape Analysis ──
    # Split into thirds and identify the shape
    third = max(1, n // 3)
    start_eng = float(np.mean(engagement[:third]))
    mid_eng = float(np.mean(engagement[third:2*third]))
    end_eng = float(np.mean(engagement[2*third:]))

    if start_eng > mid_eng and start_eng > end_eng:
        shape = "front-loaded"
    elif end_eng > mid_eng and end_eng > start_eng:
        shape = "back-loaded"
    elif mid_eng < start_eng and mid_eng < end_eng:
        shape = "u-shape"
    elif mid_eng > start_eng and mid_eng > end_eng:
        shape = "peak-middle"
    else:
        shape = "flat"

    shape_insights = {
        "front-loaded": {
            "title": "Front-Loaded Engagement",
            "severity": "medium",
            "message": f"Brain engagement is strongest in the first {third}s, then declines. Viewers get hooked early but lose interest before {offer} is reinforced.",
            "action": f"Move {offer} earlier — viewers are most receptive in the opening seconds. Or add a second hook in the final third to re-engage before the {cta_label}."
        },
        "back-loaded": {
            "title": "Slow Build — Late Peak",
            "severity": "medium",
            "message": f"Engagement builds toward the end, peaking after {2*third}s. The opening is the weakest section — viewers in a feed may scroll past before {offer} lands.",
            "action": f"Tease the payoff upfront: show a flash of the climax in the first 2 seconds to buy time. Or restructure to reveal {offer} immediately."
        },
        "u-shape": {
            "title": "Mid-Section Drop",
            "severity": "medium",
            "message": f"Strong opening and closing, but engagement dips in the middle ({third}s-{2*third}s). Attention drops between the hook and the {cta_label}.",
            "action": f"Add a pattern interrupt between {third}s-{2*third}s: scene change, new speaker, text overlay, or tempo shift. Bridge the gap between hook and {cta_label}."
        },
        "peak-middle": {
            "title": "Center-Weighted Engagement",
            "severity": "positive",
            "message": f"The core of the ad ({third}s-{2*third}s) drives the strongest response — this is where {offer} resonates most.",
            "action": f"The structure works. Ensure the {cta_label} appears during or immediately after this peak, not after engagement fades."
        },
        "flat": {
            "title": "Flat Response — Low Contrast",
            "severity": "medium",
            "message": f"Brain engagement is uniform throughout — no strong peaks or valleys. Flat response correlates with lower recall for {brand}.",
            "action": "Create contrast: add a surprising moment, an emotional beat, a visual disruption, or a change in pacing. Memorable ads need at least one peak."
        },
    }
    insights.append({"type": "shape", **shape_insights[shape]})

    # ── 2. Hook Verdict ──
    hook_3s = float(np.mean(engagement[:min(3, n)]))
    hook_vs_avg = (hook_3s - mean_eng) / (mean_eng + 1e-8) * 100

    if hook_vs_avg < -15:
        # Find what's WEAK in the hook
        hook_regions = {r: float(np.mean(region_activations[r][:min(3, n)])) for r in region_stats}
        weakest_hook_region = min(hook_regions, key=hook_regions.get) if hook_regions else None
        weakness_label = {
            "visual_processing": "visuals aren't grabbing attention",
            "attention": "nothing demands focus",
            "emotional_response": "no emotional hook",
            "language_processing": "the copy/audio doesn't land immediately",
            "motor_planning": "no impulse to engage",
        }.get(weakest_hook_region, "the opening lacks stimulus")

        insights.append({
            "type": "hook", "severity": "high",
            "title": f"Weak Hook — {abs(hook_vs_avg):.0f}% Below Average",
            "message": f"The first 3 seconds underperform the rest of the ad. Specifically, {weakness_label} in the opening.",
            "action": "Restructure the first 2 seconds: lead with motion, a face, bold text, or a sound that demands attention. The hook determines 80% of view-through rate."
        })
    elif hook_vs_avg > 15:
        # Find what's STRONG in the hook
        hook_regions = {r: float(np.mean(region_activations[r][:min(3, n)])) for r in region_stats}
        strongest_hook_region = max(hook_regions, key=hook_regions.get) if hook_regions else None
        strength_label = {
            "visual_processing": "strong visual opening",
            "attention": "immediately grabs focus",
            "emotional_response": "emotional hook",
            "language_processing": "compelling opening line/audio",
            "motor_planning": "action-triggering opener",
        }.get(strongest_hook_region, "effective opening")

        insights.append({
            "type": "hook", "severity": "positive",
            "title": f"Strong Hook — {hook_vs_avg:.0f}% Above Average",
            "message": f"The opening outperforms the rest of the ad, driven by {strength_label}. This hook pattern is worth replicating.",
            "action": "Document this hook structure — it works. Test variations of the same opening pattern across other creatives."
        })

    # ── 3. Engagement Drops — find the worst dip ──
    if n >= 6:
        # Find the biggest drop (3-second rolling window)
        window = 3
        rolling = [float(np.mean(engagement[max(0,i-1):i+window-1])) for i in range(n - window + 1)]
        worst_window_idx = int(np.argmin(rolling))
        worst_window_val = rolling[worst_window_idx]
        drop_pct = (mean_eng - worst_window_val) / (mean_eng + 1e-8) * 100

        if drop_pct > 15:
            drop_start = worst_window_idx + hrf_offset
            drop_end = worst_window_idx + window + hrf_offset
            # What region drops the most here?
            drop_regions = {}
            for rname in region_stats:
                region_at_drop = float(np.mean(region_activations[rname][worst_window_idx:worst_window_idx+window]))
                region_overall = region_stats[rname]["mean"]
                drop_regions[rname] = (region_overall - region_at_drop) / (region_overall + 1e-8)
            biggest_drop_region = max(drop_regions, key=drop_regions.get) if drop_regions else None
            drop_label = {
                "visual_processing": "visual interest drops — the scene may be static or repetitive",
                "attention": "cognitive focus fades — nothing new to process",
                "emotional_response": "emotional connection weakens",
                "language_processing": "the message loses clarity or relevance",
                "motor_planning": "action intent drops — the ad feels passive",
            }.get(biggest_drop_region, "engagement drops")

            insights.append({
                "type": "drop", "severity": "high",
                "title": f"Attention Drop at {drop_start}s-{drop_end}s ({drop_pct:.0f}% below avg)",
                "message": f"The weakest moment in the ad. Specifically, {drop_label} during this window.",
                "action": f"Review what happens at {drop_start}s-{drop_end}s in the video. Add a scene change, new audio cue, or text overlay to re-engage."
            })

    # ── 4. Peak Moment — what makes it work ──
    peak_s = peak_idx + hrf_offset
    # Find what region is strongest at the peak
    peak_regions = {r: float(region_activations[r][peak_idx]) for r in region_stats}
    peak_driver = max(peak_regions, key=peak_regions.get) if peak_regions else None
    peak_label = {
        "visual_processing": "a strong visual moment",
        "attention": "high cognitive focus",
        "emotional_response": "an emotional beat",
        "language_processing": "a compelling message or audio cue",
        "motor_planning": "an action trigger",
    }.get(peak_driver, "a peak in brain activity")

    insights.append({
        "type": "peak", "severity": "positive",
        "title": f"Peak at {peak_s}s — Driven by {peak_label.title()}",
        "message": f"The strongest brain response occurs at {peak_s}s, primarily driven by {peak_label}. This is the most impactful moment in the ad.",
        "action": f"Align {offer} with this moment. If the {cta_label} isn't near {peak_s}s, consider moving it closer."
    })

    # ── 5. CTA Window ──
    if n >= 5:
        last_3 = float(np.mean(engagement[-3:]))
        cta_vs_avg = (last_3 - mean_eng) / (mean_eng + 1e-8) * 100
        if cta_vs_avg < -15:
            insights.append({
                "type": "cta", "severity": "high",
                "title": f"{cta_label} Lands on Low Engagement ({cta_vs_avg:.0f}%)",
                "message": f"The final 3 seconds are {abs(cta_vs_avg):.0f}% below average. Viewers' attention has faded when the {cta_label} appears.",
                "action": f"Option A: Move the {cta_label} to the peak moment ({peak_s}s). Option B: Add a re-engagement spike before the {cta_label} — a bold visual, urgency text, or direct question."
            })
        elif cta_vs_avg > 10:
            insights.append({
                "type": "cta", "severity": "positive",
                "title": f"Strong {cta_label} Window",
                "message": f"Engagement is {cta_vs_avg:.0f}% above average at the end. Viewers are attentive when the {cta_label} appears.",
                "action": f"Make sure the {cta_label} is visually prominent and uses action language. This audience is primed to act."
            })

    # ── 6. Region Dominance (only if there's a clear leader) ──
    if len(sorted_regions) >= 2:
        top_name, top_stats = sorted_regions[0]
        second_name, second_stats = sorted_regions[1]
        dominance = (top_stats["mean"] - second_stats["mean"]) / (second_stats["mean"] + 1e-8) * 100

        if dominance > 10:
            labels = {
                "visual_processing": ("Visually-Driven", "imagery and motion", "Test different visual treatments — color, pacing, composition — since the audience responds primarily to what they see."),
                "attention": ("Cognitively Demanding", "focused thinking", "Balance information density with breathing room. Too much cognitive load causes fatigue and drop-off."),
                "emotional_response": ("Emotion-Led", "feelings over logic", "Use emotion-aligned CTAs ('Don't miss out', 'Join the movement') — they outperform rational CTAs for this type of creative."),
                "language_processing": ("Copy-Driven", "words and audio", "The copy is doing the heavy lifting. A/B test different scripts, voiceover styles, or text overlays."),
                "motor_planning": ("Action-Primed", "impulse to act", "This is ideal for performance ads. Place the CTA at the moment of highest motor activation for maximum conversion."),
            }
            if top_name in labels:
                label, driver, action = labels[top_name]
                insights.append({
                    "type": "dominance", "severity": "positive",
                    "title": f"{label} Creative ({dominance:.0f}% stronger than #{2})",
                    "message": f"This ad's engagement is primarily driven by {driver}. The {labels.get(top_name, ('','',''))[0].lower()} signal is {dominance:.0f}% stronger than the next closest region.",
                    "action": action,
                })

    return insights


MAX_DURATION = 60  # Analyze up to 60s of video
# VJEPA2 consumes 64 frames per 4 s clip: 16 fps is its native rate. At the old
# 2 fps every frame the model saw was one of 8 duplicated eight times, a domain
# shift from the real video it was trained on. Decoding is cached, so the extra
# frames cost almost nothing.
TRIM_FPS = 16


def trim_video(input_path: str, max_seconds: int = MAX_DURATION) -> str:
    """Trim video to max_seconds and reduce FPS for faster encoding."""
    output_path = input_path.replace(".mp4", "_trimmed.mp4")
    try:
        result = subprocess.run(
            ["ffmpeg", "-y", "-i", input_path, "-t", str(max_seconds),
             "-c:v", "libx264", "-preset", "ultrafast", "-crf", "28",
             "-vf", f"fps={TRIM_FPS},scale='min(320,iw)':-2",  # native VJEPA2 rate, 320px width
             "-c:a", "aac", "-b:a", "64k", "-ar", "16000", output_path],
            capture_output=True, text=True, timeout=30,
        )
        if result.returncode == 0 and Path(output_path).exists():
            size = Path(output_path).stat().st_size
            print(f"Trimmed to {max_seconds}s @{TRIM_FPS}fps 320px: {size} bytes")
            return output_path
    except Exception as e:
        print(f"Trim failed ({e}), using original")
    return input_path


# One analysis at a time on the GPU, whether it came from the UI/sync API or
# from the async queue below. Threads waiting here are the queue.
ANALYSIS_LOCK = threading.Lock()


def analyze_video(video_url: str, filename: str = "video.mp4", ad_context_json: str = "{}") -> str:
    """Main analysis — uses full TRIBE v2 pipeline with WhisperX on GPU. Serialized."""
    with ANALYSIS_LOCK:
        return _analyze_video_locked(video_url, filename, ad_context_json)


def _analyze_video_locked(video_url: str, filename: str = "video.mp4", ad_context_json: str = "{}") -> str:
    try:
        # Parse ad context for tailored insights
        try:
            ad_context = json.loads(ad_context_json) if ad_context_json else {}
        except Exception:
            ad_context = {}

        load_model()
        start_time = time.time()
        # HRF offset: TRIBE v2 internally accounts for ~5s hemodynamic delay,
        # but we report seconds as video time (0-based) so users can match
        # insights to actual moments in their video.
        hrf_offset = 0

        # Download video
        print(f"Downloading: {video_url[:100]}")
        with tempfile.NamedTemporaryFile(suffix=".mp4", delete=False) as f:
            resp = requests.get(video_url, timeout=120)
            resp.raise_for_status()
            f.write(resp.content)
            video_path = f.name
        print(f"Downloaded {len(resp.content)} bytes")

        # Trim to MAX_DURATION seconds and downscale for faster encoding
        video_path = trim_video(video_path)

        # Build events — try full pipeline (WhisperX + audio + video), fallback to audio-only
        try:
            print("Building events with full pipeline (WhisperX)...")
            # Keyword argument: the library accepts exactly one of text/audio/video.
            # Passing three positionals set all three and it refused, which is why
            # every analysis before 2026-09-24 silently fell back to audio-only.
            df = model.get_events_dataframe(video_path=video_path)
            n_words = int((df["type"] == "Word").sum()) if "type" in df else 0
            print(f"Events: {len(df)} rows, {n_words} words (full pipeline with transcription)")
        except Exception as whisperx_err:
            print(f"WhisperX pipeline failed, falling back to audio-only: {type(whisperx_err).__name__}: {whisperx_err}")
            traceback.print_exc()
            df = build_events_without_whisperx(video_path)
            print(f"Events: {len(df)} rows (audio-only fallback)")

        # Run brain prediction
        print("Running brain prediction...")
        preds, segments = model.predict(events=df)
        print(f"Predictions: {preds.shape}")

        # Compute engagement
        region_activations = compute_regions(preds)
        engagement = compute_engagement(region_activations)
        mean_eng = float(np.mean(engagement))
        overall_score = compute_score(engagement)

        # Build timeline (exclude last 2 points — attention accumulation artifact)
        valid_len = max(1, len(engagement) - 2)
        timeline = [
            {
                "second": t + hrf_offset,
                "engagement": round(float(engagement[t]), 3),
                "regions": {r: round(float(region_activations[r][t]), 4)
                            for r in FUNCTIONAL_REGIONS if r in region_activations},
            }
            for t in range(valid_len)
        ]

        # Region summary with artifact-corrected peaks
        region_summary = find_peak_region_seconds(region_activations, hrf_offset)

        # Find peak/lowest in the valid range (excluding tail artifact)
        valid_engagement = engagement[:valid_len]
        peak_idx = int(np.argmax(valid_engagement))
        low_idx = int(np.argmin(valid_engagement))

        insights = generate_insights(valid_engagement, region_activations, hrf_offset, ad_context)
        elapsed = time.time() - start_time

        report = {
            "pipeline": "full" if ("type" in df and (df["type"] == "Word").any()) else "audio_only",
            "filename": filename,
            "duration_seconds": valid_len + hrf_offset,
            "analyzed_seconds": valid_len,
            "overall_score": overall_score,
            "mean_engagement": round(mean_eng, 3),
            "peak": {
                "value": round(float(valid_engagement[peak_idx]), 3),
                "second": peak_idx + hrf_offset,
            },
            "lowest": {
                "value": round(float(valid_engagement[low_idx]), 3),
                "second": low_idx + hrf_offset,
            },
            "region_summary": region_summary,
            "timeline": timeline,
            "insights": insights,
            "inference_time_seconds": round(elapsed, 1),
            "model": "facebook/tribev2",
        }
        Path(video_path).unlink(missing_ok=True)
        print(f"Analysis complete: score={overall_score}, {len(engagement)} timesteps, {elapsed:.1f}s")
        return json.dumps(report)

    except Exception as e:
        error_msg = f"Error: {type(e).__name__}: {str(e)}"
        print(error_msg)
        traceback.print_exc()
        return json.dumps({"error": error_msg, "traceback": traceback.format_exc()})


# ---------------------------------------------------------------------------
# Async API for AdVisor: accept a job, return at once, analyse in the
# background, POST the report to the caller's callback URL.
#
# Only hosts listed in ADVISOR_CALLBACK_HOSTS (comma-separated, default: the
# AdVisor Supabase project) may receive callbacks, so this public endpoint
# cannot be pointed at arbitrary servers. The caller supplies job_ref and
# job_token; both are echoed back and the caller verifies the token, so no
# shared secret has to live in this Space.
# ---------------------------------------------------------------------------
ALLOWED_CALLBACK_HOSTS = {
    h.strip().lower()
    for h in os.environ.get("ADVISOR_CALLBACK_HOSTS", "emyiuzsgvnnzgtqwmttb.supabase.co").split(",")
    if h.strip()
}
CALLBACK_RETRIES = 4


def _post_callback(callback_url: str, payload: dict) -> bool:
    delay = 5
    for attempt in range(1, CALLBACK_RETRIES + 1):
        try:
            resp = requests.post(
                callback_url,
                json=payload,
                headers={"Content-Type": "application/json", "X-Job-Token": payload.get("job_token", "")},
                timeout=60,
            )
            if resp.ok:
                print(f"[async] callback delivered for {payload.get('job_ref')} ({resp.status_code})")
                return True
            print(f"[async] callback HTTP {resp.status_code} for {payload.get('job_ref')}: {resp.text[:200]}")
            if 400 <= resp.status_code < 500 and resp.status_code != 429:
                return False  # our payload is wrong; retrying will not help
        except Exception as e:  # network
            print(f"[async] callback attempt {attempt} failed: {e}")
        time.sleep(delay)
        delay = min(delay * 2, 60)
    return False


def _run_async_job(video_url: str, filename: str, ad_context_json: str, callback_url: str, job_ref: str, job_token: str):
    started = time.time()
    try:
        raw = analyze_video(video_url, filename, ad_context_json)
        report = json.loads(raw)
        if isinstance(report, dict) and report.get("error"):
            payload = {"job_ref": job_ref, "job_token": job_token, "status": "failed", "error": str(report.get("error"))[:1000]}
        else:
            payload = {"job_ref": job_ref, "job_token": job_token, "status": "complete", "report": report}
    except Exception as e:  # never let a job die silently
        payload = {"job_ref": job_ref, "job_token": job_token, "status": "failed", "error": f"{type(e).__name__}: {e}"[:1000]}
    payload["elapsed_seconds"] = round(time.time() - started, 1)
    _post_callback(callback_url, payload)


def analyze_video_async(video_url: str, filename: str, ad_context_json: str, callback_url: str, job_ref: str, job_token: str) -> str:
    """Queue an analysis and return immediately. The report arrives at callback_url."""
    try:
        host = (urlparse(callback_url).hostname or "").lower()
        if urlparse(callback_url).scheme != "https" or host not in ALLOWED_CALLBACK_HOSTS:
            return json.dumps({"accepted": False, "error": f"callback host not allowed: {host or '(none)'}"})
        if not video_url or not job_ref or not job_token:
            return json.dumps({"accepted": False, "error": "video_url, job_ref and job_token are required"})
        t = threading.Thread(
            target=_run_async_job,
            args=(video_url, filename or "video.mp4", ad_context_json or "{}", callback_url, job_ref, job_token),
            daemon=True,
            name=f"tribe-async-{job_ref[:12]}",
        )
        t.start()
        waiting = sum(1 for th in threading.enumerate() if th.name.startswith("tribe-async-"))
        print(f"[async] accepted job {job_ref} ({waiting} in flight)")
        return json.dumps({"accepted": True, "job_ref": job_ref, "in_flight": waiting})
    except Exception as e:
        return json.dumps({"accepted": False, "error": f"{type(e).__name__}: {e}"})


app = gr.Blocks(title="TRIBE v2 Neural Engagement Analyzer")
with app:
    gr.Markdown("# TRIBE v2 Neural Engagement Analyzer")
    gr.Markdown("Predict brain response to video content, second by second.")
    with gr.Row():
        video_url_input = gr.Textbox(label="Video URL", placeholder="https://example.com/video.mp4")
        filename_input = gr.Textbox(label="Filename", value="video.mp4")
    ad_context_input = gr.Textbox(label="Ad Context (JSON)", value="{}", visible=False)
    analyze_btn = gr.Button("Analyze", variant="primary")
    output = gr.JSON(label="Neural Engagement Report")
    analyze_btn.click(fn=analyze_video, inputs=[video_url_input, filename_input, ad_context_input], outputs=output)

    # Hidden controls back the async API (gradio_api/call/analyze_video_async).
    with gr.Row(visible=False):
        async_video_url = gr.Textbox(label="video_url")
        async_filename = gr.Textbox(label="filename", value="video.mp4")
        async_ad_context = gr.Textbox(label="ad_context_json", value="{}")
        async_callback_url = gr.Textbox(label="callback_url")
        async_job_ref = gr.Textbox(label="job_ref")
        async_job_token = gr.Textbox(label="job_token")
        async_btn = gr.Button("Queue")
        async_out = gr.Textbox(label="ack")
    async_btn.click(
        fn=analyze_video_async,
        inputs=[async_video_url, async_filename, async_ad_context, async_callback_url, async_job_ref, async_job_token],
        outputs=async_out,
        api_name="analyze_video_async",
    )
app.queue()
if __name__ == "__main__":
    app.launch()