"""Indexer for ManimGL source code and examples."""

import ast
import logging
from dataclasses import dataclass, field
from pathlib import Path
from typing import Any

from .store import RAGStore, get_store

logger = logging.getLogger(__name__)


@dataclass
class APISignature:
    """Extracted API signature."""

    name: str
    module: str
    signature: str
    docstring: str
    parameters: list[dict[str, Any]] = field(default_factory=list)
    category: str = "unknown"  # mobject, animation, scene, util


@dataclass
class SceneExample:
    """Extracted scene example."""

    name: str
    source: str
    code: str
    description: str
    animations_used: list[str] = field(default_factory=list)
    mobjects_used: list[str] = field(default_factory=list)


@dataclass
class AnimationPattern:
    """Animation pattern template."""

    name: str
    description: str
    code: str
    category: str  # transform, creation, indication, movement, etc.
    tags: list[str] = field(default_factory=list)


class ManimIndexer:
    """Index ManimGL source code for RAG retrieval."""

    # Known animation classes for categorization
    ANIMATION_CLASSES = {
        "Transform",
        "ReplacementTransform",
        "TransformFromCopy",
        "MoveToTarget",
        "ApplyMethod",
        "FadeIn",
        "FadeOut",
        "FadeTransform",
        "ShowCreation",
        "Uncreate",
        "Write",
        "DrawBorderThenFill",
        "GrowFromCenter",
        "GrowFromPoint",
        "Indicate",
        "Flash",
        "CircleIndicate",
        "WiggleOutThenIn",
        "Rotating",
        "Rotate",
        "AnimationGroup",
        "Succession",
        "LaggedStart",
        "LaggedStartMap",
        "UpdateFromFunc",
        "UpdateFromAlphaFunc",
        "MoveAlongPath",
        "Homotopy",
        "ChangingDecimal",
        "ChangeDecimalToValue",
        "CountInFrom",
    }

    MOBJECT_CLASSES = {
        "Circle",
        "Square",
        "Rectangle",
        "Triangle",
        "Polygon",
        "RegularPolygon",
        "Line",
        "Arrow",
        "DoubleArrow",
        "Vector",
        "Arc",
        "Dot",
        "Ellipse",
        "Text",
        "Tex",
        "TexText",
        "MathTex",
        "Title",
        "BulletedList",
        "NumberLine",
        "Axes",
        "ThreeDAxes",
        "NumberPlane",
        "CoordinateSystem",
        "Graph",
        "ParametricCurve",
        "FunctionGraph",
        "Sphere",
        "Cube",
        "Cylinder",
        "Cone",
        "Prism",
        "Surface",
        "VGroup",
        "Group",
        "VMobject",
        "Mobject",
        "ValueTracker",
        "ComplexValueTracker",
        "Matrix",
        "IntegerMatrix",
        "DecimalMatrix",
        "SVGMobject",
        "ImageMobject",
    }

    def __init__(self, store: RAGStore | None = None):
        """Initialize indexer.

        Args:
            store: RAG store instance, or uses global store
        """
        self.store = store or get_store()
        self.signatures: list[APISignature] = []
        self.examples: list[SceneExample] = []
        self.patterns: list[AnimationPattern] = []

    def index_manimgl_source(self, source_path: Path) -> dict[str, int]:
        """Index the manimgl source directory.

        Args:
            source_path: Path to manimlib/ directory

        Returns:
            Stats dict with counts
        """
        logger.info(f"Indexing manimgl source: {source_path}")

        stats = {
            "files_processed": 0,
            "signatures_extracted": 0,
            "examples_extracted": 0,
            "errors": 0,
        }

        # Index all Python files
        for py_file in source_path.rglob("*.py"):
            if "__pycache__" in str(py_file):
                continue
            try:
                self._index_file(py_file)
                stats["files_processed"] += 1
            except Exception as e:
                logger.warning(f"Error indexing {py_file}: {e}")
                stats["errors"] += 1

        stats["signatures_extracted"] = len(self.signatures)
        stats["examples_extracted"] = len(self.examples)

        return stats

    def _index_file(self, filepath: Path) -> None:
        """Index a single Python file."""
        code = filepath.read_text(encoding="utf-8", errors="ignore")

        try:
            tree = ast.parse(code)
        except SyntaxError:
            logger.warning(f"Syntax error in {filepath}")
            return

        module_name = self._get_module_name(filepath)

        for node in ast.walk(tree):
            if isinstance(node, ast.ClassDef):
                self._extract_class(node, module_name, code)
            elif isinstance(node, ast.FunctionDef) and node.col_offset == 0:
                # Only top-level functions
                self._extract_function(node, module_name, code)

    def _get_module_name(self, filepath: Path) -> str:
        """Extract module name from file path."""
        parts = filepath.parts
        try:
            manimlib_idx = parts.index("manimlib")
            module_parts = parts[manimlib_idx:]
            module = ".".join(module_parts)
            return module.replace(".py", "")
        except ValueError:
            return filepath.stem

    def _extract_class(self, node: ast.ClassDef, module: str, code: str) -> None:
        """Extract class as API signature."""
        docstring = ast.get_docstring(node) or ""

        # Determine category
        category = "unknown"
        bases = [self._get_base_name(b) for b in node.bases]

        if any(b in self.ANIMATION_CLASSES or "Animation" in b for b in bases):
            category = "animation"
        elif any(b in self.MOBJECT_CLASSES or "Mobject" in b or "VMobject" in b for b in bases):
            category = "mobject"
        elif any("Scene" in b for b in bases):
            category = "scene"
            # Also extract as example if it's a Scene
            self._extract_scene_example(node, module, code)

        # Extract __init__ signature
        init_params = []
        for item in node.body:
            if isinstance(item, ast.FunctionDef) and item.name == "__init__":
                init_params = self._extract_params(item)
                break

        signature = f"class {node.name}({', '.join(bases)})"
        if init_params:
            param_str = ", ".join(p["name"] for p in init_params if p["name"] != "self")
            signature = f"{node.name}({param_str})"

        self.signatures.append(
            APISignature(
                name=node.name,
                module=module,
                signature=signature,
                docstring=docstring[:500],  # Truncate long docstrings
                parameters=init_params,
                category=category,
            )
        )

    def _extract_function(
        self,
        node: ast.FunctionDef,
        module: str,
        code: str,  # noqa: ARG002
    ) -> None:
        """Extract function as API signature."""
        docstring = ast.get_docstring(node) or ""
        params = self._extract_params(node)

        param_str = ", ".join(p["name"] for p in params)
        signature = f"{node.name}({param_str})"

        self.signatures.append(
            APISignature(
                name=node.name,
                module=module,
                signature=signature,
                docstring=docstring[:500],
                parameters=params,
                category="util",
            )
        )

    def _extract_params(self, node: ast.FunctionDef) -> list[dict[str, Any]]:
        """Extract function parameters."""
        params = []
        args = node.args

        # Get defaults starting positions
        num_defaults = len(args.defaults)
        default_offset = len(args.args) - num_defaults

        for i, arg in enumerate(args.args):
            param = {
                "name": arg.arg,
                "type": None,
                "default": None,
            }

            # Type annotation
            if arg.annotation:
                param["type"] = ast.unparse(arg.annotation)

            # Default value
            if i >= default_offset:
                default_node = args.defaults[i - default_offset]
                try:
                    param["default"] = ast.unparse(default_node)
                except Exception:
                    param["default"] = "..."

            params.append(param)

        # **kwargs
        if args.kwarg:
            params.append({"name": f"**{args.kwarg.arg}", "type": None, "default": None})

        return params

    def _extract_scene_example(self, node: ast.ClassDef, module: str, code: str) -> None:
        """Extract Scene class as example."""
        # Get the full class source
        try:
            start = node.lineno - 1
            end = node.end_lineno or start + 1
            lines = code.split("\n")
            class_code = "\n".join(lines[start:end])
        except Exception:
            class_code = ""

        docstring = ast.get_docstring(node) or ""

        # Find animations and mobjects used
        animations_used = set()
        mobjects_used = set()

        for child in ast.walk(node):
            if isinstance(child, ast.Call):
                func_name = self._get_call_name(child)
                if func_name in self.ANIMATION_CLASSES:
                    animations_used.add(func_name)
                elif func_name in self.MOBJECT_CLASSES:
                    mobjects_used.add(func_name)

        self.examples.append(
            SceneExample(
                name=node.name,
                source=module,
                code=class_code,
                description=docstring or f"Scene class: {node.name}",
                animations_used=list(animations_used),
                mobjects_used=list(mobjects_used),
            )
        )

    def _get_base_name(self, base: ast.expr) -> str:
        """Get base class name from AST node."""
        if isinstance(base, ast.Name):
            return base.id
        elif isinstance(base, ast.Attribute):
            return base.attr
        return ""

    def _get_call_name(self, call: ast.Call) -> str:
        """Get function/class name from Call node."""
        if isinstance(call.func, ast.Name):
            return call.func.id
        elif isinstance(call.func, ast.Attribute):
            return call.func.attr
        return ""

    def add_animation_patterns(self) -> int:
        """Add built-in animation pattern templates."""
        patterns = [
            AnimationPattern(
                name="riemann_sum",
                description="Animated Riemann sum approximation of an integral",
                code="""class RiemannSumDemo(Scene):
    def construct(self):
        axes = Axes(x_range=[-1, 5], y_range=[-1, 5])
        func = axes.get_graph(lambda x: 0.5 * x**2, color=BLUE)

        # Create rectangles for Riemann sum
        rects = axes.get_riemann_rectangles(
            func, x_range=[0, 4], dx=0.5,
            stroke_width=1, fill_opacity=0.5
        )

        self.play(ShowCreation(axes), ShowCreation(func))
        self.play(ShowCreation(rects))

        # Animate refining the approximation
        for dx in [0.25, 0.1, 0.05]:
            new_rects = axes.get_riemann_rectangles(
                func, x_range=[0, 4], dx=dx,
                stroke_width=1, fill_opacity=0.5
            )
            self.play(Transform(rects, new_rects))
        self.wait()""",
                category="calculus",
                tags=["integral", "riemann", "approximation", "calculus"],
            ),
            AnimationPattern(
                name="matrix_transform",
                description="Visualize how a matrix transforms the plane",
                code="""class MatrixTransformDemo(Scene):
    def construct(self):
        plane = NumberPlane()
        vectors = VGroup(
            Vector([1, 0], color=RED),
            Vector([0, 1], color=GREEN),
        )

        self.play(ShowCreation(plane), ShowCreation(vectors))

        # Define transformation matrix
        matrix = [[2, 1], [1, 2]]

        # Apply transformation
        self.play(
            plane.animate.apply_matrix(matrix),
            vectors[0].animate.apply_matrix(matrix),
            vectors[1].animate.apply_matrix(matrix),
            run_time=2
        )
        self.wait()""",
                category="linear_algebra",
                tags=["matrix", "transform", "linear", "vectors"],
            ),
            AnimationPattern(
                name="fourier_series",
                description="Animated Fourier series approximation",
                code="""class FourierSeriesDemo(Scene):
    def construct(self):
        axes = Axes(x_range=[-PI, PI], y_range=[-2, 2])

        # Target function (square wave)
        def square_wave(x):
            return 1 if x > 0 else -1

        # Fourier approximation
        def fourier_approx(x, n_terms):
            result = 0
            for n in range(1, n_terms + 1, 2):
                result += (4 / (n * PI)) * np.sin(n * x)
            return result

        target = axes.get_graph(square_wave, discontinuities=[0])
        self.play(ShowCreation(axes), ShowCreation(target))

        # Show successive approximations
        approx = axes.get_graph(lambda x: fourier_approx(x, 1), color=YELLOW)
        self.play(ShowCreation(approx))

        for n in [3, 5, 9, 15, 31]:
            new_approx = axes.get_graph(lambda x: fourier_approx(x, n), color=YELLOW)
            self.play(Transform(approx, new_approx))
        self.wait()""",
                category="calculus",
                tags=["fourier", "series", "approximation", "wave"],
            ),
            AnimationPattern(
                name="derivative_tangent",
                description="Show derivative as tangent line slope",
                code="""class DerivativeTangentDemo(Scene):
    def construct(self):
        axes = Axes(x_range=[-3, 3], y_range=[-1, 9])
        func = axes.get_graph(lambda x: x**2, color=BLUE)

        self.play(ShowCreation(axes), ShowCreation(func))

        # Animated tangent line
        x_tracker = ValueTracker(-2)

        def get_tangent_line():
            x = x_tracker.get_value()
            slope = 2 * x  # derivative of x^2
            point = axes.c2p(x, x**2)
            line = axes.get_tangent_line(x, func, length=4, color=YELLOW)
            return line

        tangent = always_redraw(get_tangent_line)
        dot = always_redraw(lambda: Dot(axes.c2p(x_tracker.get_value(), x_tracker.get_value()**2)))

        self.play(ShowCreation(tangent), ShowCreation(dot))
        self.play(x_tracker.animate.set_value(2), run_time=4)
        self.wait()""",
                category="calculus",
                tags=["derivative", "tangent", "slope", "calculus"],
            ),
            AnimationPattern(
                name="physics_projectile",
                description="Projectile motion with trajectory",
                code="""class ProjectileMotionDemo(Scene):
    def construct(self):
        axes = Axes(x_range=[0, 10], y_range=[0, 5])

        v0 = 5  # initial velocity
        theta = PI / 4  # launch angle
        g = 9.8

        def trajectory(t):
            x = v0 * np.cos(theta) * t
            y = v0 * np.sin(theta) * t - 0.5 * g * t**2
            return axes.c2p(x, max(0, y))

        ball = Dot(color=RED).move_to(axes.c2p(0, 0))
        path = TracedPath(ball.get_center, stroke_color=YELLOW)

        self.play(ShowCreation(axes))
        self.add(path)
        self.play(
            UpdateFromFunc(ball, lambda m: m.move_to(trajectory(self.time))),
            run_time=2,
            rate_func=linear
        )
        self.wait()""",
                category="physics",
                tags=["projectile", "motion", "physics", "trajectory"],
            ),
            AnimationPattern(
                name="complex_plane",
                description="Operations on the complex plane",
                code="""class ComplexPlaneDemo(Scene):
    def construct(self):
        plane = ComplexPlane()
        plane.add_coordinate_labels()

        z1 = complex(2, 1)
        z2 = complex(-1, 2)

        dot1 = Dot(plane.n2p(z1), color=RED)
        dot2 = Dot(plane.n2p(z2), color=BLUE)

        self.play(ShowCreation(plane))
        self.play(ShowCreation(dot1), ShowCreation(dot2))

        # Show addition
        sum_z = z1 + z2
        dot_sum = Dot(plane.n2p(sum_z), color=GREEN)
        line1 = Line(plane.n2p(0), plane.n2p(z1), color=RED)
        line2 = Line(plane.n2p(z1), plane.n2p(sum_z), color=BLUE)

        self.play(ShowCreation(line1), ShowCreation(line2))
        self.play(ShowCreation(dot_sum))
        self.wait()""",
                category="complex_analysis",
                tags=["complex", "plane", "addition", "vectors"],
            ),
            AnimationPattern(
                name="graph_traversal",
                description="Animated graph traversal (BFS/DFS)",
                code="""class GraphTraversalDemo(Scene):
    def construct(self):
        # Create graph vertices
        vertices = VGroup(*[Dot() for _ in range(6)])
        vertices.arrange_in_grid(2, 3, buff=1.5)

        # Create edges
        edges = VGroup(
            Line(vertices[0].get_center(), vertices[1].get_center()),
            Line(vertices[1].get_center(), vertices[2].get_center()),
            Line(vertices[0].get_center(), vertices[3].get_center()),
            Line(vertices[1].get_center(), vertices[4].get_center()),
            Line(vertices[2].get_center(), vertices[5].get_center()),
        )

        labels = VGroup(*[Text(str(i)).scale(0.5).next_to(v, UP) for i, v in enumerate(vertices)])

        self.play(ShowCreation(vertices), ShowCreation(edges), Write(labels))

        # BFS traversal
        order = [0, 1, 3, 2, 4, 5]
        for i in order:
            self.play(vertices[i].animate.set_color(YELLOW).scale(1.5), run_time=0.5)
            self.play(vertices[i].animate.set_color(GREEN).scale(1/1.5), run_time=0.3)
        self.wait()""",
                category="algorithms",
                tags=["graph", "bfs", "dfs", "traversal", "algorithms"],
            ),
            AnimationPattern(
                name="3d_surface",
                description="3D surface plot with rotation",
                code="""class Surface3DDemo(ThreeDScene):
    def construct(self):
        axes = ThreeDAxes()

        surface = Surface(
            lambda u, v: axes.c2p(u, v, np.sin(u) * np.cos(v)),
            u_range=[-PI, PI],
            v_range=[-PI, PI],
            resolution=(30, 30),
        )
        surface.set_fill_by_value(axes=axes, colors=[(BLUE, -1), (GREEN, 0), (RED, 1)])

        self.set_camera_orientation(phi=60 * DEGREES, theta=-45 * DEGREES)
        self.play(ShowCreation(axes))
        self.play(ShowCreation(surface), run_time=2)

        # Rotate camera
        self.begin_ambient_camera_rotation(rate=0.2)
        self.wait(5)
        self.stop_ambient_camera_rotation()""",
                category="3d",
                tags=["3d", "surface", "plot", "rotation"],
            ),
        ]

        self.patterns.extend(patterns)
        return len(patterns)

    def build_index(self) -> dict[str, int]:
        """Build vector index from collected data."""
        stats = {
            "examples": 0,
            "signatures": 0,
            "patterns": 0,
        }

        # Index examples
        if self.examples:
            example_data = []
            for ex in self.examples:
                text = f"{ex.name}: {ex.description}\nAnimations: {', '.join(ex.animations_used)}\nMobjects: {', '.join(ex.mobjects_used)}"
                example_data.append(
                    {
                        "vector": self.store.embed_single(text),
                        "text": text,
                        "name": ex.name,
                        "code": ex.code,
                        "source": ex.source,
                        "animations": ex.animations_used,
                        "mobjects": ex.mobjects_used,
                    }
                )
            self.store.create_table(self.store.EXAMPLES_TABLE, example_data)
            stats["examples"] = len(example_data)

        # Index API signatures
        if self.signatures:
            sig_data = []
            for sig in self.signatures:
                text = f"{sig.signature}\n{sig.docstring}"
                sig_data.append(
                    {
                        "vector": self.store.embed_single(text),
                        "text": text,
                        "name": sig.name,
                        "module": sig.module,
                        "signature": sig.signature,
                        "docstring": sig.docstring,
                        "category": sig.category,
                        "parameters": str(sig.parameters),
                    }
                )
            self.store.create_table(self.store.API_TABLE, sig_data)
            stats["signatures"] = len(sig_data)

        # Index patterns
        if self.patterns:
            pattern_data = []
            for pat in self.patterns:
                text = f"{pat.name}: {pat.description}\nTags: {', '.join(pat.tags)}"
                pattern_data.append(
                    {
                        "vector": self.store.embed_single(text),
                        "text": text,
                        "name": pat.name,
                        "description": pat.description,
                        "code": pat.code,
                        "category": pat.category,
                        "tags": pat.tags,
                    }
                )
            self.store.create_table(self.store.PATTERNS_TABLE, pattern_data)
            stats["patterns"] = len(pattern_data)

        return stats


async def index_manimgl(source_path: str | Path) -> dict[str, Any]:
    """Index manimgl source and build RAG database.

    Args:
        source_path: Path to manimlib/ source directory

    Returns:
        Stats about indexed content
    """
    source_path = Path(source_path)

    indexer = ManimIndexer()

    # Index source code
    source_stats = indexer.index_manimgl_source(source_path)

    # Add built-in patterns
    pattern_count = indexer.add_animation_patterns()

    # Build vector index
    index_stats = indexer.build_index()

    return {
        "source": source_stats,
        "index": index_stats,
        "patterns_added": pattern_count,
    }
