"""LLM-powered Manim scene generation from a topic description."""

import os
import re
import subprocess
import sys
from pathlib import Path
from typing import Optional

import click

SKILL_SOURCE = Path(__file__).parent / "skill"

PROVIDERS = {
    "anthropic": {
        "env_var": "ANTHROPIC_API_KEY",
        "default_model": "claude-sonnet-4-5",
        "label": "Anthropic (Claude)",
    },
    "openai": {
        "env_var": "OPENAI_API_KEY",
        "default_model": "gpt-4o",
        "label": "OpenAI",
        "base_url": None,
    },
    "google": {
        "env_var": "GOOGLE_API_KEY",
        "default_model": "gemini-2.0-flash",
        "label": "Google Gemini",
        "base_url": "https://generativelanguage.googleapis.com/v1beta/openai/",
    },
    "groq": {
        "env_var": "GROQ_API_KEY",
        "default_model": "llama-3.3-70b-versatile",
        "label": "Groq",
        "base_url": "https://api.groq.com/openai/v1",
    },
    "mistral": {
        "env_var": "MISTRAL_API_KEY",
        "default_model": "mistral-large-latest",
        "label": "Mistral",
        "base_url": "https://api.mistral.ai/v1",
    },
}

CORE_RULES = [
    "troubleshooting.md",
    "scene-planning.md",
    "visual-design-principles.md",
    "equations.md",
    "animations.md",
]

MAX_TOKENS = 8192


def _build_system_prompt() -> str:
    """Build the LLM system prompt from bundled skill files."""
    parts: list[str] = []

    skill_md = SKILL_SOURCE / "SKILL.md"
    if skill_md.exists():
        parts.append(skill_md.read_text())

    rules_dir = SKILL_SOURCE / "rules"
    for name in CORE_RULES:
        f = rules_dir / name
        if f.exists():
            heading = f.stem.replace("-", " ").title()
            parts.append(f"\n\n## {heading}\n\n{f.read_text()}")

    parts.append(
        "\n\n---\n\n"
        "Return ONLY a complete, runnable Python file. "
        "Do not include any explanation outside the code. "
        "The file must define at least one Scene subclass. "
        "Use descriptive comments inside the code."
    )
    return "\n".join(parts)


def _extract_code(response: str) -> str:
    """Extract Python code from an LLM response."""
    match = re.search(r"```(?:python)?\n(.*?)```", response, re.DOTALL)
    if match:
        return match.group(1).strip()
    return response.strip()


def _resolve_api_key(provider: str, api_key: Optional[str]) -> str:
    """Resolve API key from argument or environment variable."""
    key = (api_key or "").strip() or os.environ.get(PROVIDERS[provider]["env_var"], "").strip()
    if not key:
        env_var = PROVIDERS[provider]["env_var"]
        click.echo(f"No API key found. Set {env_var} or pass --api-key.")
        sys.exit(1)
    return key


def _call_anthropic(topic: str, model: str, api_key: str) -> str:
    try:
        import anthropic
    except ImportError:
        click.echo("Run: pip install anthropic  (or: pip install '3brown1blue[anthropic]')")
        sys.exit(1)

    try:
        client = anthropic.Anthropic(api_key=api_key, timeout=60.0)
        msg = client.messages.create(
            model=model,
            max_tokens=MAX_TOKENS,
            system=_build_system_prompt(),
            messages=[{"role": "user", "content": f"Create a Manim scene that explains: {topic}"}],
        )
    except Exception as e:
        click.echo(f"Anthropic API error: {e}")
        sys.exit(1)

    if msg.stop_reason == "max_tokens":
        click.echo("Warning: response was cut off (max tokens reached). Output may be incomplete.")

    return msg.content[0].text


def _call_openai_compatible(
    topic: str, model: str, api_key: str, base_url: Optional[str]
) -> str:
    try:
        from openai import OpenAI
    except ImportError:
        click.echo("Run: pip install openai  (or: pip install '3brown1blue[openai]')")
        sys.exit(1)

    kwargs: dict = {"api_key": api_key, "timeout": 60.0}
    if base_url:
        kwargs["base_url"] = base_url

    try:
        client = OpenAI(**kwargs)
        response = client.chat.completions.create(
            model=model,
            messages=[
                {"role": "system", "content": _build_system_prompt()},
                {"role": "user", "content": f"Create a Manim scene that explains: {topic}"},
            ],
            max_tokens=MAX_TOKENS,
        )
    except Exception as e:
        click.echo(f"API error: {e}")
        sys.exit(1)

    choice = response.choices[0]
    if choice.finish_reason == "length":
        click.echo("Warning: response was cut off (max tokens reached). Output may be incomplete.")

    content = choice.message.content
    if not content:
        click.echo("Provider returned empty content. Try again or check your API quota.")
        sys.exit(1)

    return content


def _generate_scene(topic: str, provider: str, model: str, api_key: str) -> str:
    if provider == "anthropic":
        return _call_anthropic(topic, model, api_key)
    cfg = PROVIDERS[provider]
    return _call_openai_compatible(topic, model, api_key, cfg.get("base_url"))


def _render_scene(scene_file: Path, quality: str) -> None:
    """Detect the last Scene subclass in the file and run manim."""
    source = scene_file.read_text()
    matches = re.findall(r"^class\s+(\w+)\s*\(.*Scene.*\)", source, re.MULTILINE)
    if not matches:
        click.echo("Could not detect a Scene class -- skipping render.")
        click.echo(f"Run manually: manim -pq{quality} {scene_file} <SceneName>")
        return

    scene_name = matches[-1]
    cmd = ["manim", f"-pq{quality}", str(scene_file), scene_name]
    click.echo(f"Rendering:  {' '.join(cmd)}\n")
    result = subprocess.run(cmd, check=False)
    if result.returncode != 0:
        click.echo(f"\nRender failed (exit {result.returncode}).")
        sys.exit(result.returncode)


# ── Click command ──────────────────────────────────────────────────────────

@click.command()
@click.argument("topic")
@click.option(
    "--provider", "-p",
    type=click.Choice(list(PROVIDERS)),
    default=None,
    help="LLM provider to use.",
)
@click.option("--model", "-m", default=None, help="Model name (uses provider default).")
@click.option("--api-key", "-k", default=None, help="API key (or set the provider env var).")
@click.option("--output", "-o", default="scene.py", show_default=True, help="Output file.")
@click.option("--render", is_flag=True, help="Run manim on the generated file.")
@click.option(
    "--quality", "-q",
    type=click.Choice(["l", "m", "h", "k"]),
    default="l",
    show_default=True,
    help="Render quality: l=low/fast, m=medium, h=1080p, k=4K.",
)
def generate(
    topic: str,
    provider: Optional[str],
    model: Optional[str],
    api_key: Optional[str],
    output: str,
    render: bool,
    quality: str,
) -> None:
    """Generate a Manim scene that explains TOPIC using an LLM.

    \b
    Examples:
      3brown1blue generate "backpropagation"
      3brown1blue generate "Fourier transform" --provider openai --render
      3brown1blue generate "attention mechanism" -p anthropic -o attention.py --render -q h
    """
    if provider is None:
        click.echo("\nWhich LLM provider?\n")
        items = list(PROVIDERS.items())
        for i, (key, cfg) in enumerate(items, 1):
            click.echo(f"  {i}. {cfg['label']}")
        click.echo()
        choice = click.prompt("Provider", type=click.IntRange(1, len(items)))
        provider = items[choice - 1][0]

    cfg = PROVIDERS[provider]
    model = model or cfg["default_model"]
    api_key = _resolve_api_key(provider, api_key)

    click.echo(f"\nGenerating: \"{topic}\"")
    click.echo(f"Provider:   {cfg['label']}  ({model})")

    raw = _generate_scene(topic, provider, model, api_key)
    code = _extract_code(raw)

    out_path = Path(output)
    out_path.write_text(code)
    click.echo(f"Saved:      {out_path}")

    if render:
        _render_scene(out_path, quality)
