"""Utility functions for Deep Research workflow integration."""

import os
from datetime import datetime
from typing import List, Optional, Callable
from langchain_core.messages import MessageLikeRepresentation, AIMessage, ToolMessage
from langchain_core.tools import tool

# Simple think tool for strategic reflection
@tool
def think_tool(reflection: str) -> str:
    """Tool for strategic reflection on research progress and decision-making.

    Use this tool after each search to analyze results and plan next steps systematically.
    This creates a deliberate pause in the research workflow for quality decision-making.

    When to use:
    - After receiving search results: What key information did I find?
    - Before deciding next steps: Do I have enough to answer comprehensively?
    - When assessing research gaps: What specific information am I still missing?
    - Before concluding research: Can I provide a complete answer now?

    Reflection should address:
    1. Analysis of current findings - What concrete information have I gathered?
    2. Gap assessment - What crucial information is still missing?
    3. Quality evaluation - Do I have sufficient evidence/examples for a good answer?
    4. Strategic decision - Should I continue searching, backtrack, or provide my answer?

    Args:
        reflection: Your detailed reflection on research progress, findings, gaps, and next steps

    Returns:
        Confirmation that reflection was recorded for decision-making
    """
    return f"Reflection recorded: {reflection}"

def get_today_str() -> str:
    """Get today's date as a formatted string."""
    return datetime.now().strftime("%Y-%m-%d")

def get_api_key_for_model(model_name: str, config: dict) -> Optional[str]:
    """Get the appropriate API key for a model.
    
    Args:
        model_name: Name of the model (e.g., "gpt-4o", "claude-3-sonnet")
        config: Configuration dictionary that may contain API keys
        
    Returns:
        API key string or None
    """
    # Extract provider from model name
    if "gpt" in model_name.lower() or "openai" in model_name.lower():
        return os.environ.get("OPENAI_API_KEY")
    elif "claude" in model_name.lower() or "anthropic" in model_name.lower():
        return os.environ.get("ANTHROPIC_API_KEY")
    else:
        # Try common environment variables
        return os.environ.get("OPENAI_API_KEY") or os.environ.get("ANTHROPIC_API_KEY")

def get_model_token_limit(model_name: str) -> Optional[int]:
    """Get the token limit for a model.
    
    Args:
        model_name: Name of the model
        
    Returns:
        Token limit or None if unknown
    """
    model_limits = {
        "gpt-4o": 128000,
        "gpt-4o-mini": 128000,
        "gpt-4": 8192,
        "gpt-4-turbo": 128000,
        "gpt-5": 200000,
        "gpt-5-nano": 128000,
        "claude-3-sonnet": 200000,
        "claude-3-haiku": 200000,
        "claude-4-sonnet": 200000,
        "claude-4-sonnet-20250514": 200000,
    }
    
    # Try exact match first
    if model_name in model_limits:
        return model_limits[model_name]
    
    # Try partial matches
    for model_key, limit in model_limits.items():
        if model_key in model_name.lower():
            return limit
    
    # Default fallback
    return 8192

def is_token_limit_exceeded(error: Exception, model_name: str) -> bool:
    """Check if an error indicates token limit exceeded.
    
    Args:
        error: Exception that occurred
        model_name: Name of the model
        
    Returns:
        True if error indicates token limit exceeded
    """
    error_str = str(error).lower()
    token_limit_indicators = [
        "token limit",
        "context length",
        "too many tokens",
        "maximum context",
        "input is too long",
        "context_length_exceeded"
    ]
    
    return any(indicator in error_str for indicator in token_limit_indicators)

def remove_up_to_last_ai_message(messages: List[MessageLikeRepresentation]) -> List[MessageLikeRepresentation]:
    """Remove messages up to (but not including) the last AI message.
    
    This is used for token limit handling - keeps the most recent AI message
    and any subsequent tool/human messages.
    
    Args:
        messages: List of messages
        
    Returns:
        Truncated list of messages
    """
    if not messages:
        return messages
    
    # Find the last AI message
    last_ai_index = -1
    for i in range(len(messages) - 1, -1, -1):
        if isinstance(messages[i], AIMessage):
            last_ai_index = i
            break
    
    # If no AI message found, return last few messages
    if last_ai_index == -1:
        return messages[-5:] if len(messages) > 5 else messages
    
    # Return from last AI message onwards
    return messages[last_ai_index:]

def get_notes_from_tool_calls(messages: List[MessageLikeRepresentation]) -> List[str]:
    """Extract notes from tool call messages.
    
    Args:
        messages: List of messages that may contain tool calls
        
    Returns:
        List of extracted notes/content
    """
    notes = []
    
    for message in messages:
        if isinstance(message, ToolMessage):
            if message.content and isinstance(message.content, str):
                notes.append(message.content)
    
    return notes

# Note: get_search_tools function removed - now using external tools passed to workflow

async def get_all_tools(config: dict) -> List[Callable]:
    """Get all available tools for research (placeholder for MCP integration).
    
    Currently returns search tools only. Can be extended for MCP tools.
    
    Args:
        config: Configuration dictionary
        
    Returns:
        List of available tools
    """
    tools = []
    
    # Add basic search tools (will be replaced with actual search engine tools)
    # This is a placeholder - actual implementation will use the search engine
    # from the evaluation context
    
    # Add think tool
    tools.append(think_tool)
    
    return tools

def openai_websearch_called(message: MessageLikeRepresentation) -> bool:
    """Check if OpenAI native web search was called (placeholder).
    
    Args:
        message: Message to check
        
    Returns:
        False (not implemented for this integration)
    """
    return False

def anthropic_websearch_called(message: MessageLikeRepresentation) -> bool:
    """Check if Anthropic native web search was called (placeholder).
    
    Args:
        message: Message to check
        
    Returns:
        False (not implemented for this integration)
    """
    return False
