import os
import time
import tempfile
import traceback
import subprocess
from concurrent.futures import ThreadPoolExecutor, as_completed

import requests
import gradio as gr
from openai import OpenAI

try:
    import spaces
    _HAS_SPACES_GPU = True
except ImportError:
    _HAS_SPACES_GPU = False

from video_captioner import VideoCaptioner
from caption_generator import (
    generate_and_judge_caption,
    format_scene_data,
    STYLE_PROMPTS,
)

# --------------------------------------------------------------------------- #
# Config (models are fine as plain env vars; SECRETS -- the API keys -- must
# be set as HF Space "Secrets", never committed to a file).
# --------------------------------------------------------------------------- #
SCENE_MODEL = os.environ.get("SCENE_MODEL", "accounts/fireworks/models/minimax-m3")
CAPTION_MODEL = os.environ.get("CAPTION_MODEL", "accounts/fireworks/models/minimax-m3")
JUDGE_MODEL = os.environ.get("JUDGE_MODEL", CAPTION_MODEL)
REFINE_MODEL = os.environ.get("REFINE_MODEL", CAPTION_MODEL)

MAX_RETRIES = int(os.environ.get("MAX_RETRIES", "3"))
ENABLE_JUDGE = os.environ.get("ENABLE_JUDGE", "true").lower() in ("1", "true", "yes")
MAX_REFINE_ITERATIONS = int(os.environ.get("MAX_REFINE_ITERATIONS", "1"))
MAX_DURATION_SECONDS = int(os.environ.get("MAX_DURATION_SECONDS", "120"))

STYLE_LABELS = {
    "formal": "📄 Formal",
    "sarcastic": "😏 Sarcastic",
    "humorous_tech": "💻 Humorous (Tech)",
    "humorous_non_tech": "😂 Humorous (Non-Tech)",
}

FIREWORKS_API_KEY = os.environ.get("FIREWORKS_API_KEY")
OPENROUTER_API_KEY = os.environ.get("OPENROUTER_API_KEY")

_client = None
_captioner = None
_init_error = None

try:
    if not FIREWORKS_API_KEY:
        raise RuntimeError(
            "FIREWORKS_API_KEY is not set. Add it under Space Settings -> "
            "Variables and secrets -> New secret."
        )
    _client = OpenAI(
        api_key=FIREWORKS_API_KEY,
        base_url="https://api.fireworks.ai/inference/v1",
    )
    _captioner = VideoCaptioner(
        api_key=FIREWORKS_API_KEY,
        model=SCENE_MODEL,
        max_retries=MAX_RETRIES,
        audio_api=OPENROUTER_API_KEY,
    )
except Exception as e:
    _init_error = str(e)


# --------------------------------------------------------------------------- #
# Helpers
# --------------------------------------------------------------------------- #
def download_video(url: str, dest_dir: str) -> str:
    local_path = os.path.join(dest_dir, "clip.mp4")
    with requests.get(url, stream=True, timeout=120) as r:
        r.raise_for_status()
        with open(local_path, "wb") as f:
            for chunk in r.iter_content(chunk_size=1024 * 1024):
                if chunk:
                    f.write(chunk)
    return local_path


def get_duration_seconds(path: str) -> float:
    cmd = [
        "ffprobe", "-v", "error",
        "-show_entries", "format=duration",
        "-of", "default=noprint_wrappers=1:nokey=1",
        path,
    ]
    out = subprocess.run(cmd, capture_output=True, text=True, check=True)
    return float(out.stdout.strip())


EMPTY = ("", "", "", "")


if _HAS_SPACES_GPU:
    @spaces.GPU(duration=5)
    def _zerogpu_touch():
        """No-op. ZeroGPU Spaces require at least one @spaces.GPU-decorated
        function to be detected at startup, even though this app's actual
        work (video download, transcription, captioning) is all done via
        remote API calls and never touches the local GPU."""
        return True
else:
    def _zerogpu_touch():
        return True


def run_pipeline(video_url: str, progress=gr.Progress()):
    """Generator: yields (status_markdown, formal, sarcastic, tech, non_tech)
    repeatedly so the UI updates live instead of showing a blank screen."""

    if _init_error:
        yield f"❌ App is misconfigured: {_init_error}", *EMPTY
        return

    if not video_url or not video_url.strip():
        yield "⚠️ Please paste a direct video URL first.", *EMPTY
        return

    log = []

    def status(msg: str) -> str:
        log.append(msg)
        return "\n\n".join(log)

    results = {}

    try:
        _zerogpu_touch()
        with tempfile.TemporaryDirectory() as tmp_dir:
            progress(0.03, desc="Downloading video")
            yield status("📥 Downloading video..."), *EMPTY

            video_path = download_video(video_url.strip(), tmp_dir)

            duration = get_duration_seconds(video_path)
            if duration > MAX_DURATION_SECONDS:
                yield status(
                    f"❌ This clip is {duration:.0f}s long. Please use a video "
                    f"under {MAX_DURATION_SECONDS}s (2 minutes)."
                ), *EMPTY
                return

            progress(0.15, desc="Extracting frames & audio")
            yield status(
                f"✅ Downloaded ({duration:.0f}s clip)\n\n"
                "🎞️ Extracting frames and transcribing audio..."
            ), *EMPTY

            with ThreadPoolExecutor(max_workers=2) as ex:
                transcript_future = ex.submit(_captioner.get_transcript, video_path)
                frames_future = ex.submit(_captioner.extract_frames, video_path)
                transcript = transcript_future.result()
                frames_b64 = frames_future.result()

            progress(0.3, desc="Analyzing scene")
            yield status(
                f"✅ Extracted {len(frames_b64)} frames"
                + (", transcribed audio" if transcript else ", no speech detected")
                + "\n\n🧠 Analyzing the scene with the vision model..."
            ), *EMPTY

            messages = _captioner.build_messages(transcript, frames_b64)
            scene_json = _captioner.get_scene_json(messages)
            scene_json["audio_transcript"] = transcript
            scene_text = format_scene_data(scene_json)

            progress(0.4, desc="Generating captions")
            yield status(
                f"✅ Scene analyzed: {scene_json.get('scene', '') or 'done'}\n\n"
                "✍️ Writing captions in 4 styles (generate → judge → refine)..."
            ), *EMPTY

            styles = list(STYLE_PROMPTS.keys())
            done = 0
            with ThreadPoolExecutor(max_workers=len(styles)) as ex:
                future_to_style = {
                    ex.submit(
                        generate_and_judge_caption,
                        style, scene_text,
                        _client, CAPTION_MODEL,
                        _client, JUDGE_MODEL,
                        _client, REFINE_MODEL,
                        MAX_RETRIES, ENABLE_JUDGE, MAX_REFINE_ITERATIONS,
                    ): style
                    for style in styles
                }
                for future in as_completed(future_to_style):
                    style = future_to_style[future]
                    try:
                        outcome = future.result()
                        caption = outcome["caption"]
                    except Exception as e:
                        caption = f"[Fallback] {style.replace('_', ' ').title()} caption."
                        print(f"[{style}] failed: {e}")

                    results[style] = caption
                    done += 1
                    progress(0.4 + 0.55 * (done / len(styles)), desc=f"{style} done")
                    yield (
                        status(f"✅ {STYLE_LABELS[style]} caption ready ({done}/{len(styles)})"),
                        results.get("formal", ""),
                        results.get("sarcastic", ""),
                        results.get("humorous_tech", ""),
                        results.get("humorous_non_tech", ""),
                    )

            progress(1.0, desc="Done")
            yield (
                status("🏁 All captions ready!"),
                results.get("formal", ""),
                results.get("sarcastic", ""),
                results.get("humorous_tech", ""),
                results.get("humorous_non_tech", ""),
            )

    except requests.exceptions.RequestException as e:
        yield status(f"❌ Could not download that video URL: {e}"), *EMPTY
    except subprocess.CalledProcessError:
        yield status("❌ Could not read that file as a video (ffprobe failed)."), *EMPTY
    except Exception as e:
        traceback.print_exc()
        yield status(f"❌ Something went wrong: {e}"), *EMPTY


# --------------------------------------------------------------------------- #
# UI
# --------------------------------------------------------------------------- #
with gr.Blocks(title="AI Video Captioner", theme=gr.themes.Soft()) as demo:
    gr.Markdown(
        "# 🎬 AI Video Captioner\n"
        "Paste a **direct link** to a video under 2 minutes (a `.mp4` URL, etc.) "
        "and get an AI-generated caption in 4 different styles: formal, sarcastic, "
        "tech-humor, and everyday humor. Each caption is graded and refined by an "
        "LLM judge before it's shown."
    )

    with gr.Row():
        video_url = gr.Textbox(
            label="Video URL",
            placeholder="https://storage.googleapis.com/amd-hackathon-clips/1860079-uhd_2560_1440_25fps.mp4",
            scale=4,
        )
        run_btn = gr.Button("Generate Captions", variant="primary", scale=1)

    status_box = gr.Markdown("Paste a video URL above and click **Generate Captions**.")

    with gr.Row():
        formal_out = gr.Textbox(label=STYLE_LABELS["formal"], lines=4, interactive=False)
        sarcastic_out = gr.Textbox(label=STYLE_LABELS["sarcastic"], lines=4, interactive=False)
    with gr.Row():
        tech_out = gr.Textbox(label=STYLE_LABELS["humorous_tech"], lines=4, interactive=False)
        nontech_out = gr.Textbox(label=STYLE_LABELS["humorous_non_tech"], lines=4, interactive=False)

    run_btn.click(
        fn=run_pipeline,
        inputs=[video_url],
        outputs=[status_box, formal_out, sarcastic_out, tech_out, nontech_out],
    )
    video_url.submit(
        fn=run_pipeline,
        inputs=[video_url],
        outputs=[status_box, formal_out, sarcastic_out, tech_out, nontech_out],
    )

demo.queue(max_size=10)

if __name__ == "__main__":
    demo.launch(server_name="0.0.0.0", server_port=7860, share=False)