import asyncio
from typing import AsyncGenerator

from google.adk.agents import BaseAgent, LlmAgent, SequentialAgent
from google.adk.agents.invocation_context import InvocationContext
from google.adk.events import Event
from google.adk.runners import Runner
from google.adk.sessions import InMemorySessionService
from google.genai.types import Content, Part


# 天気情報を取得するエージェント
class WeatherAgent(BaseAgent):
    def __init__(self) -> None:
        super().__init__(
            name="weather_agent",
            description="指定された場所の現在の天気情報を取得します。",
        )

    async def _run_async_impl(
        self, ctx: InvocationContext
    ) -> AsyncGenerator[Event, None]:
        # ここではダミーの天気情報を返します。
        yield Event(
            author="weather_agent",
            content=Content(
                role="model",
                parts=[Part(text="12/8の東京の天気は晴れ、気温は25度です。")],
            ),
            partial=False,
            invocation_id=ctx.invocation_id,
        )


# 日記を書くエージェント
class DiaryWriterAgent(LlmAgent):
    def __init__(self) -> None:
        super().__init__(
            name="diary_writer_agent",
            description="提供された天気情報を基に日記を書きます。",
            model="gemini-2.5-flash",
            instruction="あなたは優れた日記作家です。提供された天気情報を基に、感情豊かな日記を書いてください。",
        )


# 天気情報と日記を書くエージェントを組み合わせたシーケンシャルエージェント
class DiaryAgent(SequentialAgent):
    def __init__(self) -> None:
        weather_agent = WeatherAgent()
        diary_writer_agent = DiaryWriterAgent()
        super().__init__(
            name="diary_agent",
            description="天気情報を取得し、その情報を基に日記を書きます。",
            sub_agents=[weather_agent, diary_writer_agent],
        )


def main() -> None:
    session_service = InMemorySessionService()
    diary_agent = DiaryAgent()
    app_name = "write daily"
    user_id = "user_123"

    runner = Runner(
        agent=diary_agent,
        app_name=app_name,
        session_service=session_service,
    )

    content = Content(
        role="user",
        parts=[
            Part(
                text="12/8は渋谷に買い物に行きました。欲しいものがたくさんあって、とても楽しかったです。",
            )
        ],
    )

    session = asyncio.run(
        session_service.create_session(app_name=app_name, user_id=user_id)
    )

    events = runner.run(
        user_id="user_123",
        session_id=session.id,
        new_message=content,
    )

    for event in events:
        if event.is_final_response() and event.content:
            print(f"{event.author}の出力:")
            for part in event.content.parts:
                print(part.text)


if __name__ == "__main__":
    main()

# 天気情報を取得するエージェント
class WeatherAgent(BaseAgent):
    def __init__(self) -> None:
        super().__init__(
            name="weather_agent",
            description="指定された場所の現在の天気情報を取得します。",
        )

    async def _run_async_impl(
        self, ctx: InvocationContext
    ) -> AsyncGenerator[Event, None]:
        # ここではダミーの天気情報を返します。
        yield Event(
            author="weather_agent",
            content=Content(
                role="model",
                parts=[Part(text="12/8の東京の天気は晴れ、気温は25度です。")],
            ),
            partial=False,
            invocation_id=ctx.invocation_id,
        )

# 日記を書くエージェント
class DiaryWriterAgent(LlmAgent):
    def __init__(self) -> None:
        super().__init__(
            name="diary_writer_agent",
            description="提供された天気情報を基に日記を書きます。",
            model="gemini-2.5-flash",
            instruction="あなたは優れた日記作家です。提供された天気情報を基に、感情豊かな日記を書いてください。",
        )

# 天気情報と日記を書くエージェントを組み合わせたシーケンシャルエージェント
class DiaryAgent(SequentialAgent):
    def __init__(self) -> None:
        weather_agent = WeatherAgent()
        diary_writer_agent = DiaryWriterAgent()
        super().__init__(
            name="diary_agent",
            description="天気情報を取得し、その情報を基に日記を書きます。",
            sub_agents=[weather_agent, diary_writer_agent],
        )

def main() -> None:
    session_service = InMemorySessionService()
    diary_agent = DiaryAgent()
    app_name = "write daily"
    user_id = "user_123"

    runner = Runner(
        agent=diary_agent,
        app_name=app_name,
        session_service=session_service,
    )

    content = Content(
        role="user",
        parts=[
            Part(
                text="12/8は渋谷に買い物に行きました。欲しいものがたくさんあって、とても楽しかったです。",
            )
        ],
    )

    session = asyncio.run(
        session_service.create_session(app_name=app_name, user_id=user_id)
    )

    events = runner.run(
        user_id="user_123",
        session_id=session.id,
        new_message=content,
    )

    for event in events:
        if event.is_final_response() and event.content:
            print("各エージェントの出力:")
            for part in event.content.parts:
                print(part.text)

weather_agentの出力:
12/8の東京の天気は晴れ、気温は25度です。
diary_writer_agentの出力:
12月8日

今日は、冬の足音が聞こえ始めるはずの12月だというのに、東京はまさかの快晴、気温はなんと25度！まるで春のような陽気に包まれた一日でした。厚手のコートは必要なく、身軽な格好で渋谷の街を闊歩できるなんて、なんだか不思議な感覚です。こんなにも穏やかな日差しの中で、心ゆくまでショッピングを楽しめるなんて、本当に贅沢な時間でした。

渋谷の街は、どこか浮き足立つようなクリスマスの気配と、この季節外れの暖かさが相まって、いつも以上にキラキラと輝いて見えました。ウィンドウに並ぶ魅力的な商品たちの一つ一つが、私の心を強く惹きつけます。あれも欲しい、これも素敵！と、目がいくつあっても足りないくらい。

特に目的を決めずに歩き始めたはずなのに、気づけば時間が経つのも忘れて、たくさんのお店を巡っていました。お洋服、雑貨、アクセサリー…どれもこれもが、今の私に「おいでよ！」と手招きしているよう。実際に手に取ったり、鏡の前で合わせてみたりするたびに、胸が高鳴るのを感じました。まさに、宝探しをしているような気分。

欲しいものがたくさん見つかって、今日は本当に大満足！両手にいっぱいのショップバッグを抱えて、充実感に満たされた笑顔で家路につきました。冬なのに暖かな日差しが降り注ぐ中、こんなにも楽しい時間を過ごせるなんて、最高の贈り物をもらった気分です。心も体もぽかぽかと温かくなった、幸せな一日でした。