"""
CoT (Chain of Thought) Agent - Generate itinerary by enhancing the reasoning process
Based on Direct method, add "Let's think step by step" to the prompt to enhance the reasoning process
"""

import json
from typing import Dict, Any, Tuple, List


class CoTAgent:
    """Agent based on Chain of Thought reasoning"""
    
    def __init__(self, backbone_llm=None, **kwargs):
        """
        Initialize CoTAgent
        
        Args:
            backbone_llm: Backend LLM model
            **kwargs: Other parameters
        """
        self.backbone_llm = backbone_llm
        self.method = "CoT"
        
    def _enhance_messages_with_cot(self, messages: List[Dict[str, str]], locale: str) -> List[Dict[str, str]]:
        """
        Add CoT prompt to messages to enhance the reasoning process
        
        Args:
            messages: Original message list
            locale: Language
        Returns:
            List[Dict[str, str]]: Enhanced message list
        """
        if not messages:
            return messages
            
        # Find the last user message
        last_user_message = None
        for i in range(len(messages) - 1, -1, -1):
            if messages[i].get("role") == "user":
                last_user_message = messages[i]
                break
        
        if last_user_message:
            # Add CoT prompt before the user message
            if locale.split("-")[0] == "zh":
                cot_prompt = "让我们一步一步地思考。"
            else:
                cot_prompt = "Let's think step by step."
            if not last_user_message["content"].endswith(cot_prompt):
                last_user_message["content"] = last_user_message["content"] + "\n" + cot_prompt
        
        return messages
        
    def run(self, query_data: Dict[str, Any], load_cache: bool = True, 
            oralce_translation: bool = False, preference_search: bool = False) -> Tuple[bool, Dict[str, Any]]:
        """
        Run CoT agent to generate itinerary
        
        Args:
            query_data: Query data
            load_cache: Whether to load cache (not implemented yet)
            oralce_translation: Whether to use oracle translation
            preference_search: Whether to perform preference search
            
        Returns:
            Tuple[bool, Dict]: (Whether successful, plan data)
        """
        try:
            # Extract messages
            messages = query_data.get("messages", [])
            locale = query_data.get("locale", "en-US")
            
            # Ensure messages is a list format
            if isinstance(messages, str):
                try:
                    messages = json.loads(messages)
                except json.JSONDecodeError:
                    print(f"  Warning: Failed to parse messages string, creating basic message")
                    messages = []
            
            if not messages or not isinstance(messages, list):
                # If there are no messages or the format is incorrect, build a basic prompt
                user_query = query_data.get("userQuery", "")
                if not user_query:
                    return False, {"error": "No user query or messages provided"}
                
                messages = [
                    {
                        "role": "user", 
                        "content": user_query
                    }
                ]
            
            # Use CoT to enhance messages
            enhanced_messages = self._enhance_messages_with_cot(messages, locale)
            
            print(f"  CoT agent: Calling LLM with {len(enhanced_messages)} enhanced messages")
            print(f"  CoT agent: Enhanced prompt with 'Let's think step by step'")
            
            # Call LLM
            if self.backbone_llm:
                response = self.backbone_llm(enhanced_messages, one_line=False, json_mode=False)
            else:
                response = "Mock response: This is a test response from CoTAgent with step-by-step reasoning"
            
            # Build result
            result = {
                "method": "CoT",
                "llm_response": response,
                "messages": messages,
                "enhanced_messages": enhanced_messages,  # Save original messages for comparison
                "cot_enhanced": True,
                "status": "completed"
            }
            
            print(f"  CoT agent: Generated response ({len(str(response))} chars)")
            
            return True, result
            
        except Exception as e:
            print(f"  CoT agent error: {e}")
            return False, {
                "error": str(e),
                "method": "CoT",
                "status": "failed"
            }
    
    def solve(self, query_data: Dict[str, Any], prob_idx: str = None, 
              oracle_verifier: bool = True) -> Tuple[bool, Dict[str, Any]]:
        """
        Compatible with LLM-modulo's solve interface
        
        Args:
            query_data: Query data
            prob_idx: Problem index
            oracle_verifier: Whether to use oracle verifier
            
        Returns:
            Tuple[bool, Dict]: (Whether successful, plan data)
        """
        return self.run(query_data) 