#!/usr/bin/env python3
import asyncio
import os
import sys

sys.path.insert(0, "/app/src")

from claude import AgentOptions, ChatClient, ClaudeAgentClient, ToolResult, ToolUse

PROMPT = os.environ.get("TEST_PROMPT", "Say hi")

AGENT_CONFIGS = [
    {"label": "haiku", "model": "claude-haiku-4-5-20251001", "effort": None},
    {"label": "sonnet", "model": "claude-sonnet-4-6", "effort": None},
    {"label": "sonnet+effort=low", "model": "claude-sonnet-4-6", "effort": "low"},
    {"label": "sonnet+effort=med", "model": "claude-sonnet-4-6", "effort": "medium"},
    {"label": "sonnet+effort=high", "model": "claude-sonnet-4-6", "effort": "high"},
    {"label": "opus", "model": "claude-opus-4-6", "effort": None},
]


async def run_agent(cfg):
    opts = AgentOptions(model=cfg["model"], effort=cfg["effort"])
    async with ClaudeAgentClient(options=opts) as client:
        text = ""
        async for chunk in client.send_message_stream(PROMPT):
            if chunk.text_delta:
                text += chunk.text_delta
        print(f"  [{cfg['label']}] {text.strip()[:60]}")


async def run_chat():
    async with ChatClient() as c:
        # non-streaming
        r = await c.chat(PROMPT)
        print(f"  [chat]    {r.content[0]['text'].strip()[:60]}")

        # streaming collect
        text = await c.collect(PROMPT)
        print(f"  [collect] {text.strip()[:60]}")

        # raw stream (collect)
        parts = []
        async for chunk in c.stream(PROMPT):
            if chunk.text_delta:
                parts.append(chunk.text_delta)
        print(f"  [stream]  {''.join(parts).strip()[:60]}")

        # real-time streaming (print chunks as they arrive)
        sys.stdout.write("  [rt-strm] ")
        sys.stdout.flush()
        chunks = 0
        async for chunk in c.stream("Write a haiku about code"):
            if chunk.text_delta:
                chunks += 1
                sys.stdout.write(chunk.text_delta)
                sys.stdout.flush()
        print(f" ({chunks} chunks)")

        # system prompt override
        r = await c.chat(PROMPT, system="You are a pirate. Be brief.")
        print(f"  [system]  {r.content[0]['text'].strip()[:60]}")

        # multi-turn
        r = await c.chat(
            [
                {"role": "user", "content": "My name is Alice"},
                {"role": "assistant", "content": "Hi Alice!"},
                {"role": "user", "content": "What is my name?"},
            ]
        )
        print(f"  [multi]   {r.content[0]['text'].strip()[:60]}")

        # model override
        text = await c.collect(PROMPT, model="claude-haiku-4-5-20251001", max_tokens=256)
        print(f"  [haiku]   {text.strip()[:60]}")


CALC_TOOL = {
    "name": "calculator",
    "description": "Evaluates a mathematical expression and returns the result.",
    "input_schema": {
        "type": "object",
        "properties": {
            "expression": {
                "type": "string",
                "description": "The math expression to evaluate, e.g. '2 + 3 * 4'",
            }
        },
        "required": ["expression"],
    },
}


async def calculator_handler(name: str, input_data: dict) -> str:
    if name != "calculator":
        return f"Unknown tool: {name}"
    expr = input_data.get("expression", "")
    try:
        result = eval(expr, {"__builtins__": {}}, {})
        return str(result)
    except Exception as e:
        raise ValueError(f"calc error: {e}") from e


async def run_chat_tools():
    async with ChatClient() as c:
        tools = [CALC_TOOL]

        # --- Manual tool loop: call chat() with tools, inspect tool_calls ---
        r = await c.chat(
            "What is 7 * 13? Use the calculator tool.",
            tools=tools,
        )
        if r.has_tool_use:
            names = [tc.name for tc in r.tool_calls]
            print(f"  [tool-req] stop_reason=tool_use, tools={names}")
        else:
            print(f"  [tool-req] no tool use (unexpected): {r.text[:60]}")

        # --- Automatic run() loop ---
        r = await c.run(
            "What is (12 + 8) * 3? Use the calculator tool and tell me the answer.",
            tools=tools,
            tool_handler=calculator_handler,
        )
        print(f"  [tool-run] {r.text.strip()[:80]}")

        # --- Streaming run_stream() loop ---
        sys.stdout.write("  [tool-strm] ")
        sys.stdout.flush()
        async for event in c.run_stream(
            "What is 99 * 42? Use the calculator tool and give the answer.",
            tools=tools,
            tool_handler=calculator_handler,
        ):
            if isinstance(event, ToolUse):
                sys.stdout.write(f"[call:{event.name}]")
                sys.stdout.flush()
            elif isinstance(event, ToolResult):
                sys.stdout.write(f"[result:{event.content}]")
                sys.stdout.flush()
            elif hasattr(event, "text_delta") and event.text_delta:
                sys.stdout.write(event.text_delta)
                sys.stdout.flush()
        print()


async def main():
    mode = os.environ.get("TEST_MODE", "all")

    if mode in ("all", "agent"):
        print("=== ClaudeAgentClient ===")
        configs = AGENT_CONFIGS if mode == "all" else [AGENT_CONFIGS[1]]
        for cfg in configs:
            print(f">>> {cfg['label']}: model={cfg['model']} effort={cfg['effort']}")
            await run_agent(cfg)
            print()

    if mode in ("all", "chat"):
        print("=== ChatClient ===")
        await run_chat()
        print()

    if mode in ("all", "chat", "tools"):
        print("=== ChatClient Tools ===")
        await run_chat_tools()
        print()

    if mode == "single":
        print("=== ClaudeAgentClient (sonnet) ===")
        await run_agent(AGENT_CONFIGS[1])


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