"""
pipeline.py — PDF → Gemini → Manim render → return artifacts
"""

from __future__ import annotations

import subprocess
import textwrap
import time
import zipfile
from pathlib import Path
from typing import Callable

from queue_manager import State
from utils import extract_pdf_text, generate_manim_code, sanitize_manim_code

MEDIA_ROOT   = Path("media/videos")
JOBS_ROOT    = Path("jobs")
ARTIFACTS    = Path("artifacts")


def run_pipeline(
    job_id: str,
    pdf_path: str,
    gemini_api_key: str,
    status_cb: Callable[[State, str, str | None], None],
) -> dict:
    """
    Run the full pipeline and return:
        {
            "video_path": str,   # absolute path to the rendered .mp4
            "zip_path":   str,   # absolute path to artifacts_<job_id>.zip
            "code":       str,   # generated Manim source
        }
    """
    job_dir     = JOBS_ROOT / job_id
    script_path = job_dir / f"{job_id}.py"
    video_path  = MEDIA_ROOT / job_id / "720p30" / "OutputVideo.mp4"
    zip_path    = ARTIFACTS / f"artifacts_{job_id}.zip"

    job_dir.mkdir(parents=True, exist_ok=True)
    ARTIFACTS.mkdir(parents=True, exist_ok=True)

    # 1. Extract PDF text
    status_cb(State.RUNNING, "📖 Extracting text from PDF…", None)
    pdf_text = extract_pdf_text(pdf_path)
    if not pdf_text.strip():
        raise ValueError("PDF appears empty or has no selectable text.")

    # 2. Generate Manim code via Gemini
    status_cb(State.RUNNING, "🤖 Generating Manim code with Gemini…", None)
    prompt     = _build_prompt(pdf_text)
    raw_code   = generate_manim_code(prompt, gemini_api_key)
    manim_code = sanitize_manim_code(raw_code)

    status_cb(State.RUNNING, "✏️ Code generated — starting render…", manim_code)
    script_path.write_text(manim_code, encoding="utf-8")

    # 3. Render
    status_cb(State.RUNNING, "🎬 Rendering animation (this may take a few minutes)…", None)
    _render_manim(script_path)

    if not video_path.exists():
        raise FileNotFoundError(f"Rendered video not found at {video_path}")

    # 4. Package artifacts zip
    status_cb(State.RUNNING, "📦 Packaging artifacts zip…", None)
    with zipfile.ZipFile(zip_path, "w", zipfile.ZIP_DEFLATED) as zf:
        zf.write(script_path, arcname=f"{job_id}.py")
        zf.write(video_path,  arcname="OutputVideo.mp4")

    return {
        "video_path": str(video_path.resolve()),
        "zip_path":   str(zip_path.resolve()),
        "code":       manim_code,
    }


# ── Helpers ───────────────────────────────────────────────────────────────────

def _build_prompt(pdf_text: str) -> str:
    truncated = pdf_text[:12_000]
    return textwrap.dedent(f"""
        You are an expert Manim animator. Given the following document content,
        create a concise, visually engaging Manim animation that summarises the
        key ideas. Use ONLY the class name `OutputVideo` extending `Scene`.

        Requirements:
        - Class name MUST be exactly `OutputVideo(Scene)`.
        - Use only standard Manim Community v0.18+ API.
        - Output ONLY valid Python code. No explanations, no markdown fences.
        - Keep runtime under 90 seconds.
        - Avoid custom LaTeX preambles; prefer Text() over MathTex() where possible.
        - Animations should be clear, readable, and professional.

        Document content:
        ---
        {truncated}
        ---
    """).strip()


def _render_manim(script_path: Path, max_retries: int = 2) -> None:
    cmd = [
        "manim",
        str(script_path),
        "OutputVideo",
        "-qm",
        "--media_dir", "media",
        "--disable_caching",
    ]
    stderr_tail = ""
    for attempt in range(max_retries + 1):
        result = subprocess.run(cmd, capture_output=True, text=True, timeout=600)
        if result.returncode == 0:
            return
        stderr_tail = result.stderr[-3000:]
        if attempt < max_retries:
            time.sleep(5 * (attempt + 1))

    raise RuntimeError(
        f"Manim render failed after {max_retries + 1} attempts.\nSTDERR: {stderr_tail}"
    )