"""
Simple tool-calling chat with ExCom AI
Works around vLLM's tool_choice requirements
"""

from openai import OpenAI
import json
from datetime import datetime
import math

# Configure OpenAI client for your vLLM endpoint
client = OpenAI(
    base_url="https://plarnholt-excom-ai-demo.hf.space/v1",
    api_key="not-needed"
)

# Define tools
tools = [
    {
        "type": "function",
        "function": {
            "name": "calculator",
            "description": "Evaluates a mathematical expression",
            "parameters": {
                "type": "object",
                "properties": {
                    "expression": {
                        "type": "string",
                        "description": "Math expression to evaluate, e.g., '2 + 2 * 3'"
                    }
                },
                "required": ["expression"]
            }
        }
    },
    {
        "type": "function",
        "function": {
            "name": "get_current_time",
            "description": "Returns the current date and time",
            "parameters": {
                "type": "object",
                "properties": {}
            }
        }
    },
    {
        "type": "function",
        "function": {
            "name": "get_weather",
            "description": "Gets the weather for a city (simulated)",
            "parameters": {
                "type": "object",
                "properties": {
                    "city": {
                        "type": "string",
                        "description": "City name"
                    }
                },
                "required": ["city"]
            }
        }
    },
    {
        "type": "function",
        "function": {
            "name": "word_counter",
            "description": "Counts the number of words in the given text",
            "parameters": {
                "type": "object",
                "properties": {
                    "text": {
                        "type": "string",
                        "description": "Text to count words in"
                    }
                },
                "required": ["text"]
            }
        }
    },
    {
        "type": "function",
        "function": {
            "name": "reverse_text",
            "description": "Reverses the given text",
            "parameters": {
                "type": "object",
                "properties": {
                    "text": {
                        "type": "string",
                        "description": "Text to reverse"
                    }
                },
                "required": ["text"]
            }
        }
    }
]

# Tool implementations
def calculator(expression: str) -> str:
    """Evaluates a mathematical expression"""
    try:
        result = eval(expression, {"__builtins__": {}, "math": math})
        return str(result)
    except Exception as e:
        return f"Error: {str(e)}"

def get_current_time() -> str:
    """Returns the current date and time"""
    return datetime.now().strftime("%Y-%m-%d %H:%M:%S")

def get_weather(city: str) -> str:
    """Gets simulated weather for a city"""
    weather_data = {
        "paris": "18°C, sunny",
        "london": "15°C, cloudy",
        "new york": "22°C, partly cloudy",
        "tokyo": "25°C, clear",
        "berlin": "16°C, rainy",
        "sydney": "23°C, sunny",
        "mumbai": "30°C, hot and humid",
        "toronto": "12°C, cool",
    }
    return weather_data.get(city.lower(), f"Weather data not available for {city}")

def word_counter(text: str) -> str:
    """Counts words in text"""
    words = text.split()
    return f"Word count: {len(words)}"

def reverse_text(text: str) -> str:
    """Reverses text"""
    return text[::-1]

# Function dispatcher
def execute_tool(tool_name: str, arguments: dict) -> str:
    """Execute a tool by name with given arguments"""
    if tool_name == "calculator":
        return calculator(arguments["expression"])
    elif tool_name == "get_current_time":
        return get_current_time()
    elif tool_name == "get_weather":
        return get_weather(arguments["city"])
    elif tool_name == "word_counter":
        return word_counter(arguments["text"])
    elif tool_name == "reverse_text":
        return reverse_text(arguments["text"])
    else:
        return f"Unknown tool: {tool_name}"

def chat(user_message: str, messages: list = None):
    """Send a message and handle tool calls"""
    if messages is None:
        messages = []

    # Add user message
    messages.append({"role": "user", "content": user_message})

    try:
        # Call the model with tools (no tool_choice parameter)
        response = client.chat.completions.create(
            model="excom-ai",
            messages=messages,
            tools=tools,
            temperature=0.4
        )

        assistant_message = response.choices[0].message

        # Check if model wants to use tools
        if assistant_message.tool_calls:
            # Add assistant's tool call request to messages
            messages.append({
                "role": "assistant",
                "content": assistant_message.content,
                "tool_calls": [
                    {
                        "id": tc.id,
                        "type": "function",
                        "function": {
                            "name": tc.function.name,
                            "arguments": tc.function.arguments
                        }
                    }
                    for tc in assistant_message.tool_calls
                ]
            })

            # Execute each tool call
            for tool_call in assistant_message.tool_calls:
                function_name = tool_call.function.name
                function_args = json.loads(tool_call.function.arguments)

                print(f"🔧 Calling tool: {function_name}({function_args})")

                # Execute the tool
                tool_result = execute_tool(function_name, function_args)

                # Add tool result to messages
                messages.append({
                    "role": "tool",
                    "tool_call_id": tool_call.id,
                    "name": function_name,
                    "content": tool_result
                })

            # Get final response from model
            final_response = client.chat.completions.create(
                model="excom-ai",
                messages=messages,
                temperature=0.4
            )

            return final_response.choices[0].message.content, messages
        else:
            # No tools needed, return direct response
            return assistant_message.content, messages

    except Exception as e:
        return f"Error: {e}", messages

def main():
    print("=" * 60)
    print("ExCom AI - Simple Tool Calling Chat")
    print("=" * 60)
    print("Available tools:")
    print("  • calculator - Evaluate math expressions")
    print("  • get_current_time - Get current date/time")
    print("  • get_weather - Get weather for cities")
    print("  • word_counter - Count words in text")
    print("  • reverse_text - Reverse text")
    print("\nType 'quit', 'exit', or 'q' to end.")
    print("Type 'reset' to clear conversation history.")
    print("=" * 60)
    print()

    messages = []

    while True:
        try:
            user_input = input("You: ").strip()

            if user_input.lower() in ['quit', 'exit', 'q']:
                print("Goodbye!")
                break

            if user_input.lower() == 'reset':
                messages = []
                print("🔄 Conversation history cleared.\n")
                continue

            if not user_input:
                continue

            response, messages = chat(user_input, messages)
            print(f"Assistant: {response}\n")

        except KeyboardInterrupt:
            print("\n\nGoodbye!")
            break
        except Exception as e:
            print(f"❌ Error: {e}\n")
            # Reset messages on unexpected error
            messages = []

if __name__ == "__main__":
    main()