#!/usr/bin/env -S uv run --script
# /// script
# requires-python = ">=3.11"
# dependencies = ["kodelet-sdk"]
# [tool.uv.sources]
# kodelet-sdk = { path = "../..", editable = true }
# ///
from __future__ import annotations

import asyncio
import os
import sys

from kodelet_sdk import BaseModel, Client, Extension

DEFAULT_MESSAGE = "Use calculator to add 123 and 456"


class CalculatorInput(BaseModel):
    a: int
    b: int


ext = Extension(name="calculator", version="0.1.0")


@ext.tool("calculator", description="Add two integers", input_schema=CalculatorInput)
async def calculator(input: CalculatorInput) -> str:
    print(f"Local calculator called: {input.a} + {input.b}", file=sys.stderr)
    return str(input.a + input.b)


async def main() -> None:
    message = " ".join(sys.argv[1:]).strip() or DEFAULT_MESSAGE
    command = os.environ.get("KODELET_BIN", "kodelet")
    client = Client(command=command)

    try:
        session = await client.create_session(extensions=[ext])
        response = await session.run_and_wait(message=message)
        print(response.content)
        print(f"\nConversation: {response.conversation_id or session.id}", file=sys.stderr)
    finally:
        await client.close()


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