import asyncio
import traceback

from lambda_agent_api.agent_api import AgentApi, AgentBusyException
from lambda_agent_api.server import (
    MsgEventTextChunk,
    MsgEventToolCallChunk,
    MsgEventToolResult,
    MsgEventSendFile,
)


async def main():
    chat_id = input("Chat id: ") or 0
    api = AgentApi("agent-1", "ws://localhost:8000/", chat_id=chat_id)

    try:
        await api.connect()
    except AgentBusyException:
        print(f"Чат {chat_id} занят другим клиентом")
        return

    while True:
        try:
            prompt = await asyncio.get_event_loop().run_in_executor(None, input, ">>> ")
            attachments_input = await asyncio.get_event_loop().run_in_executor(
                None, input, "Attachments (comma-separated, empty for none): "
            )
            attachments = (
                [a.strip() for a in attachments_input.split(",") if a.strip()]
                if attachments_input.strip()
                else None
            )

            print("Agent: ", end="")
            is_tool = False
            async for chunk in api.send_message(prompt, attachments):
                match chunk:
                    case MsgEventTextChunk():
                        is_tool = False
                        print(chunk.text, end="", flush=True)
                    case MsgEventToolCallChunk():
                        if not is_tool:
                            print(
                                f"\n\n### TOOL CALL: ({chunk.tool_name}) ",
                                end="",
                                flush=True,
                            )
                            is_tool = True
                        print(chunk.args_chunk, end="", flush=True)
                    case MsgEventToolResult():
                        is_tool = False
                        print(f"\nResult: {chunk.result}\n\n", end="", flush=True)
                    case MsgEventSendFile():
                        print(f"\n[SEND FILE] {chunk.path}\n", end="", flush=True)

            print("\n")
        except KeyboardInterrupt:
            break
        except:
            traceback.print_exc()

    await api.close()


asyncio.run(main())