"""Prompt management system for different evaluation strategies."""

import os
import json
from pathlib import Path
from typing import Dict, List, Optional, Any
from enum import Enum


class PromptStrategy(Enum):
    """Different prompting strategies."""
    BASELINE = "baseline"
    CHAIN_OF_THOUGHT = "cot"
    FEW_SHOT = "few_shot"
    ZERO_SHOT = "zero_shot"
    CUSTOM = "custom"


class PromptManager:
    """Manager for loading and organizing prompts."""
    
    def __init__(self, prompt_directory: Optional[str] = None):
        """Initialize prompt manager.
        
        Args:
            prompt_directory: Directory containing prompt files
        """
        self.prompt_directory = Path(prompt_directory).expanduser() if prompt_directory else None
        self.prompts: Dict[str, Dict[str, str]] = {}
        self._load_default_prompts()
        
        if self.prompt_directory and self.prompt_directory.exists():
            self._load_prompts_from_directory()
    
    def _load_default_prompts(self):
        """Load default prompts."""
        self.prompts[PromptStrategy.BASELINE.value] = {
            "system": """You are a helpful assistant that can answer questions and help with tasks.

## Rules
1. You need to act as an agent that uses tools to help the user answer the question.
2. Make sure you get all the information you need.
3. Once you have all the information, think step by step and answer the question.
4. If the question is not solvable, you should say "I don't know".

## Format
Your final answer should start with "#### ". If you want to answer 27, please say "#### 27"."""
        }
        
        self.prompts[PromptStrategy.CHAIN_OF_THOUGHT.value] = {
            "system": """You are a problem-solving agent. Your mission is to solve math word problems by finding all necessary facts using search tools.

## Strategy: The Research Loop
Think of your work as a loop: Analyze -> Search -> Attempt to Solve. Repeat until you have enough information.

Step 1: Analyze the Gaps
- What is the main question?
- What facts do I have?
- What crucial information is missing?

Step 2: Plan and Execute Your Search
- Use search(query) for new searches
- Use next_page() to get more results from the last search
- Balance between exploring new queries and getting more results

Step 3: Attempt to Solve
- Try to calculate the answer with gathered information
- Show your work step by step
- If stuck, identify what's missing and return to Step 1

## Final Answer Format
Your final answer must begin with #### and contain only the numerical solution.
Example: #### 42"""
        }
        
        self.prompts[PromptStrategy.ZERO_SHOT.value] = {
            "system": """You are an AI assistant. Answer the question using the available tools.

Your final answer should start with "#### " followed by the answer."""
        }
    
    def _load_prompts_from_directory(self):
        """Load prompts from directory structure."""
        if not self.prompt_directory:
            return
        
        # Look for prompt files in the directory
        for strategy in PromptStrategy:
            strategy_file = self.prompt_directory / f"{strategy.value}.txt"
            if strategy_file.exists():
                with open(strategy_file, 'r') as f:
                    self.prompts[strategy.value] = {"system": f.read().strip()}
            
            # Also check for JSON format with multiple prompts
            json_file = self.prompt_directory / f"{strategy.value}.json"
            if json_file.exists():
                with open(json_file, 'r') as f:
                    self.prompts[strategy.value] = json.load(f)
    
    def get_prompt(self, strategy: str = "baseline", prompt_type: str = "system") -> str:
        """Get a specific prompt.
        
        Args:
            strategy: The prompting strategy to use
            prompt_type: Type of prompt (system, user, etc.)
            
        Returns:
            The prompt text
        """
        if strategy not in self.prompts:
            strategy = PromptStrategy.BASELINE.value
        
        prompts = self.prompts.get(strategy, {})
        return prompts.get(prompt_type, "")
    
    def list_strategies(self) -> List[str]:
        """List available prompt strategies."""
        return list(self.prompts.keys())
    
    def add_custom_prompt(self, name: str, prompt: str, prompt_type: str = "system"):
        """Add a custom prompt.
        
        Args:
            name: Name for the custom prompt
            prompt: The prompt text
            prompt_type: Type of prompt
        """
        if name not in self.prompts:
            self.prompts[name] = {}
        self.prompts[name][prompt_type] = prompt
    
    def save_prompts(self, output_dir: str):
        """Save all prompts to a directory.
        
        Args:
            output_dir: Directory to save prompts to
        """
        output_path = Path(output_dir).expanduser()
        output_path.mkdir(parents=True, exist_ok=True)
        
        for strategy, prompts in self.prompts.items():
            # Save as JSON for complex prompts
            json_file = output_path / f"{strategy}.json"
            with open(json_file, 'w') as f:
                json.dump(prompts, f, indent=2)
            
            # Also save system prompt as text for easy editing
            if "system" in prompts:
                text_file = output_path / f"{strategy}.txt"
                with open(text_file, 'w') as f:
                    f.write(prompts["system"])
