#!/usr/bin/env python3
"""
Codex Integration Example for Qwen3-4B Tool Calling
This script demonstrates how to use the model with Codex
"""

import requests
import json
import time
from typing import List, Dict, Any, Optional

class CodexClient:
    """Client for interacting with Codex-compatible server"""
    
    def __init__(self, base_url: str = "http://localhost:8000"):
        self.base_url = base_url
        self.session = requests.Session()
        self.model_name = "Qwen3-4B-Function-Calling-Pro"
    
    def chat_completion(self, messages: List[Dict[str, str]], tools: Optional[List[Dict]] = None, 
                       temperature: float = 0.7, max_tokens: int = 512) -> Dict[str, Any]:
        """Send chat completion request to Codex server"""
        payload = {
            "model": self.model_name,
            "messages": messages,
            "temperature": temperature,
            "max_tokens": max_tokens,
            "stop": ["<|im_end|>", "<|im_start|>", "<tool_call>", "</tool_call>"]
        }
        
        if tools:
            payload["tools"] = tools
        
        try:
            response = self.session.post(
                f"{self.base_url}/v1/chat/completions",
                json=payload,
                headers={"Content-Type": "application/json"},
                timeout=60
            )
            response.raise_for_status()
            return response.json()
        except requests.exceptions.RequestException as e:
            print(f"Error making request: {e}")
            return {"error": str(e)}
    
    def extract_tool_calls(self, response: Dict[str, Any]) -> List[Dict[str, Any]]:
        """Extract tool calls from Codex response"""
        tool_calls = []
        if "choices" in response and len(response["choices"]) > 0:
            message = response["choices"][0]["message"]
            if "tool_calls" in message:
                tool_calls = message["tool_calls"]
        return tool_calls
    
    def is_server_running(self) -> bool:
        """Check if the Codex server is running"""
        try:
            response = self.session.get(f"{self.base_url}/health", timeout=5)
            return response.status_code == 200
        except:
            return False

def demo_codex_integration():
    """Demonstrate Codex integration"""
    print("🔌 Codex Integration Demo")
    print("=" * 50)
    
    # Initialize client
    codex = CodexClient()
    
    # Check if server is running
    if not codex.is_server_running():
        print("❌ Codex server is not running!")
        print("💡 Start the server with: ./run_model.sh server")
        return
    
    print("✅ Connected to Codex server")
    
    # Define tools for Codex
    tools = [
        {
            "type": "function",
            "function": {
                "name": "get_weather",
                "description": "Get current weather for a location",
                "parameters": {
                    "type": "object",
                    "properties": {
                        "location": {
                            "type": "string",
                            "description": "City name or location"
                        }
                    },
                    "required": ["location"]
                }
            }
        },
        {
            "type": "function",
            "function": {
                "name": "search_hotels",
                "description": "Search for hotels in a city",
                "parameters": {
                    "type": "object",
                    "properties": {
                        "city": {
                            "type": "string",
                            "description": "City name"
                        },
                        "check_in": {
                            "type": "string",
                            "description": "Check-in date (YYYY-MM-DD)"
                        },
                        "check_out": {
                            "type": "string",
                            "description": "Check-out date (YYYY-MM-DD)"
                        }
                    },
                    "required": ["city"]
                }
            }
        },
        {
            "type": "function",
            "function": {
                "name": "calculate",
                "description": "Perform mathematical calculations",
                "parameters": {
                    "type": "object",
                    "properties": {
                        "expression": {
                            "type": "string",
                            "description": "Mathematical expression to evaluate"
                        }
                    },
                    "required": ["expression"]
                }
            }
        }
    ]
    
    # Test cases
    test_cases = [
        "What's the weather like in London?",
        "Find me a hotel in Paris for next week",
        "Calculate 25 + 17",
        "Get weather for New York and find hotels there"
    ]
    
    for i, message in enumerate(test_cases, 1):
        print(f"\n📝 Test {i}: {message}")
        print("-" * 40)
        
        # Send request
        messages = [{"role": "user", "content": message}]
        response = codex.chat_completion(messages, tools=tools)
        
        if "error" in response:
            print(f"❌ Error: {response['error']}")
            continue
        
        # Extract response
        if "choices" in response and len(response["choices"]) > 0:
            assistant_message = response["choices"][0]["message"]
            print(f"Response: {assistant_message.get('content', 'No content')}")
            
            # Check for tool calls
            tool_calls = codex.extract_tool_calls(response)
            if tool_calls:
                print(f"\n🔧 Tool Calls ({len(tool_calls)}):")
                for j, tool_call in enumerate(tool_calls, 1):
                    print(f"  {j}. {tool_call['function']['name']}")
                    print(f"     Arguments: {tool_call['function'].get('arguments', {})}")
            else:
                print("\n❌ No tool calls detected")
        else:
            print("❌ No response received")

def interactive_codex_chat():
    """Interactive chat with Codex server"""
    print("💬 Interactive Codex Chat")
    print("=" * 50)
    
    codex = CodexClient()
    
    if not codex.is_server_running():
        print("❌ Codex server is not running!")
        print("💡 Start the server with: ./run_model.sh server")
        return
    
    print("✅ Connected to Codex server")
    print("Type 'quit' to exit")
    print("-" * 50)
    
    while True:
        try:
            user_input = input("\nYou: ").strip()
            if user_input.lower() in ['quit', 'exit', 'q']:
                break
            
            if not user_input:
                continue
            
            # Send request
            messages = [{"role": "user", "content": user_input}]
            response = codex.chat_completion(messages)
            
            if "error" in response:
                print(f"❌ Error: {response['error']}")
                continue
            
            # Display response
            if "choices" in response and len(response["choices"]) > 0:
                assistant_message = response["choices"][0]["message"]
                print(f"\nAssistant: {assistant_message.get('content', 'No content')}")
                
                # Check for tool calls
                tool_calls = codex.extract_tool_calls(response)
                if tool_calls:
                    print(f"\n🔧 Tool Calls ({len(tool_calls)}):")
                    for i, tool_call in enumerate(tool_calls, 1):
                        print(f"  {i}. {tool_call['function']['name']}")
                        print(f"     Arguments: {tool_call['function'].get('arguments', {})}")
        
        except KeyboardInterrupt:
            print("\n\nGoodbye! 👋")
            break
        except Exception as e:
            print(f"Error: {e}")

if __name__ == "__main__":
    import sys
    
    if len(sys.argv) > 1 and sys.argv[1] == "demo":
        demo_codex_integration()
    else:
        interactive_codex_chat()