"""DSPy modules for error analysis and fix generation."""

from __future__ import annotations

from typing import Any, Dict, List, Optional

from autodebug.optimization.signatures import (
    AnalyzeError,
    ClassifyError,
    GenerateFix,
)

# DSPy is optional
_dspy_available = False
try:
    import dspy
    _dspy_available = True
except ImportError:
    dspy = None


if _dspy_available:

    class ErrorAnalyzer(dspy.Module):
        """Module for analyzing production errors."""

        def __init__(self):
            super().__init__()
            self.analyze = dspy.ChainOfThought(AnalyzeError)

        def forward(
            self,
            exception_type: str,
            message: str,
            stack_trace: str,
            io_transcript: str = "",
            code_context: str = "",
        ):
            """Analyze error and return root cause.

            Args:
                exception_type: Type of exception raised.
                message: Error message.
                stack_trace: Full stack trace.
                io_transcript: I/O operations before error.
                code_context: Relevant source code.

            Returns:
                Analysis result with root_cause, category, confidence.
            """
            return self.analyze(
                exception_type=exception_type,
                message=message,
                stack_trace=stack_trace,
                io_transcript=io_transcript or "No I/O transcript available",
                code_context=code_context or "No code context available",
            )

    class FixGenerator(dspy.Module):
        """Module for generating error fixes."""

        def __init__(self):
            super().__init__()
            self.analyze = dspy.ChainOfThought(AnalyzeError)
            self.generate = dspy.ChainOfThought(GenerateFix)

        def forward(
            self,
            exception_type: str,
            message: str,
            stack_trace: str,
            io_transcript: str = "",
            code_context: str = "",
            similar_fixes: str = "",
        ):
            """Analyze error and generate fix.

            Args:
                exception_type: Type of exception raised.
                message: Error message.
                stack_trace: Full stack trace.
                io_transcript: I/O operations before error.
                code_context: Relevant source code.
                similar_fixes: Past fixes that worked.

            Returns:
                Fix result with fix_strategy, unified_diff, explanation.
            """
            # First analyze the error
            analysis = self.analyze(
                exception_type=exception_type,
                message=message,
                stack_trace=stack_trace,
                io_transcript=io_transcript or "No I/O transcript available",
                code_context=code_context or "No code context available",
            )

            # Then generate a fix based on analysis
            return self.generate(
                exception_type=exception_type,
                root_cause=analysis.root_cause,
                code_context=code_context or "No code context available",
                similar_fixes=similar_fixes or "No similar fixes available",
            )

    class ErrorClassifier(dspy.Module):
        """Module for classifying errors."""

        def __init__(self):
            super().__init__()
            self.classify = dspy.Predict(ClassifyError)

        def forward(
            self,
            exception_type: str,
            message: str,
            stack_trace: str,
        ):
            """Classify error for routing.

            Args:
                exception_type: Type of exception raised.
                message: Error message.
                stack_trace: Stack trace.

            Returns:
                Classification with category, severity, is_user_facing.
            """
            return self.classify(
                exception_type=exception_type,
                message=message,
                stack_trace=stack_trace,
            )

    class HikaflowPipeline(dspy.Module):
        """Full pipeline for error analysis and fix generation."""

        def __init__(self):
            super().__init__()
            self.classifier = ErrorClassifier()
            self.analyzer = ErrorAnalyzer()
            self.generator = FixGenerator()

        def forward(
            self,
            error: Dict[str, Any],
            similar_fixes: Optional[List[Dict[str, Any]]] = None,
        ):
            """Run full pipeline on error.

            Args:
                error: Error dictionary with exception_type, message, stack_trace, etc.
                similar_fixes: List of similar past fixes.

            Returns:
                Pipeline result with classification, analysis, and fix.
            """
            # Format similar fixes
            similar_str = ""
            if similar_fixes:
                similar_str = "\n".join([
                    f"- {f.get('exception_type', 'Unknown')}: {f.get('fix_type', 'unknown')} "
                    f"({f.get('success_rate', 0):.0%} success rate)"
                    for f in similar_fixes
                ])

            # Run classification
            classification = self.classifier(
                exception_type=error.get("exception_type", "Error"),
                message=error.get("message", ""),
                stack_trace=error.get("stack_trace", ""),
            )

            # Run fix generation (includes analysis)
            fix_result = self.generator(
                exception_type=error.get("exception_type", "Error"),
                message=error.get("message", ""),
                stack_trace=error.get("stack_trace", ""),
                io_transcript=error.get("io_transcript", ""),
                code_context=error.get("code_context", ""),
                similar_fixes=similar_str,
            )

            return {
                "classification": {
                    "category": classification.category,
                    "severity": classification.severity,
                    "is_user_facing": classification.is_user_facing,
                },
                "fix": {
                    "strategy": fix_result.fix_strategy,
                    "unified_diff": fix_result.unified_diff,
                    "explanation": fix_result.explanation,
                },
            }

else:
    # Stub classes when DSPy not available
    class ErrorAnalyzer:
        """Stub for ErrorAnalyzer when DSPy not installed."""

        def forward(self, **kwargs):
            raise ImportError("DSPy not installed. Run: pip install dspy-ai")

    class FixGenerator:
        """Stub for FixGenerator when DSPy not installed."""

        def forward(self, **kwargs):
            raise ImportError("DSPy not installed. Run: pip install dspy-ai")

    class ErrorClassifier:
        """Stub for ErrorClassifier when DSPy not installed."""

        def forward(self, **kwargs):
            raise ImportError("DSPy not installed. Run: pip install dspy-ai")

    class HikaflowPipeline:
        """Stub for HikaflowPipeline when DSPy not installed."""

        def forward(self, **kwargs):
            raise ImportError("DSPy not installed. Run: pip install dspy-ai")
