"""LLM-powered hypothesis generation for autonomous debugging.

This module generates fix hypotheses using LLMs and produces
applicable patches for validation.

Uses the validation module (Plan 52) to ensure generated fixes
compile, type-check, and apply cleanly before being suggested.
"""

from __future__ import annotations

import hashlib
import json
import os
import re
from dataclasses import dataclass
from typing import TYPE_CHECKING, Any, Dict, List, Optional

try:
    import anthropic
    HAS_ANTHROPIC = True
except ImportError:
    HAS_ANTHROPIC = False

try:
    import openai
    HAS_OPENAI = True
except ImportError:
    HAS_OPENAI = False

if TYPE_CHECKING:
    from autodebug.validation.fix_validator import ValidationResult


@dataclass
class FilePatch:
    """A patch for a single file."""

    path: str
    original_content: Optional[str]
    patched_content: str
    diff: str

    @property
    def line_count(self) -> int:
        """Count lines changed in the patch."""
        additions = len([l for l in self.diff.splitlines() if l.startswith("+")])
        deletions = len([l for l in self.diff.splitlines() if l.startswith("-")])
        return additions + deletions


@dataclass
class Patch:
    """A complete patch that may span multiple files."""

    files: List[FilePatch]
    raw_diff: str

    @property
    def hash(self) -> str:
        """Compute a stable hash of the patch."""
        content = self.raw_diff.encode("utf-8")
        return hashlib.sha256(content).hexdigest()

    @property
    def line_count(self) -> int:
        """Total lines changed across all files."""
        return sum(f.line_count for f in self.files)

    @property
    def file_count(self) -> int:
        """Number of files modified."""
        return len(self.files)

    @classmethod
    def from_unified_diff(cls, diff_text: str) -> Patch:
        """Parse a unified diff into a Patch object."""
        files = []
        current_file = None
        current_diff_lines = []

        for line in diff_text.splitlines():
            if line.startswith("--- a/"):
                if current_file and current_diff_lines:
                    files.append(FilePatch(
                        path=current_file,
                        original_content=None,
                        patched_content="",
                        diff="\n".join(current_diff_lines),
                    ))
                current_file = line[6:]  # Remove "--- a/"
                current_diff_lines = [line]
            elif line.startswith("+++ b/"):
                current_diff_lines.append(line)
            elif current_file:
                current_diff_lines.append(line)

        # Add last file
        if current_file and current_diff_lines:
            files.append(FilePatch(
                path=current_file,
                original_content=None,
                patched_content="",
                diff="\n".join(current_diff_lines),
            ))

        return cls(files=files, raw_diff=diff_text)


@dataclass
class Hypothesis:
    """A hypothesis for fixing a failure."""

    root_cause: str
    confidence: float
    patch: Patch
    explanation: str
    strategy: str = "fix_the_code"
    validation: Optional[ValidationResult] = None
    patched_code: Optional[str] = None

    @property
    def id(self) -> str:
        """Unique identifier for this hypothesis."""
        return self.patch.hash[:12]

    @property
    def is_validated(self) -> bool:
        """Check if hypothesis has been validated."""
        return self.validation is not None

    @property
    def is_valid(self) -> bool:
        """Check if hypothesis passed validation."""
        return self.validation is not None and self.validation.is_valid


# Prompt templates
HYPOTHESIS_PROMPT = """You are an expert Python debugger. A test is failing deterministically.

## Failure Information
Exception: {exception_type}: {exception_message}
Location: {failure_file}:{failure_line} in {failure_function}

## Stack Trace
{stack_trace}

## Failing Code Context
```python
{failure_code}
```

## Failing Function
```python
{failing_function_source}
```

{dependency_context}

## Task
Generate {num_hypotheses} candidate fixes, ranked by likelihood. For each fix:
1. Explain the likely root cause
2. Provide a minimal patch (unified diff format)
3. Explain why this fix should work

{similar_fixes_section}
{prior_failures_section}

Important guidelines:
- Focus on the ACTUAL bug, not just adding error handling
- Prefer minimal changes that fix the root cause
- Consider edge cases and null checks
- Use the strategy hint: {strategy}

Format your response as JSON:
```json
{{
  "hypotheses": [
    {{
      "root_cause": "Brief explanation of why the failure occurs",
      "confidence": 0.8,
      "patch": "--- a/path/to/file.py\\n+++ b/path/to/file.py\\n@@ -42,6 +42,8 @@ def function():\\n     existing line\\n+    new line\\n     existing line",
      "explanation": "Why this fix works"
    }}
  ]
}}
```
"""

DIVERSITY_STRATEGIES = [
    ("fix_the_code", "Fix the implementation to handle this case correctly"),
    ("fix_the_test", "Adjust the test assertion if expectations are wrong"),
    ("add_null_check", "Add null/None checks to prevent the error"),
    ("add_error_handling", "Add try/except to handle the error gracefully"),
    ("fix_the_data", "Fix test fixtures or input data"),
]


# Model configuration
MODEL_HIERARCHY = [
    {
        "provider": "anthropic",
        "model": "claude-sonnet-4-20250514",
        "max_tokens": 8000,
        "use_for": "complex",
    },
    {
        "provider": "openai",
        "model": "gpt-5",
        "max_tokens": 8000,
        "use_for": "fallback",
    },
    {
        "provider": "anthropic",
        "model": "claude-3-5-haiku-20241022",
        "max_tokens": 4000,
        "use_for": "triage",
    },
]


class CostTracker:
    """Track LLM API costs."""

    # Approximate costs per 1K tokens (USD)
    COSTS = {
        "claude-sonnet-4-20250514": {"input": 0.003, "output": 0.015},
        "claude-3-5-haiku-20241022": {"input": 0.0008, "output": 0.004},
        "gpt-5": {"input": 0.00125, "output": 0.01},
        "gpt-5-mini": {"input": 0.00025, "output": 0.002},
    }

    def __init__(self):
        self.total_input_tokens = 0
        self.total_output_tokens = 0
        self.total_cost = 0.0

    def record(self, model: str, input_tokens: int, output_tokens: int):
        """Record token usage and cost."""
        self.total_input_tokens += input_tokens
        self.total_output_tokens += output_tokens

        costs = self.COSTS.get(model, {"input": 0.01, "output": 0.03})
        cost = (input_tokens / 1000 * costs["input"]) + (output_tokens / 1000 * costs["output"])
        self.total_cost += cost

    def __str__(self) -> str:
        return f"Tokens: {self.total_input_tokens}in/{self.total_output_tokens}out, Cost: ${self.total_cost:.4f}"


def _format_stack_trace(stack: List[Dict[str, Any]]) -> str:
    """Format stack trace for LLM."""
    lines = []
    for frame in stack[:10]:  # Limit to 10 frames
        file = frame.get("file", "?")
        line = frame.get("line", "?")
        func = frame.get("function", "?")
        code = frame.get("code", "")
        lines.append(f"  File \"{file}\", line {line}, in {func}")
        if code:
            lines.append(f"    {code}")
    return "\n".join(lines)


def _format_dependency_context(context: Dict[str, Any]) -> str:
    """Format dependency context for LLM."""
    parts = []

    if context.get("last_http"):
        http = context["last_http"]
        parts.append("## Last HTTP Request")
        parts.append(f"Method: {http.get('method', '?')}")
        parts.append(f"URL: {http.get('url', '?')}")
        parts.append(f"Status: {http.get('status', '?')}")
        if http.get("response_body"):
            body = str(http["response_body"])[:500]
            parts.append(f"Response: {body}")
        parts.append("")

    if context.get("last_sql"):
        sql = context["last_sql"]
        parts.append("## Last SQL Query")
        parts.append(f"Query: {sql.get('query', '?')[:200]}")
        if sql.get("params"):
            parts.append(f"Params: {sql.get('params')}")
        parts.append("")

    return "\n".join(parts) if parts else ""


def _call_anthropic(
    prompt: str,
    model: str,
    max_tokens: int,
    cost_tracker: Optional[CostTracker] = None,
) -> str:
    """Call Anthropic API (or DeepSeek/OpenAI-compatible if configured)."""
    # Prefer DeepSeek/OpenAI-compatible path
    try:
        from autodebug.llm_config import get_llm_config, create_openai_client
        config = get_llm_config()
        if config.is_openai_compatible:
            client = create_openai_client(config)
            response = client.chat.completions.create(
                model=config.default_model,
                max_tokens=max_tokens,
                messages=[{"role": "user", "content": prompt}],
            )
            if cost_tracker and response.usage:
                cost_tracker.record(
                    config.default_model,
                    response.usage.prompt_tokens,
                    response.usage.completion_tokens,
                )
            return response.choices[0].message.content
    except (ImportError, ValueError):
        pass

    if not HAS_ANTHROPIC:
        raise ImportError("anthropic package not installed")

    api_key = os.environ.get("ANTHROPIC_API_KEY")
    if not api_key:
        raise ValueError("ANTHROPIC_API_KEY environment variable not set")

    client = anthropic.Anthropic(api_key=api_key, timeout=120.0)
    response = client.messages.create(
        model=model,
        max_tokens=max_tokens,
        messages=[{"role": "user", "content": prompt}],
    )

    if cost_tracker:
        cost_tracker.record(
            model,
            response.usage.input_tokens,
            response.usage.output_tokens,
        )

    return response.content[0].text


def _call_openai(
    prompt: str,
    model: str,
    max_tokens: int,
    cost_tracker: Optional[CostTracker] = None,
) -> str:
    """Call OpenAI API."""
    if not HAS_OPENAI:
        raise ImportError("openai package not installed")

    api_key = os.environ.get("OPENAI_API_KEY")
    if not api_key:
        raise ValueError("OPENAI_API_KEY environment variable not set")

    client = openai.OpenAI(api_key=api_key, timeout=120.0)
    response = client.chat.completions.create(
        model=model,
        max_tokens=max_tokens,
        messages=[{"role": "user", "content": prompt}],
    )

    if cost_tracker and response.usage:
        cost_tracker.record(
            model,
            response.usage.prompt_tokens,
            response.usage.completion_tokens,
        )

    return response.choices[0].message.content


def _call_llm(
    prompt: str,
    model_config: Dict[str, Any],
    cost_tracker: Optional[CostTracker] = None,
) -> str:
    """Call LLM based on provider configuration."""
    provider = model_config["provider"]
    model = model_config["model"]
    max_tokens = model_config["max_tokens"]

    if provider == "anthropic":
        return _call_anthropic(prompt, model, max_tokens, cost_tracker)
    elif provider == "openai":
        return _call_openai(prompt, model, max_tokens, cost_tracker)
    else:
        raise ValueError(f"Unknown provider: {provider}")


def _parse_llm_response(response: str) -> List[Dict[str, Any]]:
    """Parse LLM response into hypothesis dictionaries."""
    # Extract JSON from response
    json_match = re.search(r"```json\s*(.*?)\s*```", response, re.DOTALL)
    if json_match:
        json_str = json_match.group(1)
    else:
        # Try to find raw JSON
        json_match = re.search(r"\{.*\"hypotheses\".*\}", response, re.DOTALL)
        if json_match:
            json_str = json_match.group(0)
        else:
            return []

    try:
        data = json.loads(json_str)
        return data.get("hypotheses", [])
    except json.JSONDecodeError:
        return []


def generate_hypotheses(
    context: Dict[str, Any],
    num_hypotheses: int = 3,
    strategies: Optional[List[str]] = None,
    model_config: Optional[Dict[str, Any]] = None,
    cost_tracker: Optional[CostTracker] = None,
) -> List[Hypothesis]:
    """Generate fix hypotheses using LLM.

    Args:
        context: Context dictionary from extract_context()
        num_hypotheses: Number of hypotheses to generate
        strategies: Optional list of strategies to try
        model_config: Optional model configuration override
        cost_tracker: Optional cost tracker

    Returns:
        List of Hypothesis objects
    """
    if strategies is None:
        strategies = [s[0] for s in DIVERSITY_STRATEGIES[:num_hypotheses]]

    if model_config is None:
        model_config = MODEL_HIERARCHY[0]

    hypotheses = []

    for strategy in strategies[:num_hypotheses]:
        strategy_desc = dict(DIVERSITY_STRATEGIES).get(strategy, strategy)

        # L-04: Include similar historical fixes as few-shot examples (RAG)
        similar_fixes_section = ""
        similar_fixes = context.get("_similar_fixes", [])
        if similar_fixes:
            parts = ["## Similar Historical Fixes (for reference)"]
            for fix in similar_fixes[:3]:
                parts.append(f"- Exception: {fix.get('exception_type', '?')}")
                parts.append(f"  Strategy: {fix.get('strategy', '?')}")
                if fix.get("patch_summary"):
                    parts.append(f"  Fix summary: {fix['patch_summary'][:200]}")
                parts.append("")
            similar_fixes_section = "\n".join(parts)

        # M-02: Include prior failure feedback for refinement rounds
        prior_failures_section = ""
        prior_failures = context.get("_prior_failures", "")
        if prior_failures:
            prior_failures_section = (
                "## Prior Attempts That Failed\n"
                "These approaches were tried but did not work:\n"
                f"{prior_failures}\n"
                "Please generate DIFFERENT fixes that avoid these failure modes."
            )

        # Build prompt
        prompt = HYPOTHESIS_PROMPT.format(
            exception_type=context.get("exception_type", "Unknown"),
            exception_message=context.get("exception_message", "")[:500],
            failure_file=context.get("failure_file", "?"),
            failure_line=context.get("failure_line", "?"),
            failure_function=context.get("failure_function", "?"),
            stack_trace=_format_stack_trace(context.get("stack_trace", [])),
            failure_code=context.get("failure_code", "# No code available"),
            failing_function_source=context.get("failing_function_source", "# No source available")[:2000],
            dependency_context=_format_dependency_context(context),
            num_hypotheses=1,  # One per strategy
            strategy=strategy_desc,
            similar_fixes_section=similar_fixes_section,
            prior_failures_section=prior_failures_section,
        )

        try:
            response = _call_llm(prompt, model_config, cost_tracker)
            parsed = _parse_llm_response(response)

            for h in parsed:
                patch = Patch.from_unified_diff(h.get("patch", ""))
                hypotheses.append(Hypothesis(
                    root_cause=h.get("root_cause", "Unknown"),
                    confidence=float(h.get("confidence", 0.5)),
                    patch=patch,
                    explanation=h.get("explanation", ""),
                    strategy=strategy,
                ))
        except Exception as e:
            # Log error but continue with other strategies
            print(f"Error generating hypothesis for {strategy}: {e}")
            continue

    return hypotheses


def generate_hypotheses_with_fallback(
    context: Dict[str, Any],
    num_hypotheses: int = 3,
    cost_tracker: Optional[CostTracker] = None,
    max_rounds: int = 2,
    validation_callback: Optional[callable] = None,
) -> List[Hypothesis]:
    """Generate hypotheses with model fallback and iterative refinement.

    Implements a 2-round generate-validate-refine loop (M-02):
    if the first round of hypotheses fails validation, their
    failure reasons are fed back into a refinement prompt to
    generate improved hypotheses in round 2.

    Args:
        context: Context dictionary
        num_hypotheses: Number of hypotheses to generate
        cost_tracker: Optional cost tracker
        max_rounds: Maximum refinement rounds (default: 2)
        validation_callback: Optional callable(List[Hypothesis]) -> List[tuple[Hypothesis, str]]
            that returns (hypothesis, failure_reason) pairs for failed hypotheses.
            If None, no refinement loop is used.

    Returns:
        List of Hypothesis objects
    """
    # Filter models based on available API keys
    has_anthropic = bool(os.environ.get("ANTHROPIC_API_KEY"))
    has_openai = bool(os.environ.get("OPENAI_API_KEY"))

    available_models = [
        m for m in MODEL_HIERARCHY
        if (m["provider"] == "anthropic" and has_anthropic)
        or (m["provider"] == "openai" and has_openai)
    ]

    if not available_models:
        print("Error: No API keys available (ANTHROPIC_API_KEY or OPENAI_API_KEY)")
        return []

    all_hypotheses: List[Hypothesis] = []

    for model_config in available_models:
        try:
            hypotheses = generate_hypotheses(
                context,
                num_hypotheses=num_hypotheses,
                model_config=model_config,
                cost_tracker=cost_tracker,
            )
            all_hypotheses.extend(hypotheses)

            # M-02: Iterative refinement loop.
            # If a validation callback is provided, check which hypotheses
            # failed and generate refined replacements using failure feedback.
            if validation_callback and all_hypotheses and max_rounds > 1:
                for round_num in range(1, max_rounds):
                    failures = validation_callback(all_hypotheses)
                    if not failures:
                        break  # All hypotheses passed, no refinement needed

                    # Build a refinement context with failure reasons
                    failure_summaries = []
                    for failed_hyp, reason in failures[:3]:  # Top 3 failures
                        failure_summaries.append(
                            f"- Strategy '{failed_hyp.strategy}': {reason}"
                        )

                    refine_context = dict(context)
                    refine_context["_prior_failures"] = "\n".join(failure_summaries)

                    refined = generate_hypotheses(
                        refine_context,
                        num_hypotheses=min(len(failures), num_hypotheses),
                        model_config=model_config,
                        cost_tracker=cost_tracker,
                    )

                    # Replace failed hypotheses with refined ones
                    failed_ids = {h.id for h, _ in failures}
                    all_hypotheses = [
                        h for h in all_hypotheses if h.id not in failed_ids
                    ] + refined

            return all_hypotheses
        except Exception as e:
            print(f"Model {model_config['model']} failed: {e}, trying fallback...")
            continue

    return []


def validate_hypotheses(
    hypotheses: List[Hypothesis],
    source_code: str,
    file_path: Optional[str] = None,
    use_ty: bool = True,
    auto_fix: bool = True,
) -> List[Hypothesis]:
    """Validate hypotheses and filter to only valid ones.

    Uses the validation module (Plan 52) to ensure fixes compile,
    type-check, and apply cleanly.

    Args:
        hypotheses: List of hypotheses to validate
        source_code: Original source code that will be patched
        file_path: Optional path for context
        use_ty: Use Ty type checker (faster) instead of mypy
        auto_fix: Attempt to auto-fix validation failures

    Returns:
        List of validated hypotheses (may be fewer than input)
    """
    try:
        from autodebug.validation import AutoFixer, FixValidator
    except ImportError:
        # Validation module not available, return all hypotheses
        print("Warning: validation module not available, skipping validation")
        return hypotheses

    validator = FixValidator(use_ty=use_ty)
    fixer = AutoFixer() if auto_fix else None
    validated = []

    for hypothesis in hypotheses:
        # Run validation
        result = validator.validate(hypothesis, source_code, file_path)

        if result.is_valid:
            # Passed validation
            hypothesis.validation = result
            hypothesis.patched_code = result.patched_code
            validated.append(hypothesis)
        elif auto_fix and fixer and result.patched_code:
            # Try auto-fix
            fixed_code = fixer.try_auto_fix(result.patched_code, result.errors)
            if fixed_code:
                # Re-validate the fixed code
                revalidation = validator.validate_code(fixed_code)
                if revalidation.is_valid:
                    hypothesis.validation = revalidation
                    hypothesis.patched_code = fixed_code
                    validated.append(hypothesis)
                    continue

            # Auto-fix didn't help
            print(f"Hypothesis {hypothesis.id} failed validation: {result.summary()}")
        else:
            print(f"Hypothesis {hypothesis.id} failed validation: {result.summary()}")

    return validated


def generate_and_validate_hypotheses(
    context: Dict[str, Any],
    source_code: str,
    num_hypotheses: int = 3,
    file_path: Optional[str] = None,
    use_ty: bool = True,
    auto_fix: bool = True,
    cost_tracker: Optional[CostTracker] = None,
) -> List[Hypothesis]:
    """Generate hypotheses and validate them.

    Combines generation and validation into a single function.

    Args:
        context: Context dictionary from extract_context()
        source_code: Original source code
        num_hypotheses: Number of hypotheses to generate
        file_path: Optional file path for context
        use_ty: Use Ty type checker
        auto_fix: Attempt to auto-fix validation failures
        cost_tracker: Optional cost tracker

    Returns:
        List of validated Hypothesis objects
    """
    # Generate hypotheses
    hypotheses = generate_hypotheses_with_fallback(
        context,
        num_hypotheses=num_hypotheses,
        cost_tracker=cost_tracker,
    )

    if not hypotheses:
        return []

    # Validate
    return validate_hypotheses(
        hypotheses,
        source_code,
        file_path=file_path,
        use_ty=use_ty,
        auto_fix=auto_fix,
    )


def format_hypothesis_report(hypotheses: List[Hypothesis]) -> str:
    """Format hypotheses as human-readable report.

    Args:
        hypotheses: List of hypotheses

    Returns:
        Formatted report string
    """
    if not hypotheses:
        return "No hypotheses generated."

    lines = [
        "",
        "Generated Hypotheses",
        "=" * 50,
        "",
    ]

    for i, h in enumerate(hypotheses, 1):
        lines.append(f"### Hypothesis {i} (confidence: {h.confidence:.0%})")
        lines.append(f"Strategy: {h.strategy}")
        lines.append(f"Root cause: {h.root_cause}")

        # Add validation status if available
        if h.validation:
            status = "VALID" if h.validation.is_valid else "INVALID"
            lines.append(f"Validation: {status} ({h.validation.duration_ms:.1f}ms)")
            if not h.validation.is_valid and h.validation.errors:
                lines.append("Validation Errors:")
                for err in h.validation.errors[:3]:  # Show first 3 errors
                    lines.append(f"  - {err}")

        lines.append("")
        lines.append("Patch:")
        lines.append("```diff")
        lines.append(h.patch.raw_diff[:1000])
        if len(h.patch.raw_diff) > 1000:
            lines.append("... (truncated)")
        lines.append("```")
        lines.append("")
        lines.append(f"Explanation: {h.explanation}")
        lines.append("")
        lines.append("-" * 40)
        lines.append("")

    return "\n".join(lines)
