"""
Direct Agent - Directly call the large model through the prompt to generate itinerary
"""

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


class DirectAgent:
    """Agent that directly calls the large model"""
    
    def __init__(self, backbone_llm=None, **kwargs):
        """
        Initialize DirectAgent
        
        Args:
            backbone_llm: Backend LLM model
            **kwargs: Other parameters
        """
        self.backbone_llm = backbone_llm
        self.method = "Direct"
        
    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]]:
        """
        Directly run LLM 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", [])
            
            # 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
                    }
                ]
            
            print(f"  Direct agent: Calling LLM with {len(messages)} messages")
            
            # Directly call LLM
            if self.backbone_llm:
                response = self.backbone_llm(messages, one_line=False, json_mode=False)
            else:
                response = "Mock response: This is a test response from DirectAgent"
            
            # Build result
            result = {
                "method": "Direct",
                "llm_response": response,
                "messages": messages,
                "status": "completed"
            }
            
            print(f"  Direct agent: Generated response ({len(str(response))} chars)")
            
            return True, result
            
        except Exception as e:
            print(f"  Direct agent error: {e}")
            return False, {
                "error": str(e),
                "method": "Direct",
                "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) 