from functioncalming import get_completion
from pydantic import BaseModel, field_validator
import asyncio

def get_weather(city: str, country: str = "US") -> str:
    """Get current weather for a city"""
    return f"Sunny, 72°F in {city}, {country}"

class PersonInfo(BaseModel):
    """Extract person information from text"""
    name: str
    age: int
    occupation: str
    
    @field_validator('age')
    def validate_age(cls, v):
        if v < 0 or v > 150:
            raise ValueError("Age must be between 0 and 150")
        return v

async def main():
    # Function calling example
    response = await get_completion(
        user_message="What's the weather like in Paris?",
        tools=[get_weather],
        model="gpt-4o"
    )
    print(response.tool_call_results[0])  # "Sunny, 72°F in Paris, US"
    
    # Structured output example with validation
    response = await get_completion(
        user_message="Extract info: John is a 30-year-old teacher",
        tools=[PersonInfo],
        model="gpt-4o"
    )
    person = response.tool_call_results[0]
    print(f"{person.name}, age {person.age}, works as {person.occupation}")

if __name__ == "__main__":
    asyncio.run(main())

class EmailAddress(BaseModel):
    email: str
    
    @field_validator('email')
    def validate_email(cls, v):
        if '@' not in v:
            raise ValueError("Email must contain @ symbol")
        return v

response = await get_completion(
    user_message="My email is john.doe.gmail.com",  # Missing @
    tools=[EmailAddress],
    retries=2  # Will automatically retry up to 2 times
)

# response.messages contains clean history as if the model got it right first time
# response.messages_raw contains the actual history with retries
print(f"Retries needed: {response.retries_done}")
print(f"Final result: {response.tool_call_results[0].email}")

def calculate_tip(bill_amount: float, tip_percentage: float = 15.0) -> float:
    """Calculate tip amount"""
    return bill_amount * (tip_percentage / 100)

class Receipt(BaseModel):
    """Generate a structured receipt"""
    items: list[str]
    subtotal: float
    tip: float
    total: float

response = await get_completion(
    user_message="Calculate tip for $45.50 bill and create receipt for coffee and sandwich",
    tools=[calculate_tip, Receipt],
    model="gpt-4o"
)

response = await get_completion(
    user_message="Hello world",
    model="gpt-4o"
)

print(f"Cost: ${response.cost:.4f}")
print(f"Tokens used: {response.usage.total_tokens}")
print(f"Model: {response.model}")
print(f"Unknown costs: {response.unknown_costs}")  # True for very new models

# When you have many tools with complex schemas
response = await get_completion(
    user_message="Parse this document",
    tools=[ComplexTool1, ComplexTool2, ComplexTool3, ...],
    abbreviate_tools=True  # First shows tool names only, then full schema for chosen tool
)

response = await get_completion(
    user_message="Get weather for NYC and LA, then create a travel comparison",
    tools=[get_weather, TravelComparison],
    model="gpt-4o"
)

# response.tool_call_results contains results from all tool calls
for i, result in enumerate(response.tool_call_results):
    print(f"Tool {i+1} result: {result}")

class CodeGeneration(BaseModel):
    """Generate code with strict structure"""
    language: str
    code: str
    explanation: str

# Uses Structured Outputs automatically on compatible models
response = await get_completion(
    user_message="Write a Python function to reverse a string",
    tools=[CodeGeneration],
    model="gpt-4o"  # Structured Outputs supported
)

from openai import AsyncOpenAI
from functioncalming.client import set_openai_client

custom_client = AsyncOpenAI(api_key="your-key", timeout=30.0)

with set_openai_client(custom_client):
    response = await get_completion(
        user_message="Hello",
        tools=[SomeTool]
    )

from functioncalming.client import calm_middleware


@calm_middleware
async def log_requests(*, model, messages, tools, **kwargs):
    print(f"Making request to {model}")
    try:
        completion = yield  # OpenAI API call happens here
        print(f"Got response with {completion.usage.total_tokens} tokens")
    except Exception as e:
        print(f"Call to OpenAI failed: {e}")
        # Note: you should always re-raise here - of course, you can still wrap the call to get_completion in a try/except
        raise e
    


response = await get_completion(
    user_message="Hello",
    tools=[SomeTool],
    middleware=log_requests
)

from functioncalming.types import EscapedOutput

class Config(BaseModel):
    name: str
    debug_mode: bool = False  # Defaults work automatically

def generate_report():
    # Return complex data that shouldn't be shown to model
    return EscapedOutput(
        result_for_model="Report generated successfully",
        data={"complex": "internal", "data": [1, 2, 3]}
    )

response = await get_completion(
    user_message="Generate config and report",
    tools=[Config, generate_report]
)

# Model sees "Report generated successfully"
# You get the full data object
report_data = response.tool_call_results[1].data

from functioncalming.context import calm_context

def my_tool(param: str) -> str:
    context = calm_context.get()
    tool_call_id = context.tool_call.id
    function_name = context.tool_call.function.name
    return f"Called {function_name} with ID {tool_call_id}"

from functioncalming.client import register_model

register_model(
    model_name="gpt-new-model",
    supports_structured_outputs=True,
    cost_per_1mm_input_tokens=2.0,
    cost_per_1mm_output_tokens=8.0
)

from functioncalming.utils import ToolCallError

def strict_function(value: int) -> str:
    if value < 0:
        # This will trigger a retry
        raise ToolCallError("Value must be positive")
    return str(value)