"""RAG-enhanced prompt building for fix generation.

Uses similar past fixes to improve LLM prompts for generating
better fixes.
"""

from __future__ import annotations

import logging
from typing import TYPE_CHECKING, Any, Dict, List, Optional

if TYPE_CHECKING:
    from autodebug.learning.search import FixSearchService

logger = logging.getLogger(__name__)


# Common fix patterns for cold start (no historical data)
BOOTSTRAP_PATTERNS = [
    {
        "exception_type": "KeyError",
        "pattern": r"KeyError: '(\w+)'",
        "root_cause": "Dictionary key does not exist. Common causes: typo, missing data from API, or uninitialized dictionary.",
        "fix_hint": "Use .get() with default, or check key existence first.",
    },
    {
        "exception_type": "TypeError",
        "pattern": r"'NoneType' object is not subscriptable",
        "root_cause": "Attempting to index None. Variable is None when expected to have a value.",
        "fix_hint": "Add null check before access, or trace why value is None.",
    },
    {
        "exception_type": "AttributeError",
        "pattern": r"'NoneType' object has no attribute",
        "root_cause": "Calling method on None object. Variable unexpectedly None.",
        "fix_hint": "Add null check or use Optional pattern.",
    },
    {
        "exception_type": "IndexError",
        "pattern": r"list index out of range",
        "root_cause": "Accessing list with invalid index. List may be empty or shorter than expected.",
        "fix_hint": "Check list length before access, or handle empty case.",
    },
    {
        "exception_type": "AssertionError",
        "pattern": r"assert .* == .*",
        "root_cause": "Assertion failed. Could be logic error, test setup issue, or flaky test.",
        "fix_hint": "Review assertion logic and test data setup.",
    },
    {
        "exception_type": "ValueError",
        "pattern": r"invalid literal for int\(\)",
        "root_cause": "String cannot be converted to integer. Input validation missing.",
        "fix_hint": "Add input validation or try/except with meaningful error.",
    },
    {
        "exception_type": "FileNotFoundError",
        "pattern": r"No such file or directory",
        "root_cause": "File path doesn't exist. Path may be relative or test setup incomplete.",
        "fix_hint": "Use pathlib for path handling, ensure test fixtures exist.",
    },
    {
        "exception_type": "ConnectionError",
        "pattern": r"Connection refused|Failed to establish",
        "root_cause": "Network connection failed. Service may be down or port blocked.",
        "fix_hint": "Add retry logic, mock external services in tests.",
    },
]


class EnhancedPromptBuilder:
    """Build LLM prompts with RAG-enhanced examples."""

    def __init__(
        self,
        search_service: Optional[FixSearchService] = None,
        max_examples: int = 3,
        include_bootstrap: bool = True,
    ):
        """Initialize prompt builder.

        Args:
            search_service: Service for finding similar fixes.
            max_examples: Maximum number of examples to include.
            include_bootstrap: Include bootstrap patterns when no data.
        """
        self.search = search_service
        self.max_examples = max_examples
        self.include_bootstrap = include_bootstrap

    async def build_fix_prompt(
        self,
        failure_context: Dict[str, Any],
        code_context: str,
        org_id: Optional[str] = None,
        include_examples: bool = True,
    ) -> str:
        """Build enhanced prompt with similar fix examples.

        Args:
            failure_context: Information about the failure.
            code_context: Code around the failure.
            org_id: Organization ID for filtering.
            include_examples: Whether to include RAG examples.

        Returns:
            Complete prompt for LLM.
        """
        # Build base prompt
        prompt = self._build_base_prompt(failure_context, code_context)

        if not include_examples:
            return prompt

        # Find similar fixes if search service available
        examples = []
        if self.search:
            try:
                similar_fixes = await self.search.find_similar_fixes(
                    failure_context=failure_context,
                    org_id=org_id,
                    top_k=self.max_examples,
                    only_accepted=True,
                )
                examples.extend(similar_fixes)
            except Exception as e:
                logger.warning(f"Failed to find similar fixes: {e}")

        # Add bootstrap patterns if needed
        if len(examples) < self.max_examples and self.include_bootstrap:
            bootstrap = self._get_bootstrap_patterns(failure_context)
            # Only add bootstrap if we have few real examples
            examples_needed = self.max_examples - len(examples)
            examples.extend(bootstrap[:examples_needed])

        # Add examples to prompt
        if examples:
            prompt += self._format_examples(examples)

        return prompt

    def _build_base_prompt(self, failure: Dict[str, Any], code: str) -> str:
        """Build base prompt without examples."""
        exception_type = failure.get("exception_type", "Exception")
        message = failure.get("exception_message", "N/A")
        file_path = failure.get("failing_file", "N/A")
        line_no = failure.get("failing_line", "N/A")
        test_name = failure.get("test_name", "N/A")

        return f"""You are an expert software engineer debugging a test failure.

## Failure Information

**Exception Type:** {exception_type}
**Message:** {message}
**File:** {file_path}:{line_no}
**Test:** {test_name}

## Code Context

```python
{code[:8000]}
```

## Task

Analyze this failure and generate a fix. Your response MUST include:

1. **Root Cause Analysis** (1-2 sentences explaining why this failed)
2. **Fix** in unified diff format:
```diff
--- a/file.py
+++ b/file.py
@@ -line,count +line,count @@
 context
-old line
+new line
 context
```
3. **Confidence Score** (0.0-1.0, how confident you are in this fix)

Focus on fixing the actual bug, not suppressing the error.
Be precise and minimal - only change what's necessary.
"""

    def _format_examples(self, examples: List[Any]) -> str:
        """Format examples for inclusion in prompt."""
        formatted = "\n\n## Similar Past Fixes (Successfully Applied)\n\n"
        formatted += "The following are similar failures that were successfully fixed. "
        formatted += "Use these as guidance, but generate a fix specific to the current failure.\n\n"

        for i, ex in enumerate(examples, 1):
            if hasattr(ex, "exception_type"):
                # SimilarFix object
                formatted += f"### Example {i}"
                if hasattr(ex, "similarity_score"):
                    formatted += f" (Similarity: {ex.similarity_score:.0%})"
                formatted += "\n"
                formatted += f"**Exception:** {ex.exception_type}\n"
                if ex.root_cause:
                    formatted += f"**Root Cause:** {ex.root_cause}\n"
                if ex.patch:
                    patch_preview = ex.patch[:600]
                    if len(ex.patch) > 600:
                        patch_preview += "\n..."
                    formatted += f"**Fix Pattern:**\n```diff\n{patch_preview}\n```\n\n"
            elif isinstance(ex, dict):
                # Bootstrap pattern
                formatted += f"### Pattern: {ex.get('exception_type', 'Unknown')}\n"
                formatted += f"**Common Root Cause:** {ex.get('root_cause', '')}\n"
                formatted += f"**Fix Approach:** {ex.get('fix_hint', '')}\n\n"

        return formatted

    def _get_bootstrap_patterns(self, failure_context: Dict[str, Any]) -> List[Dict[str, Any]]:
        """Get matching bootstrap patterns for cold start."""
        exception_type = failure_context.get("exception_type", "")

        matching = []
        for pattern in BOOTSTRAP_PATTERNS:
            if pattern["exception_type"] == exception_type:
                matching.append(pattern)

        # Also include some general patterns if no exact match
        if not matching:
            matching = BOOTSTRAP_PATTERNS[:2]

        return matching


class FixResponseParser:
    """Parse LLM responses into structured fix data."""

    @staticmethod
    def parse_response(response: str) -> Dict[str, Any]:
        """Parse LLM response into fix components.

        Args:
            response: Raw LLM response text.

        Returns:
            Parsed fix with root_cause, patch, confidence.
        """
        import re

        result = {
            "root_cause": "",
            "patch": "",
            "confidence": 0.5,
            "raw_response": response,
        }

        # Extract root cause (first paragraph or section)
        root_cause_patterns = [
            r"\*\*Root Cause[^*]*\*\*[:\s]*(.+?)(?=\n\n|\*\*|```|$)",
            r"Root Cause[:\s]*(.+?)(?=\n\n|\*\*|```|$)",
            r"^(.+?)(?=\n\n|```)",
        ]
        for pattern in root_cause_patterns:
            match = re.search(pattern, response, re.IGNORECASE | re.DOTALL)
            if match:
                result["root_cause"] = match.group(1).strip()
                break

        # Extract diff patch
        diff_patterns = [
            r"```diff\n(.*?)```",
            r"```\n(---.*?\+\+\+.*?)```",
            r"(---\s*a/.*?\n\+\+\+\s*b/.*?\n@@.*?(?=\n\n|\Z))",
        ]
        for pattern in diff_patterns:
            match = re.search(pattern, response, re.DOTALL)
            if match:
                result["patch"] = match.group(1).strip()
                break

        # Extract confidence
        confidence_patterns = [
            r"\*\*Confidence[^*]*\*\*[:\s]*([\d.]+)",
            r"Confidence[:\s]*([\d.]+)",
            r"confidence[:\s]*([\d.]+)",
        ]
        for pattern in confidence_patterns:
            match = re.search(pattern, response, re.IGNORECASE)
            if match:
                try:
                    conf = float(match.group(1))
                    if conf > 1:
                        conf = conf / 100  # Handle percentage
                    result["confidence"] = min(max(conf, 0), 1)
                except ValueError:
                    pass
                break

        return result
