"""Multi-agent pipeline for ManimGL code generation."""

import json
import logging
from dataclasses import dataclass, field
from typing import Any

from .providers import LLMProvider, get_provider

logger = logging.getLogger(__name__)


@dataclass
class AgentResult:
    """Result from an agent."""

    agent: str
    success: bool
    output: Any
    reasoning: str = ""
    errors: list[str] = field(default_factory=list)


class BaseAgent:
    """Base class for pipeline agents."""

    name: str = "base"
    system_prompt: str = ""

    def __init__(self, provider: LLMProvider | None = None):
        self.provider = provider or get_provider()

    async def run(
        self,
        input_data: Any,
        context: dict[str, Any] | None = None,  # noqa: ARG002
    ) -> AgentResult:
        """Run the agent. Override in subclasses."""
        raise NotImplementedError


class ConceptAnalyzer(BaseAgent):
    """Analyzes the user's concept and extracts key elements."""

    name = "concept_analyzer"
    system_prompt = """You are an expert at analyzing mathematical and scientific concepts for visualization.

Given a user's request for a ManimGL animation, analyze:
1. Core concepts to visualize
2. Mathematical objects needed (shapes, graphs, equations)
3. Animations that would best illustrate the concept
4. Suggested visual flow/narrative

Output JSON with this structure:
{
    "concepts": ["list of core concepts"],
    "objects": [{"type": "Circle/Graph/Tex/etc", "description": "what it represents"}],
    "animations": [{"type": "Transform/FadeIn/etc", "purpose": "why this animation"}],
    "narrative": ["step 1", "step 2", ...],
    "complexity": "simple|medium|complex",
    "duration_estimate": "estimated seconds"
}"""

    async def run(
        self,
        input_data: str,
        context: dict[str, Any] | None = None,  # noqa: ARG002
    ) -> AgentResult:
        """Analyze the concept from user request."""
        messages = [
            {"role": "system", "content": self.system_prompt},
            {"role": "user", "content": f"Analyze this animation request:\n\n{input_data}"},
        ]

        try:
            response = await self.provider.complete(messages, temperature=0.5)

            # Extract JSON from response
            try:
                # Try to parse as JSON directly
                if "```json" in response:
                    json_str = response.split("```json")[1].split("```")[0]
                elif "```" in response:
                    json_str = response.split("```")[1].split("```")[0]
                else:
                    json_str = response

                analysis = json.loads(json_str.strip())
            except json.JSONDecodeError:
                analysis = {"raw_analysis": response}

            return AgentResult(
                agent=self.name,
                success=True,
                output=analysis,
                reasoning="Analyzed concept for key visualization elements",
            )
        except Exception as e:
            return AgentResult(
                agent=self.name,
                success=False,
                output=None,
                errors=[str(e)],
            )


class ScenePlanner(BaseAgent):
    """Plans the scene structure with multiple animations."""

    name = "scene_planner"
    system_prompt = """You are an expert ManimGL scene architect.

Given a concept analysis, create a detailed scene plan that uses 2+ animation patterns.

Output JSON with this structure:
{
    "scene_name": "DescriptiveSceneName",
    "imports": ["from manimlib import *"],
    "sections": [
        {
            "name": "section_name",
            "description": "what happens",
            "objects_created": ["obj1", "obj2"],
            "animations": [
                {"type": "ShowCreation", "target": "obj1", "duration": 1},
                {"type": "Transform", "source": "obj1", "target": "obj2", "duration": 2}
            ]
        }
    ],
    "total_duration": 15,
    "animation_count": 5
}

Ensure:
- At least 2 different animation types are used
- Objects are created before being animated
- Proper pacing with wait() calls
- Professional quality transitions"""

    async def run(
        self,
        input_data: dict[str, Any],
        context: dict[str, Any] | None = None,  # noqa: ARG002
    ) -> AgentResult:
        """Plan the scene structure."""
        messages = [
            {"role": "system", "content": self.system_prompt},
            {
                "role": "user",
                "content": f"Create a scene plan for:\n\n{json.dumps(input_data, indent=2)}",
            },
        ]

        try:
            response = await self.provider.complete(messages, temperature=0.5)

            # Extract JSON
            try:
                if "```json" in response:
                    json_str = response.split("```json")[1].split("```")[0]
                elif "```" in response:
                    json_str = response.split("```")[1].split("```")[0]
                else:
                    json_str = response

                plan = json.loads(json_str.strip())
            except json.JSONDecodeError:
                plan = {"raw_plan": response}

            return AgentResult(
                agent=self.name,
                success=True,
                output=plan,
                reasoning=f"Planned scene with {plan.get('animation_count', '?')} animations",
            )
        except Exception as e:
            return AgentResult(
                agent=self.name,
                success=False,
                output=None,
                errors=[str(e)],
            )


class CodeGenerator(BaseAgent):
    """Generates ManimGL code from scene plan."""

    name = "code_generator"
    system_prompt = """You are an expert ManimGL programmer.

Generate complete, working ManimGL code from a scene plan.

Rules:
1. Always start with: from manimlib import *
2. Create a Scene class with construct() method
3. Use proper ManimGL syntax (not manim community edition)
4. Include comments explaining each section
5. Use self.play() for animations
6. Use self.wait() for pauses
7. Handle 3D scenes with ThreeDScene base class

Output ONLY the Python code, no explanations.
Wrap code in ```python ... ```"""

    def __init__(self, provider: LLMProvider | None = None, examples: list[str] | None = None):
        super().__init__(provider)
        self.examples = examples or []

    async def run(
        self,
        input_data: dict[str, Any],
        context: dict[str, Any] | None = None,  # noqa: ARG002
    ) -> AgentResult:
        """Generate code from plan."""
        # Build prompt with examples
        prompt = f"Generate ManimGL code for this scene plan:\n\n{json.dumps(input_data, indent=2)}"

        if self.examples:
            example_text = "\n\n---\n\n".join(self.examples[:3])
            prompt = f"Reference examples:\n\n{example_text}\n\n---\n\n{prompt}"

        messages = [
            {"role": "system", "content": self.system_prompt},
            {"role": "user", "content": prompt},
        ]

        try:
            response = await self.provider.complete(messages, temperature=0.3)

            # Extract code
            if "```python" in response:
                code = response.split("```python")[1].split("```")[0]
            elif "```" in response:
                code = response.split("```")[1].split("```")[0]
            else:
                code = response

            code = code.strip()

            # Basic validation
            if "class" not in code or "def construct" not in code:
                return AgentResult(
                    agent=self.name,
                    success=False,
                    output=code,
                    errors=["Generated code missing Scene class or construct method"],
                )

            return AgentResult(
                agent=self.name,
                success=True,
                output=code,
                reasoning="Generated ManimGL scene code",
            )
        except Exception as e:
            return AgentResult(
                agent=self.name,
                success=False,
                output=None,
                errors=[str(e)],
            )


class CodeReviewer(BaseAgent):
    """Reviews and improves generated code."""

    name = "code_reviewer"
    system_prompt = """You are a ManimGL code reviewer and fixer.

Review the code for:
1. Syntax errors
2. ManimGL API correctness (not manim community edition)
3. Object references (objects used before creation)
4. Animation compatibility (correct animation for object type)
5. Best practices

If issues found, output corrected code.
If code is good, output the same code with any minor improvements.

Output JSON:
{
    "issues_found": ["list of issues or empty"],
    "improvements": ["list of improvements made"],
    "code": "corrected/improved code"
}"""

    async def run(
        self,
        input_data: str,
        context: dict[str, Any] | None = None,  # noqa: ARG002
    ) -> AgentResult:
        """Review and improve code."""
        messages = [
            {"role": "system", "content": self.system_prompt},
            {
                "role": "user",
                "content": f"Review this ManimGL code:\n\n```python\n{input_data}\n```",
            },
        ]

        try:
            response = await self.provider.complete(messages, temperature=0.3)

            # Extract JSON
            try:
                if "```json" in response:
                    json_str = response.split("```json")[1].split("```")[0]
                elif "```" in response:
                    # Might be the code directly
                    json_str = response.split("```")[1].split("```")[0]
                    if "issues_found" not in json_str:
                        # It's code, not JSON
                        return AgentResult(
                            agent=self.name,
                            success=True,
                            output={
                                "code": json_str.strip(),
                                "issues_found": [],
                                "improvements": [],
                            },
                            reasoning="Code reviewed and returned",
                        )
                else:
                    json_str = response

                review = json.loads(json_str.strip())
            except json.JSONDecodeError:
                # Assume the response is the corrected code
                review = {
                    "code": input_data,
                    "issues_found": [],
                    "improvements": ["Unable to parse review response"],
                }

            return AgentResult(
                agent=self.name,
                success=True,
                output=review,
                reasoning=f"Found {len(review.get('issues_found', []))} issues",
            )
        except Exception as e:
            return AgentResult(
                agent=self.name,
                success=False,
                output={"code": input_data},
                errors=[str(e)],
            )


async def run_multi_agent_pipeline(
    user_request: str,
    examples: list[str] | None = None,
    provider: LLMProvider | None = None,
) -> dict[str, Any]:
    """Run the full multi-agent code generation pipeline.

    Args:
        user_request: User's description of the animation they want
        examples: Optional list of example code for context
        provider: Optional LLM provider

    Returns:
        Dict with pipeline results including final code
    """
    provider = provider or get_provider()
    results = {
        "request": user_request,
        "stages": [],
        "success": False,
        "code": None,
        "errors": [],
    }

    # Stage 1: Concept Analysis
    logger.info("Stage 1: Analyzing concept...")
    analyzer = ConceptAnalyzer(provider)
    analysis_result = await analyzer.run(user_request)
    results["stages"].append(analysis_result)

    if not analysis_result.success:
        results["errors"].extend(analysis_result.errors)
        return results

    # Stage 2: Scene Planning
    logger.info("Stage 2: Planning scene...")
    planner = ScenePlanner(provider)
    plan_result = await planner.run(analysis_result.output)
    results["stages"].append(plan_result)

    if not plan_result.success:
        results["errors"].extend(plan_result.errors)
        return results

    # Stage 3: Code Generation
    logger.info("Stage 3: Generating code...")
    generator = CodeGenerator(provider, examples=examples)
    code_result = await generator.run(plan_result.output)
    results["stages"].append(code_result)

    if not code_result.success:
        results["errors"].extend(code_result.errors)
        return results

    # Stage 4: Code Review
    logger.info("Stage 4: Reviewing code...")
    reviewer = CodeReviewer(provider)
    review_result = await reviewer.run(code_result.output)
    results["stages"].append(review_result)

    # Extract final code
    if review_result.success and review_result.output:
        final_code = review_result.output.get("code", code_result.output)
        results["code"] = final_code
        results["success"] = True
        results["issues_fixed"] = review_result.output.get("issues_found", [])
        results["improvements"] = review_result.output.get("improvements", [])
    else:
        # Fall back to unreviewed code
        results["code"] = code_result.output
        results["success"] = True
        results["errors"].append("Code review failed, using unreviewed code")

    return results
