import json
from dotenv import load_dotenv
from projectdavid import Entity, ToolCallRequestEvent

load_dotenv()

# Universal Function call handler
def get_flight_times(tool_name: str, arguments: dict) -> str:
    return json.dumps({
        "status": "success",
        "departure": arguments.get("departure", "UNK"),
        "arrival": arguments.get("arrival", "UNK"),
        "duration": "4h 30m",
    })

TOOL_REGISTRY = {
    "get_flight_times": get_flight_times,
}

client = Entity()

# 1. Create the Assistant and attach tools
assistant = client.assistants.create_assistant(
    name="Test Assistant",
    model="gpt-oss-120b",
    instructions="You are a helpful AI assistant, your name is Nexa.",
    tools=[{
        "type": "function",
        "function": {
            "name": "get_flight_times",
            "description": "Get flight times",
            "parameters": {
                "type": "object",
                "properties": {
                    "departure": {"type": "string"},
                    "arrival": {"type": "string"},
                },
                "required": ["departure", "arrival"],
            },
        },
    }],
)

thread = client.threads.create_thread()

message = client.messages.create_message(
    thread_id=thread.id,
    role="user",
    content="Please provide the flight times for a trip departing from Tokyo and arriving in Sydney?",
    assistant_id=assistant.id,
)

run = client.runs.create_run(assistant_id=assistant.id, thread_id=thread.id)

# 2. Setup Streaming
stream = client.synchronous_inference_stream
stream.setup(
    thread_id=thread.id,
    assistant_id=assistant.id,
    message_id=message.id,
    run_id=run.id,
)

# 3. Stream events via Ollama and handle tool execution dynamically
for event in stream.stream_events(model="ollama/qwen3:4b"):
    print(f"{event.__class__.__name__:<25} | {json.dumps(event.to_dict())}")
    
    if isinstance(event, ToolCallRequestEvent):
        handler = TOOL_REGISTRY.get(event.tool_name)
        if handler:
            print(f"\n[TOOL CALL] {event.tool_name} → executing...")
            event.execute(handler)
            print("[✓] Result submitted — resuming stream...\n")
        else:
            print(f"\n[!] Unknown tool: '{event.tool_name}'\n")