from __future__ import annotations

import os
import tempfile
import threading
import typing
from pathlib import Path

import gradio as gr
import numpy as np
import spaces
from typing_extensions import Self

# ZeroGPU currently offers a Python 3.10 runtime. Some TRIBE dependencies
# resolve postponed annotations through ``typing.Self`` (added in 3.11), so
# expose the standard typing_extensions backport before importing TRIBE code.
if not hasattr(typing, "Self"):
    typing.Self = Self  # type: ignore[attr-defined]

from brain_analysis import (
    align_events,
    analyze_predictions,
    build_summary,
    build_timeline_figure,
    load_events,
)
from tribe_runtime import predict_video


MAX_VIDEO_BYTES = int(os.getenv("MAX_VIDEO_BYTES", str(500 * 1024 * 1024)))
INFERENCE_LOCK = threading.Lock()


@spaces.GPU(duration=44)
def analyze_gameplay(video_path: str | None, events_path: str | None, progress=gr.Progress()):
    if not video_path:
        raise gr.Error("Upload a gameplay video first.")
    video = Path(video_path)
    if not video.is_file():
        raise gr.Error("The uploaded video is no longer available. Please upload it again.")
    if video.stat().st_size > MAX_VIDEO_BYTES:
        raise gr.Error("Video exceeds the 500 MB Space limit. Trim it to the scene you want to test.")

    try:
        events = load_events(events_path)
        with INFERENCE_LOCK:
            predictions, segments = predict_video(str(video), progress=progress)
        result = analyze_predictions(predictions, segments)
        timeline = align_events(result.timeline, events)
        figure = build_timeline_figure(timeline, events)
        summary = build_summary(timeline, events)

        output_dir = Path(tempfile.mkdtemp(prefix="clawbrain-"))
        csv_path = output_dir / "clawbrain_timeline.csv"
        npz_path = output_dir / "tribev2_cortical_predictions.npz"
        timeline.to_csv(csv_path, index=False)
        np.savez_compressed(
            npz_path,
            predictions=result.predictions,
            time_s=timeline["time_s"].to_numpy(),
        )
        progress(1.0, desc="Analysis ready")
        return summary, figure, timeline, [str(csv_path), str(npz_path)]
    except gr.Error:
        raise
    except Exception as exc:
        message = str(exc)
        if "gated" in message.lower() or "401" in message or "403" in message:
            raise gr.Error(
                "Model access was denied. Add HF_TOKEN as a Space secret and accept access "
                "for meta-llama/Llama-3.2-3B on Hugging Face."
            ) from exc
        raise gr.Error(f"Analysis failed: {message[:500]}") from exc


CSS = """
@import url('https://fonts.googleapis.com/css2?family=Azeret+Mono:wght@400;500;600&family=Syne:wght@600;700&display=swap');
:root { --ink:#0b0d0c; --paper:#d8d4c7; --acid:#d9ff43; --signal:#ff6b4a; --muted:#8d8c85; }
.gradio-container { background:var(--ink)!important; color:var(--paper)!important; font-family:'Azeret Mono',monospace!important; }
.gradio-container:before { content:''; position:fixed; inset:0; pointer-events:none; opacity:.19; background-image:linear-gradient(rgba(255,255,255,.025) 1px,transparent 1px),linear-gradient(90deg,rgba(255,255,255,.025) 1px,transparent 1px); background-size:32px 32px; }
.main-shell { max-width:1380px; margin:0 auto; padding:22px 18px 48px; }
.masthead { display:grid; grid-template-columns:1fr auto; gap:24px; align-items:end; border-top:1px solid #555750; border-bottom:1px solid #555750; padding:22px 0 18px; margin-bottom:18px; }
.eyebrow { color:var(--acid); letter-spacing:.18em; font-size:11px; }
.masthead h1 { margin:5px 0 0!important; color:var(--paper)!important; font:700 clamp(42px,7vw,90px)/.88 'Syne',sans-serif!important; letter-spacing:-.06em; }
.masthead p { max-width:450px; margin:0; color:#aaa89f; font-size:12px; line-height:1.65; }
.status-chip { display:inline-flex; align-items:center; gap:8px; margin-top:12px; color:#c7c5bd; font-size:10px; letter-spacing:.08em; }
.status-chip:before { content:''; width:7px; height:7px; border-radius:50%; background:var(--acid); box-shadow:0 0 16px var(--acid); animation:pulse 2s infinite; }
@keyframes pulse { 50% { opacity:.35; } }
.panel { background:rgba(16,18,17,.92)!important; border:1px solid #373a36!important; border-radius:2px!important; box-shadow:none!important; }
.panel-label { color:var(--muted); font-size:10px; letter-spacing:.14em; margin-bottom:8px; }
.primary-action { background:var(--acid)!important; color:#11130f!important; border:0!important; border-radius:1px!important; font-weight:600!important; letter-spacing:.08em!important; min-height:48px!important; }
.primary-action:hover { filter:brightness(.88); transform:translateY(-1px); }
.readout-grid { display:grid; grid-template-columns:1fr 1fr 2fr; border:1px solid #373a36; background:#0f1110; }
.readout { padding:17px; border-right:1px solid #373a36; }
.readout:last-child { border-right:0; }
.readout span,.readout small { display:block; color:#777a74; font-size:9px; letter-spacing:.12em; }
.readout strong { display:block; margin:8px 0 5px; color:var(--acid); font:600 28px/1 'Syne',sans-serif; }
.readout.wide strong { color:var(--paper); font-size:22px; }
.interpretation-note { margin-top:10px; border-left:2px solid var(--signal); padding:10px 13px; background:#171512; color:#aaa89f; font-size:10px; line-height:1.55; }
.interpretation-note b { color:var(--signal); }
.gradio-container label,.gradio-container .label-wrap { color:#aaa89f!important; font-size:10px!important; letter-spacing:.08em; }
.footer-note { color:#6f716c; font-size:9px; line-height:1.6; border-top:1px solid #373a36; padding-top:14px; margin-top:20px; }
@media(max-width:760px){ .masthead{grid-template-columns:1fr}.readout-grid{grid-template-columns:1fr 1fr}.readout.wide{grid-column:1/-1;border-top:1px solid #373a36}.readout:nth-child(2){border-right:0} }
"""


with gr.Blocks(css=CSS, title="CLAW/BRAIN · TRIBE v2") as demo:
    with gr.Column(elem_classes="main-shell"):
        gr.HTML(
            """
            <header class="masthead">
              <div><div class="eyebrow">CLAWBUSTER · CONTENT RESPONSE LAB</div><h1>CLAW/BRAIN</h1></div>
              <div><p>Video + audio cortical-response simulation for gameplay concepts, captures, and finished game scenes. Powered by Meta TRIBE v2.</p><div class="status-chip">AVERAGE-SUBJECT MODEL · RESEARCH MODE</div></div>
            </header>
            """
        )
        with gr.Row(equal_height=False):
            with gr.Column(scale=5, elem_classes="panel"):
                gr.Markdown("### 01 / STIMULUS")
                video = gr.Video(label="GAMEPLAY VIDEO", sources=["upload"], format="mp4")
            with gr.Column(scale=3, elem_classes="panel"):
                gr.Markdown("### 02 / EVENT LAYER")
                events = gr.File(
                    label="OPTIONAL EVENTS CSV",
                    file_types=[".csv"],
                    type="filepath",
                )
                gr.Markdown(
                    "CSV columns: `time_s,event,category`  \n"
                    "Example: `12.4,near miss,reward`"
                )
                run = gr.Button("RUN CORTICAL SIMULATION →", elem_classes="primary-action")
                gr.Markdown(
                    "First run downloads the TRIBE checkpoint and video/audio encoders. "
                    "GPU hardware is required for practical inference."
                )
        gr.HTML('<div class="panel-label">03 / RESPONSE READOUT</div>')
        summary = gr.HTML()
        chart = gr.Plot(show_label=False)
        with gr.Accordion("INSPECT / EXPORT MODEL OUTPUT", open=False):
            table = gr.Dataframe(label="TIMELINE", interactive=False, wrap=True)
            downloads = gr.File(label="DOWNLOADS", file_count="multiple")
        gr.HTML(
            """
            <div class="footer-note">TRIBE v2 is licensed CC BY-NC 4.0. This Space is a research prototype and is not a medical device. Commercial deployment requires appropriate model licensing, participant consent for physiological datasets, and validation of every behavioral label.</div>
            """
        )

    run.click(
        fn=analyze_gameplay,
        inputs=[video, events],
        outputs=[summary, chart, table, downloads],
        concurrency_limit=1,
    )


if __name__ == "__main__":
    demo.queue(default_concurrency_limit=1, max_size=8).launch()