"""Fix generation agent for Hikaflow.

This module provides the FixGenerator agent that generates
code fixes for test failures and flaky tests.
"""

from __future__ import annotations

import logging
from dataclasses import dataclass, field
from enum import Enum
from typing import Any, Dict, List, Optional

from autodebug.agent.base import (
    AgentCapability,
    AgentContext,
    BaseAgent,
)

logger = logging.getLogger(__name__)


class FixStrategy(str, Enum):
    """Strategy for generating fixes."""

    CONSERVATIVE = "conservative"  # Minimal changes
    RAG_BASED = "rag_based"  # Based on similar past fixes
    FRAMEWORK_SPECIFIC = "framework_specific"  # Framework patterns
    CREATIVE = "creative"  # Higher temperature exploration
    ENSEMBLE = "ensemble"  # Combine multiple approaches
    TEMPLATE_BASED = "template_based"  # Use fix templates


class FixConfidence(str, Enum):
    """Confidence level for a generated fix."""

    HIGH = "high"  # Very likely to work
    MEDIUM = "medium"  # Probably will work
    LOW = "low"  # May work, needs verification
    EXPERIMENTAL = "experimental"  # Speculative


@dataclass
class GeneratedFix:
    """A generated code fix."""

    # Fix content
    diff: str
    description: str
    file_path: str

    # Metadata
    strategy: FixStrategy
    confidence: FixConfidence
    confidence_score: float  # 0.0-1.0

    # Reasoning
    reasoning: str = ""
    root_cause_addressed: bool = True

    # Validation state
    syntax_valid: Optional[bool] = None
    types_valid: Optional[bool] = None
    tests_passed: Optional[bool] = None

    # Related data
    similar_fixes: List[str] = field(default_factory=list)
    tags: List[str] = field(default_factory=list)

    def to_dict(self) -> Dict[str, Any]:
        """Convert to dictionary."""
        return {
            "diff": self.diff,
            "description": self.description,
            "file_path": self.file_path,
            "strategy": self.strategy.value,
            "confidence": self.confidence.value,
            "confidence_score": self.confidence_score,
            "reasoning": self.reasoning,
            "root_cause_addressed": self.root_cause_addressed,
            "syntax_valid": self.syntax_valid,
            "types_valid": self.types_valid,
            "tests_passed": self.tests_passed,
            "similar_fixes": self.similar_fixes,
            "tags": self.tags,
        }


@dataclass
class FixGenerationResult:
    """Result of fix generation."""

    fixes: List[GeneratedFix]
    best_fix_index: int = 0
    total_candidates: int = 0
    elapsed_seconds: float = 0.0

    @property
    def best_fix(self) -> Optional[GeneratedFix]:
        """Get the best fix."""
        if not self.fixes:
            return None
        return self.fixes[self.best_fix_index]

    @property
    def has_high_confidence_fix(self) -> bool:
        """Check if any fix has high confidence."""
        return any(f.confidence == FixConfidence.HIGH for f in self.fixes)


class FixGenerator(BaseAgent[FixGenerationResult]):
    """Agent that generates code fixes for test failures.

    This agent:
    1. Analyzes the root cause of the failure
    2. Searches for similar past fixes
    3. Generates multiple fix candidates
    4. Ranks fixes by likelihood of success
    5. Validates syntax and types

    Example:
        generator = FixGenerator()
        context = AgentContext(
            project_path=".",
            run_id="run_123",
            error_type="KeyError",
            error_message="'user_id'",
        )
        result = await generator.execute(context)

        if result.success:
            for fix in result.metadata["result"].fixes:
                print(f"Fix: {fix.description}")
                print(f"Confidence: {fix.confidence}")
    """

    name = "fix_generator"
    description = "Generates code fixes for test failures"
    capabilities = [
        AgentCapability.READ_CODE,
        AgentCapability.SEARCH_CODE,
        AgentCapability.READ_TRANSCRIPTS,
        AgentCapability.GENERATE_FIX,
        AgentCapability.ANALYZE_ERROR,
    ]

    def __init__(
        self,
        strategies: Optional[List[FixStrategy]] = None,
        max_fixes: int = 5,
        include_rag: bool = True,
        **kwargs,
    ):
        """Initialize fix generator.

        Args:
            strategies: Strategies to use (default: all)
            max_fixes: Maximum fixes to generate
            include_rag: Whether to search for similar fixes
            **kwargs: Passed to BaseAgent
        """
        super().__init__(**kwargs)
        self.strategies = strategies or list(FixStrategy)
        self.max_fixes = max_fixes
        self.include_rag = include_rag

    async def run(self, context: AgentContext) -> FixGenerationResult:
        """Generate fixes for the failure.

        Args:
            context: Execution context with error info

        Returns:
            FixGenerationResult with generated fixes
        """
        fixes: List[GeneratedFix] = []

        # 1. Search for similar past fixes if RAG enabled
        similar_fixes = []
        if self.include_rag and FixStrategy.RAG_BASED in self.strategies:
            similar_fixes = await self._search_similar_fixes(context)
            logger.info(f"Found {len(similar_fixes)} similar past fixes")

        # 2. Generate fixes using each strategy
        for strategy in self.strategies:
            if len(fixes) >= self.max_fixes:
                break

            try:
                strategy_fixes = await self._generate_with_strategy(
                    context, strategy, similar_fixes
                )
                fixes.extend(strategy_fixes)
                logger.info(f"Strategy {strategy.value} generated {len(strategy_fixes)} fixes")
            except Exception as e:
                logger.warning(f"Strategy {strategy.value} failed: {e}")

        # 3. Deduplicate and rank fixes
        fixes = self._deduplicate_fixes(fixes)
        fixes = self._rank_fixes(fixes)

        # 4. Validate top fixes
        for fix in fixes[:3]:  # Validate top 3
            await self._validate_fix(fix, context)

        # 5. Determine best fix
        best_index = self._select_best_fix(fixes)

        return FixGenerationResult(
            fixes=fixes[:self.max_fixes],
            best_fix_index=best_index,
            total_candidates=len(fixes),
            elapsed_seconds=context.elapsed_seconds(),
        )

    async def _search_similar_fixes(
        self,
        context: AgentContext,
    ) -> List[Dict[str, Any]]:
        """Search for similar past fixes using RAG.

        Args:
            context: Execution context

        Returns:
            List of similar fix records
        """
        # This would call the hybrid search in production
        # For now, return empty list
        return []

    async def _generate_with_strategy(
        self,
        context: AgentContext,
        strategy: FixStrategy,
        similar_fixes: List[Dict[str, Any]],
    ) -> List[GeneratedFix]:
        """Generate fixes using a specific strategy.

        Args:
            context: Execution context
            strategy: Strategy to use
            similar_fixes: Similar past fixes for reference

        Returns:
            List of generated fixes
        """
        if strategy == FixStrategy.CONSERVATIVE:
            return await self._generate_conservative(context)
        elif strategy == FixStrategy.RAG_BASED:
            return await self._generate_from_rag(context, similar_fixes)
        elif strategy == FixStrategy.FRAMEWORK_SPECIFIC:
            return await self._generate_framework_specific(context)
        elif strategy == FixStrategy.CREATIVE:
            return await self._generate_creative(context)
        elif strategy == FixStrategy.TEMPLATE_BASED:
            return await self._generate_from_templates(context)
        elif strategy == FixStrategy.ENSEMBLE:
            return await self._generate_ensemble(context, similar_fixes)
        return []

    async def _generate_conservative(
        self,
        context: AgentContext,
    ) -> List[GeneratedFix]:
        """Generate minimal, conservative fixes."""
        fixes = []

        # Generate based on error type
        if context.error_type == "KeyError":
            fix = GeneratedFix(
                diff=self._create_keyerror_fix(context),
                description="Use .get() with default value to handle missing key",
                file_path=self._extract_file_from_trace(context.stack_trace or ""),
                strategy=FixStrategy.CONSERVATIVE,
                confidence=FixConfidence.HIGH,
                confidence_score=0.85,
                reasoning="KeyError can be safely fixed by using dict.get() with a default",
                tags=["keyerror", "dict", "defensive"],
            )
            fixes.append(fix)

        elif context.error_type == "AttributeError":
            fix = GeneratedFix(
                diff=self._create_attributeerror_fix(context),
                description="Add None check before attribute access",
                file_path=self._extract_file_from_trace(context.stack_trace or ""),
                strategy=FixStrategy.CONSERVATIVE,
                confidence=FixConfidence.MEDIUM,
                confidence_score=0.70,
                reasoning="AttributeError often caused by None values",
                tags=["attributeerror", "none-check", "defensive"],
            )
            fixes.append(fix)

        return fixes

    async def _generate_from_rag(
        self,
        context: AgentContext,
        similar_fixes: List[Dict[str, Any]],
    ) -> List[GeneratedFix]:
        """Generate fixes based on similar past fixes."""
        fixes = []

        for similar in similar_fixes[:3]:
            fix = GeneratedFix(
                diff=similar.get("diff", ""),
                description=f"Based on similar fix: {similar.get('description', '')}",
                file_path=self._extract_file_from_trace(context.stack_trace or ""),
                strategy=FixStrategy.RAG_BASED,
                confidence=FixConfidence.MEDIUM,
                confidence_score=similar.get("similarity", 0.6),
                reasoning=f"Similar error was fixed with this approach in {similar.get('project', 'another project')}",
                similar_fixes=[similar.get("id", "")],
                tags=["rag", "similar-fix"],
            )
            fixes.append(fix)

        return fixes

    async def _generate_framework_specific(
        self,
        context: AgentContext,
    ) -> List[GeneratedFix]:
        """Generate fixes using framework-specific patterns."""
        # Would detect framework and apply patterns
        return []

    async def _generate_creative(
        self,
        context: AgentContext,
    ) -> List[GeneratedFix]:
        """Generate creative fixes with higher temperature."""
        # Would use LLM with higher temperature
        return []

    async def _generate_from_templates(
        self,
        context: AgentContext,
    ) -> List[GeneratedFix]:
        """Generate fixes using fix templates."""
        # Would apply AST-based templates
        return []

    async def _generate_ensemble(
        self,
        context: AgentContext,
        similar_fixes: List[Dict[str, Any]],
    ) -> List[GeneratedFix]:
        """Combine multiple approaches."""
        return []

    def _deduplicate_fixes(self, fixes: List[GeneratedFix]) -> List[GeneratedFix]:
        """Remove duplicate fixes."""
        seen_diffs = set()
        unique = []

        for fix in fixes:
            # Normalize diff for comparison
            normalized = fix.diff.strip()
            if normalized not in seen_diffs:
                seen_diffs.add(normalized)
                unique.append(fix)

        return unique

    def _rank_fixes(self, fixes: List[GeneratedFix]) -> List[GeneratedFix]:
        """Rank fixes by likelihood of success."""
        return sorted(
            fixes,
            key=lambda f: (
                f.confidence_score,
                f.confidence == FixConfidence.HIGH,
                f.root_cause_addressed,
            ),
            reverse=True,
        )

    async def _validate_fix(
        self,
        fix: GeneratedFix,
        context: AgentContext,
    ) -> None:
        """Validate a fix's syntax and types.

        Attempts full-file validation first: if the target file exists,
        applies the patch to a copy and runs ast.parse on the complete
        result. Falls back to validating only the added lines in
        isolation (which may give false confidence if indentation
        context is wrong).
        """
        import ast as _ast
        import os

        added_lines = [
            line[1:] for line in (fix.diff or "").splitlines()
            if line.startswith('+') and not line.startswith('+++')
        ]

        if not added_lines:
            fix.syntax_valid = None  # Cannot verify without added lines
            fix.types_valid = None
            return

        # Try full-file validation if the target file is accessible
        full_file_validated = False
        if fix.file_path and context.project_path:
            target = os.path.join(context.project_path, fix.file_path)
            if os.path.isfile(target):
                try:
                    with open(target, "r") as f:
                        original = f.read()
                    # Apply a simplified patch: append added lines to original
                    # (a full unified-diff apply would be more accurate, but
                    # this catches most syntax issues in the combined code)
                    patched = original + "\n" + "\n".join(added_lines)
                    _ast.parse(patched)
                    fix.syntax_valid = True
                    full_file_validated = True
                except (SyntaxError, OSError):
                    fix.syntax_valid = False
                    full_file_validated = True

        # Fallback: validate added lines in isolation
        if not full_file_validated:
            try:
                _ast.parse("\n".join(added_lines))
                fix.syntax_valid = True
            except SyntaxError:
                fix.syntax_valid = False

        # Type validation would run mypy
        fix.types_valid = None  # Unknown until checked

    def _select_best_fix(self, fixes: List[GeneratedFix]) -> int:
        """Select the best fix index."""
        if not fixes:
            return 0

        # Prefer validated, high confidence fixes
        for i, fix in enumerate(fixes):
            if (
                fix.confidence == FixConfidence.HIGH
                and fix.syntax_valid is True
                and fix.root_cause_addressed
            ):
                return i

        # Fall back to first (highest ranked)
        return 0

    def _create_keyerror_fix(self, context: AgentContext) -> str:
        """Create a fix diff for KeyError."""
        # Extract key from error message
        key = context.error_message or "key"
        key = key.strip("'\"")

        return f"""--- a/file.py
+++ b/file.py
@@ -1,3 +1,5 @@
-value = data['{key}']
+value = data.get('{key}')
+if value is None:
+    raise ValueError("Missing required field: {key}")
"""

    def _create_attributeerror_fix(self, context: AgentContext) -> str:
        """Create a fix diff for AttributeError."""
        return """--- a/file.py
+++ b/file.py
@@ -1,3 +1,4 @@
+if obj is None:
+    raise ValueError("Expected object, got None")
 result = obj.method()
"""

    def _extract_file_from_trace(self, stack_trace: str) -> str:
        """Extract file path from stack trace."""
        import re

        match = re.search(r'File "([^"]+)"', stack_trace)
        if match:
            return match.group(1)
        return "unknown.py"
