"""
PromptOptimizer: Generates improved prompts based on trace analysis.

Uses M2.1 to synthesize analysis results into concrete prompt improvements,
with full reasoning transparency via interleaved thinking.
"""

import json
import os
from typing import Any

import anthropic

from reasoning_trace_optimizer.models import (
    AnalysisResult,
    OptimizationResult,
    PromptDiff,
    ReasoningTrace,
)


OPTIMIZER_SYSTEM_PROMPT = """You are an expert prompt engineer specializing in AI agent optimization.

Your task is to improve agent prompts based on reasoning trace analysis.
You have access to:
1. The original prompt that was used
2. Analysis of how the agent reasoned (its thinking trace)
3. Detected patterns and issues
4. Specific recommendations

Your goal is to create an IMPROVED prompt that:
- Addresses identified weaknesses
- Maintains existing strengths
- Prevents detected failure patterns
- Improves clarity and specificity

When optimizing, consider:
- Adding explicit guardrails for common failure modes
- Clarifying ambiguous instructions
- Adding examples for complex behaviors
- Restructuring for better context positioning
- Adding validation steps where missing

Provide the optimized prompt with clear explanations of changes."""


OPTIMIZATION_PROMPT_TEMPLATE = """Optimize the following agent prompt based on trace analysis:

## Original Task
{task}

## Original System Prompt
```
{original_prompt}
```

## Analysis Results

### Overall Score: {overall_score}/100

### Detected Patterns
{patterns}

### Weaknesses
{weaknesses}

### Recommendations
{recommendations}

### Analyzer's Reasoning
{analyzer_thinking}

---

Provide your optimization as JSON:
```json
{{
    "optimized_prompt": "<the full improved prompt>",
    "diffs": [
        {{
            "section": "<which part changed, e.g., 'instructions', 'guardrails', 'examples'>",
            "original": "<original text or 'N/A' if new>",
            "optimized": "<new/changed text>",
            "reason": "<why this change helps>"
        }}
    ],
    "key_changes": [
        "<summary of major change 1>",
        "<summary of major change 2>"
    ],
    "predicted_improvement": 15,
    "confidence": 0.75
}}
```

Think carefully about what changes will have the biggest impact on agent performance."""


class PromptOptimizer:
    """
    Optimizes agent prompts based on reasoning trace analysis.

    Uses M2.1's interleaved thinking to generate thoughtful improvements
    with full transparency into the optimization reasoning.

    Example:
        ```python
        optimizer = PromptOptimizer()
        result = optimizer.optimize(
            original_prompt=system_prompt,
            analysis=analysis_result,
            trace=reasoning_trace
        )

        print(f"Predicted improvement: {result.predicted_improvement}%")
        print(f"New prompt:\\n{result.optimized_prompt}")
        ```
    """

    def __init__(
        self,
        api_key: str | None = None,
        base_url: str = "https://api.minimax.io/anthropic",
        model: str = "MiniMax-M2.1",
    ):
        """
        Initialize PromptOptimizer with M2.1 configuration.

        Args:
            api_key: MiniMax API key
            base_url: API endpoint
            model: Model for optimization
        """
        self.model = model
        self.client = anthropic.Anthropic(
            api_key=api_key or os.environ.get("ANTHROPIC_API_KEY"),
            base_url=base_url,
        )

    def optimize(
        self,
        original_prompt: str,
        analysis: AnalysisResult,
        trace: ReasoningTrace | None = None,
        max_tokens: int = 8192,
    ) -> OptimizationResult:
        """
        Generate an optimized prompt based on analysis.

        Args:
            original_prompt: The original system prompt to improve
            analysis: Analysis results from TraceAnalyzer
            trace: Optional original trace for additional context
            max_tokens: Maximum tokens for response

        Returns:
            OptimizationResult with new prompt and change details
        """
        # Format analysis for prompt
        patterns_text = self._format_patterns(analysis)
        weaknesses_text = "\n".join(f"- {w}" for w in analysis.weaknesses)
        recommendations_text = "\n".join(f"- {r}" for r in analysis.recommendations)

        prompt = OPTIMIZATION_PROMPT_TEMPLATE.format(
            task=trace.task if trace else "Unknown task",
            original_prompt=original_prompt,
            overall_score=analysis.overall_score,
            patterns=patterns_text,
            weaknesses=weaknesses_text or "None identified",
            recommendations=recommendations_text or "None provided",
            analyzer_thinking=analysis.analyzer_thinking[:2000] if analysis.analyzer_thinking else "Not available",
        )

        # Call M2.1 for optimization
        response = self.client.messages.create(
            model=self.model,
            max_tokens=max_tokens,
            system=OPTIMIZER_SYSTEM_PROMPT,
            messages=[{"role": "user", "content": prompt}],
        )

        # Extract thinking and response
        optimizer_thinking = ""
        optimization_text = ""

        for block in response.content:
            if block.type == "thinking":
                optimizer_thinking = block.thinking
            elif block.type == "text":
                optimization_text = block.text

        # Parse the response
        result = self._parse_optimization_response(optimization_text, original_prompt)
        result.optimizer_thinking = optimizer_thinking

        return result

    def optimize_iterative(
        self,
        original_prompt: str,
        analyses: list[AnalysisResult],
        traces: list[ReasoningTrace],
    ) -> OptimizationResult:
        """
        Optimize based on multiple analysis iterations.

        Synthesizes patterns across multiple runs for more robust improvements.

        Args:
            original_prompt: The original system prompt
            analyses: List of analysis results from multiple runs
            traces: Corresponding reasoning traces

        Returns:
            OptimizationResult incorporating learnings from all iterations
        """
        # Aggregate patterns across all analyses
        all_patterns = []
        all_weaknesses = []
        all_recommendations = []
        avg_score = 0

        for analysis in analyses:
            all_patterns.extend(analysis.patterns)
            all_weaknesses.extend(analysis.weaknesses)
            all_recommendations.extend(analysis.recommendations)
            avg_score += analysis.overall_score

        avg_score /= len(analyses) if analyses else 1

        # Create aggregated analysis
        aggregated = AnalysisResult(
            trace_id="aggregated",
            patterns=all_patterns,
            overall_score=avg_score,
            weaknesses=list(set(all_weaknesses)),  # Deduplicate
            recommendations=list(set(all_recommendations)),
        )

        # Optimize based on aggregated analysis
        return self.optimize(
            original_prompt=original_prompt,
            analysis=aggregated,
            trace=traces[0] if traces else None,
        )

    def suggest_tool_improvements(
        self,
        tools: list[dict[str, Any]],
        analysis: AnalysisResult,
        trace: ReasoningTrace,
    ) -> dict[str, str]:
        """
        Suggest improvements for tool definitions based on analysis.

        Args:
            tools: Original tool definitions
            analysis: Analysis results
            trace: Original reasoning trace

        Returns:
            Dict mapping tool names to suggested description improvements
        """
        tool_issues = [
            p for p in analysis.patterns
            if p.type.value in ("tool_confusion", "tool_misuse")
        ]

        if not tool_issues:
            return {}

        prompt = f"""Based on these tool usage issues:

{self._format_patterns_for_tools(tool_issues)}

And the original tool definitions:
{json.dumps(tools, indent=2)}

Suggest improved tool descriptions. Respond as JSON:
```json
{{
    "tool_name": "improved description that addresses the confusion"
}}
```"""

        response = self.client.messages.create(
            model=self.model,
            max_tokens=2048,
            messages=[{"role": "user", "content": prompt}],
        )

        for block in response.content:
            if block.type == "text":
                try:
                    text = block.text
                    if "```json" in text:
                        text = text.split("```json")[1].split("```")[0]
                    return json.loads(text)
                except json.JSONDecodeError:
                    pass

        return {}

    def _format_patterns(self, analysis: AnalysisResult) -> str:
        """Format patterns for optimization prompt."""
        if not analysis.patterns:
            return "No significant patterns detected."

        parts = []
        for p in analysis.patterns:
            parts.append(
                f"[{p.severity.value.upper()}] {p.type.value}\n"
                f"  Description: {p.description}\n"
                f"  Evidence: {', '.join(p.evidence[:2])}\n"
                f"  Suggestion: {p.suggestion}"
            )
        return "\n\n".join(parts)

    def _format_patterns_for_tools(self, patterns: list) -> str:
        """Format tool-related patterns."""
        return "\n".join(
            f"- {p.type.value}: {p.description}" for p in patterns
        )

    def _parse_optimization_response(
        self,
        response_text: str,
        original_prompt: str,
    ) -> OptimizationResult:
        """Parse the JSON optimization response with fallback extraction."""
        result = OptimizationResult(
            original_prompt=original_prompt,
            optimized_prompt=original_prompt,  # Default to original if parsing fails
        )

        try:
            json_text = response_text
            if "```json" in response_text:
                json_text = response_text.split("```json")[1].split("```")[0]
            elif "```" in response_text:
                json_text = response_text.split("```")[1].split("```")[0]

            data = json.loads(json_text)

            result.optimized_prompt = data.get("optimized_prompt", original_prompt)
            result.predicted_improvement = data.get("predicted_improvement", 0)
            result.confidence = data.get("confidence", 0.5)
            result.key_changes = data.get("key_changes", [])

            # Parse diffs
            for d in data.get("diffs", []):
                diff = PromptDiff(
                    section=d.get("section", "unknown"),
                    original=d.get("original", ""),
                    optimized=d.get("optimized", ""),
                    reason=d.get("reason", ""),
                )
                result.diffs.append(diff)

        except (json.JSONDecodeError, KeyError) as e:
            # Fallback: try to extract optimized_prompt directly from response
            extracted_prompt = self._fallback_extract_prompt(response_text)
            if extracted_prompt and extracted_prompt != original_prompt:
                result.optimized_prompt = extracted_prompt
                result.key_changes = [f"JSON parsing failed ({type(e).__name__}), extracted prompt via fallback"]
                result.confidence = 0.3  # Lower confidence for fallback extraction
            else:
                result.key_changes = [f"Optimization parsing failed ({type(e).__name__}) - using original prompt"]

        return result

    def _fallback_extract_prompt(self, response_text: str) -> str | None:
        """
        Fallback method to extract optimized prompt when JSON parsing fails.

        Tries multiple strategies to find the prompt content.
        """
        import re

        # Strategy 1: Look for "optimized_prompt": "..." pattern
        match = re.search(r'"optimized_prompt"\s*:\s*"([^"]+)"', response_text, re.DOTALL)
        if match:
            # Unescape the string
            return match.group(1).replace('\\n', '\n').replace('\\"', '"')

        # Strategy 2: Look for content between specific markers
        markers = [
            ('## Optimized Prompt', '##'),
            ('**Optimized Prompt**', '**'),
            ('OPTIMIZED PROMPT:', '\n\n'),
            ('Here is the improved prompt:', '\n\n---'),
        ]

        for start_marker, end_marker in markers:
            if start_marker in response_text:
                start_idx = response_text.find(start_marker) + len(start_marker)
                remaining = response_text[start_idx:].strip()
                if end_marker in remaining:
                    end_idx = remaining.find(end_marker)
                    extracted = remaining[:end_idx].strip()
                    if len(extracted) > 50:  # Minimum length check
                        return extracted

        # Strategy 3: Look for a substantial code block that might be the prompt
        code_blocks = re.findall(r'```(?:text|markdown)?\n(.*?)```', response_text, re.DOTALL)
        for block in code_blocks:
            # Skip JSON blocks, look for prose blocks that could be prompts
            if not block.strip().startswith('{') and len(block) > 100:
                return block.strip()

        return None


def format_optimization_report(result: OptimizationResult) -> str:
    """Format an optimization result as a human-readable report."""
    lines = [
        "=" * 60,
        "PROMPT OPTIMIZATION REPORT",
        "=" * 60,
        "",
        f"Predicted Improvement: {result.predicted_improvement}%",
        f"Confidence: {result.confidence * 100:.0f}%",
        "",
    ]

    if result.key_changes:
        lines.append("Key Changes:")
        for change in result.key_changes:
            lines.append(f"  - {change}")
        lines.append("")

    if result.diffs:
        lines.append("Detailed Changes:")
        for diff in result.diffs:
            lines.append(f"\n  [{diff.section}]")
            if diff.original and diff.original != "N/A":
                lines.append(f"    Before: {diff.original[:100]}...")
            lines.append(f"    After: {diff.optimized[:100]}...")
            lines.append(f"    Reason: {diff.reason}")

    lines.extend([
        "",
        "=" * 60,
        "OPTIMIZED PROMPT",
        "=" * 60,
        result.optimized_prompt,
    ])

    return "\n".join(lines)