"""
Example usage of burt-logger with OpenAI API.

This example shows how to integrate burt-logger with OpenAI's API
to automatically log all LLM requests and responses for training data collection.
"""

import time
from burt_logger import LLMLogger

# Note: This example assumes you have openai installed
# pip install openai

try:
    from openai import OpenAI
except ImportError:
    print("This example requires the openai package. Install with: pip install openai")
    exit(1)


def main():
    # Initialize the logger
    logger = LLMLogger(
        endpoint="https://api.example.com/logs",  # Replace with your backend endpoint
        api_key="sk-your-backend-api-key",  # Replace with your API key
        batch_size=5,  # Send logs in batches of 5
        flush_interval=3.0,  # Or every 3 seconds
        debug=True,  # Enable debug logging to see what's happening
    )
    
    # Initialize OpenAI client
    client = OpenAI(api_key="your-openai-api-key")  # Replace with your OpenAI key
    
    print("Sending OpenAI requests and logging them...\n")
    
    # Example 1: Simple chat completion
    messages = [
        {"role": "system", "content": "You are a helpful assistant."},
        {"role": "user", "content": "What is the capital of France?"},
    ]
    
    try:
        response = client.chat.completions.create(
            model="gpt-3.5-turbo",
            messages=messages,
            temperature=0.7,
        )
        
        # Log the request and response
        logger.log(
            request={
                "provider": "openai",
                "model": "gpt-3.5-turbo",
                "messages": messages,
                "temperature": 0.7,
            },
            response={
                "id": response.id,
                "model": response.model,
                "choices": [
                    {
                        "message": {
                            "role": choice.message.role,
                            "content": choice.message.content,
                        },
                        "finish_reason": choice.finish_reason,
                    }
                    for choice in response.choices
                ],
                "usage": {
                    "prompt_tokens": response.usage.prompt_tokens,
                    "completion_tokens": response.usage.completion_tokens,
                    "total_tokens": response.usage.total_tokens,
                },
            },
            metadata={
                "user_id": "user_123",
                "session_id": "session_456",
                "environment": "production",
            },
        )
        
        print(f"Response: {response.choices[0].message.content}")
        print(f"Tokens used: {response.usage.total_tokens}\n")
        
    except Exception as e:
        print(f"Error calling OpenAI API: {e}\n")
    
    # Example 2: Multiple requests
    print("Sending multiple requests to demonstrate batching...\n")
    
    questions = [
        "What is 2 + 2?",
        "What is the largest planet in our solar system?",
        "Who wrote Romeo and Juliet?",
    ]
    
    for i, question in enumerate(questions, 1):
        try:
            messages = [{"role": "user", "content": question}]
            
            response = client.chat.completions.create(
                model="gpt-3.5-turbo",
                messages=messages,
            )
            
            logger.log(
                request={
                    "provider": "openai",
                    "model": "gpt-3.5-turbo",
                    "messages": messages,
                },
                response={
                    "choices": [
                        {"message": {"content": choice.message.content}}
                        for choice in response.choices
                    ],
                    "usage": {
                        "total_tokens": response.usage.total_tokens,
                    },
                },
            )
            
            print(f"Q{i}: {question}")
            print(f"A{i}: {response.choices[0].message.content}\n")
            
        except Exception as e:
            print(f"Error on question {i}: {e}\n")
    
    # Get statistics
    print("Logger Statistics:")
    stats = logger.get_stats()
    for key, value in stats.items():
        print(f"  {key}: {value}")
    
    # Gracefully shutdown and flush remaining logs
    print("\nShutting down logger and flushing remaining logs...")
    logger.shutdown()
    print("Done!")


if __name__ == "__main__":
    main()

