import asyncio
from typing import Any
from uuid import uuid4
import httpx
from a2a.client import A2AClient, A2ACardResolver
from a2a.types import (
    MessageSendParams,
    SendStreamingMessageRequest,
)

# チャットループ
async def interact_with_server(client: A2AClient) -> None:
    while True:
        # 入力の取得
        user_input = input('\n> ')

        # メッセージの定義
        send_message_payload: dict[str, Any] = {
            'message': {
                'role': 'user',
                'parts': [{'type': 'text', 'text': user_input}],
                'messageId': uuid4().hex,
            },
        }

        # メッセージの送信
        try:
            streaming_request = SendStreamingMessageRequest(
                id=uuid4().hex,
                params=MessageSendParams(**send_message_payload)
            )
            stream_response = client.send_message_streaming(streaming_request)
            async for chunk in stream_response:
                print(get_response_text(chunk), end='', flush=True)
                await asyncio.sleep(0.1)
        except Exception as e:
            print(f'エラーが発生しました: {e}')

# レスポンステキストの取得
def get_response_text(chunk):
    data = chunk.model_dump(mode='json', exclude_none=True)
    result = data.get('result', {})

    # アーティファクト更新イベント
    if 'artifact' in result:
        parts = result['artifact']['parts']
    # ステータス更新イベント
    elif 'status' in result:
        print('\n', result['status'])
        return ''
    # その他のイベント
    else:
        return ''

    # 最初のパートのテキストを返す
    return parts[0].get('text', '')

# メイン
async def main() -> None:
    async with httpx.AsyncClient() as httpx_client:
        # A2ACardResolverの準備
        resolver = A2ACardResolver(httpx_client, "http://localhost:10001")

        # エージェントカードの取得
        agent_card = await resolver.get_agent_card()

        # A2AClientの準備
        client = A2AClient(httpx_client=httpx_client, agent_card=agent_card)
        await interact_with_server(client)

# メインの実行
if __name__ == '__main__':
    asyncio.run(main())