"""split command: split a monolithic scene.py into a multi-file project."""

from __future__ import annotations

import ast
import re
import stat
import textwrap
from pathlib import Path
from typing import Optional

import click


def _camel_to_snake(name: str) -> str:
    """Convert CamelCase to snake_case.

    Parameters
    ----------
    name : str
        CamelCase string (e.g. "WhatIsLLMScene").

    Returns
    -------
    str
        snake_case string (e.g. "what_is_llm_scene").
    """
    # Insert underscore before uppercase letters that follow lowercase/digits
    s = re.sub(r"([a-z0-9])([A-Z])", r"\1_\2", name)
    # Insert underscore between consecutive uppercase and a following lowercase
    s = re.sub(r"([A-Z]+)([A-Z][a-z])", r"\1_\2", s)
    return s.lower()


def _class_to_filename(class_name: str, index: int) -> str:
    """Derive a scene filename from a class name and sequence index.

    Parameters
    ----------
    class_name : str
        The Scene subclass name (e.g. "HookScene").
    index : int
        1-based scene index for the filename prefix.

    Returns
    -------
    str
        Filename like "scene_01_hook.py".
    """
    snake = _camel_to_snake(class_name)
    # Strip trailing "_scene" suffix if present
    if snake.endswith("_scene"):
        snake = snake[: -len("_scene")]
    return f"scene_{index:02d}_{snake}.py"


def _is_scene_class(node: ast.ClassDef) -> bool:
    """Check whether an AST ClassDef inherits from a Scene-like base.

    Parameters
    ----------
    node : ast.ClassDef
        The class definition node.

    Returns
    -------
    bool
        True if any base class name contains "Scene".
    """
    for base in node.bases:
        if isinstance(base, ast.Name) and "Scene" in base.id:
            return True
        if isinstance(base, ast.Attribute) and "Scene" in base.attr:
            return True
    return False


def _extract_source_lines(
    source_lines: list[str],
    node: ast.AST,
) -> str:
    """Extract the source text for an AST node using line numbers.

    Parameters
    ----------
    source_lines : list[str]
        All lines of the original source file (0-indexed).
    node : ast.AST
        An AST node with lineno and end_lineno attributes.

    Returns
    -------
    str
        The source text for that node.
    """
    start = node.lineno - 1  # type: ignore[attr-defined]
    end = node.end_lineno  # type: ignore[attr-defined]
    return "\n".join(source_lines[start:end])


def _build_scene_imports() -> str:
    """Return the standard import header for a split scene file.

    Returns
    -------
    str
        Python import block that sets up sys.path and imports style.
    """
    return textwrap.dedent("""\
        import sys
        from pathlib import Path

        sys.path.insert(0, str(Path(__file__).resolve().parent))

        from utils.style import *  # noqa: F403
    """)


def _generate_render_script(
    scene_files: list[tuple[str, str]],
    quality: str = "l",
) -> str:
    """Generate a render_all.sh script for all scenes.

    Parameters
    ----------
    scene_files : list[tuple[str, str]]
        List of (filename, class_name) pairs.
    quality : str
        Default render quality flag.

    Returns
    -------
    str
        Shell script content.
    """
    lines = [
        "#!/usr/bin/env bash",
        "# Auto-generated by: 3brown1blue split",
        'set -euo pipefail',
        "",
        f'QUALITY="${{1:-{quality}}}"',
        "",
        "# Render each scene",
    ]
    for filename, class_name in scene_files:
        lines.append(f'manim -q"$QUALITY" {filename} {class_name}')
    lines.append("")
    lines.append("echo 'All scenes rendered.'")
    lines.append("")
    return "\n".join(lines)


def _parse_source(source: str) -> dict:
    """Parse a Python source file into categorized AST components.

    Parameters
    ----------
    source : str
        The full source code of the monolithic scene file.

    Returns
    -------
    dict
        Keys: "imports" (str), "style_nodes" (list of AST nodes),
        "helper_nodes" (list of AST nodes), "scene_nodes"
        (list of ast.ClassDef).

    Raises
    ------
    click.ClickException
        If the source cannot be parsed.
    """
    try:
        tree = ast.parse(source)
    except SyntaxError as exc:
        raise click.ClickException(
            f"Failed to parse source file: {exc}"
        ) from exc

    source_lines = source.splitlines()
    imports: list[str] = []
    style_nodes: list[ast.AST] = []
    helper_nodes: list[ast.AST] = []
    scene_nodes: list[ast.ClassDef] = []

    for node in ast.iter_child_nodes(tree):
        # Import statements
        if isinstance(node, (ast.Import, ast.ImportFrom)):
            imports.append(_extract_source_lines(source_lines, node))

        # Top-level assignments (style constants, dicts, etc.)
        elif isinstance(node, ast.Assign):
            style_nodes.append(node)

        # Augmented assignments (e.g. COLORS |= {...})
        elif isinstance(node, ast.AugAssign):
            style_nodes.append(node)

        # Annotated assignments (e.g. TITLE_SIZE: int = 42)
        elif isinstance(node, ast.AnnAssign):
            style_nodes.append(node)

        # Top-level function definitions (helpers)
        elif isinstance(node, ast.FunctionDef):
            helper_nodes.append(node)

        elif isinstance(node, ast.AsyncFunctionDef):
            helper_nodes.append(node)

        # Class definitions: separate Scene subclasses from utility classes
        elif isinstance(node, ast.ClassDef):
            if _is_scene_class(node):
                scene_nodes.append(node)
            else:
                # Non-Scene classes go into style/utils
                style_nodes.append(node)

        # Expression statements (docstrings, standalone comments as strings)
        elif isinstance(node, ast.Expr):
            # Module-level docstring or constant expression -- skip
            pass

    return {
        "imports": "\n".join(imports),
        "style_nodes": style_nodes,
        "helper_nodes": helper_nodes,
        "scene_nodes": scene_nodes,
    }


@click.command()
@click.argument("source", type=click.Path(exists=True, path_type=Path))
@click.option("--output-dir", "-o", default=None,
              type=click.Path(path_type=Path),
              help="Output directory (default: <source_stem>_project/).")
def split(source: Path, output_dir: Optional[Path]) -> None:
    """Split a monolithic scene.py into a multi-file project.

    SOURCE is a single Python file containing multiple Scene classes
    (typically the output of `3brown1blue generate`).

    \b
    Examples:
      3brown1blue split scene.py                    # creates ./scene_project/
      3brown1blue split scene.py -o my_video/       # custom output dir
    """
    source = source.resolve()
    source_text = source.read_text()
    source_lines = source_text.splitlines()

    if output_dir is None:
        output_dir = source.parent / f"{source.stem}_project"
    output_dir = output_dir.resolve()

    parsed = _parse_source(source_text)

    if not parsed["scene_nodes"]:
        raise click.ClickException(
            f"No Scene subclasses found in {source.name}. "
            "The file must contain at least one class inheriting from Scene."
        )

    # Create directory structure
    output_dir.mkdir(parents=True, exist_ok=True)
    utils_dir = output_dir / "utils"
    utils_dir.mkdir(exist_ok=True)

    # Write utils/__init__.py
    init_file = utils_dir / "__init__.py"
    init_file.write_text("")

    # Build style.py content
    style_parts: list[str] = ["from manim import *  # noqa: F403", ""]

    for node in parsed["style_nodes"]:
        style_parts.append(_extract_source_lines(source_lines, node))
        style_parts.append("")

    for node in parsed["helper_nodes"]:
        style_parts.append("")
        style_parts.append(_extract_source_lines(source_lines, node))
        style_parts.append("")

    style_file = utils_dir / "style.py"
    style_file.write_text("\n".join(style_parts).rstrip() + "\n")

    # Write individual scene files
    scene_files: list[tuple[str, str]] = []
    created_files: list[str] = []

    for i, scene_node in enumerate(parsed["scene_nodes"], 1):
        class_name = scene_node.name
        filename = _class_to_filename(class_name, i)
        scene_files.append((filename, class_name))

        scene_source = _extract_source_lines(source_lines, scene_node)
        file_content = _build_scene_imports() + "\n" + scene_source + "\n"

        scene_path = output_dir / filename
        scene_path.write_text(file_content)
        created_files.append(filename)

    # Generate render_all.sh
    render_script = output_dir / "render_all.sh"
    render_script.write_text(_generate_render_script(scene_files))
    # Make executable
    render_script.chmod(
        render_script.stat().st_mode | stat.S_IXUSR | stat.S_IXGRP
    )

    # Print summary
    click.echo(f"\nSplit {source.name} into {len(scene_files)} scenes:")
    click.echo(f"  Output: {output_dir}/\n")
    click.echo(f"  {'File':<35} {'Class':<30}")
    click.echo("  " + "-" * 65)
    click.echo(f"  {'utils/style.py':<35} {'(constants + helpers)':<30}")
    for filename, class_name in scene_files:
        click.echo(f"  {filename:<35} {class_name:<30}")
    click.echo(f"  {'render_all.sh':<35} {'(render script)':<30}")
    click.echo(f"\n  Total files: {len(scene_files) + 3}")
    click.echo(f"\nTo render:  cd {output_dir.name} && bash render_all.sh")
