import asyncio
import os

from agent_framework import ChatAgent
from agent_framework.azure import AzureOpenAIChatClient

from typing import Callable, Awaitable
from agent_framework import AgentRunContext, FunctionInvocationContext, ChatContext


# ---------- エージェントミドルウェア定義 ----------
async def logging_agent_middleware(
    context: AgentRunContext,
    next: Callable[[AgentRunContext], Awaitable[None]],
) -> None:
    """Agent実行の前後でログ出力"""
    print("> Before Agent Execution")
    await next(context)
    print("> After Agent Execution")

# ---------- チャットミドルウェア定義 ----------
async def logging_chat_middleware(
    context: ChatContext,
    next: Callable[[ChatContext], Awaitable[None]],
) -> None:
    """LLM呼び出し前後でログ出力"""
    print(f"  > Before LLM Call. Messages: {len(context.messages)}")
    for m in context.messages:
        print(f"    - {m.role}: {m.text}")  
    await next(context)
    print("  > After LLM Call")

# ---------- 関数ミドルウェア定義 ----------
async def logging_function_middleware(
    context: FunctionInvocationContext,
    next: Callable[[FunctionInvocationContext], Awaitable[None]],
) -> None:
    """ツール呼び出し前後でログ出力"""
    print(f"    > Before Function Call: {context.function.name}")
    await next(context)
    print("    > After Function Call")

# ---------- 関数定義 ----------
def get_weather(location: str) -> str:
    """Get weather for a location."""
    print("        - Executing get_weather function.")
    return f"Weather in {location}: 72°F and sunny"

# ---------- エージェントの作成 ----------
agent = ChatAgent(
    name="WeatherAgent",
    chat_client=AzureOpenAIChatClient(
        api_key=API_KEY,
        endpoint=ENDPOINT,
        deployment_name="gpt-4.1",
    ),
    instructions = "あなたは優れたアシスタントです。",
    tools=[get_weather],
    middleware=[
        logging_agent_middleware,
        logging_function_middleware,
        logging_chat_middleware,
    ],
)

# ---------- エージェントの実行 ----------
async def main():

    response = await agent.run("パリの気温は？")
    print(f"Assistant: {response}")

asyncio.run(main())


# ---------- 出力例 ----------
# > Before Agent Execution
#   > Before LLM Call. Messages: 2
#     - system: あなたは優れたアシスタントです。
#     - user: パリの気温は？
#   > After LLM Call
#     > Before Function Call: get_weather
#         - Executing get_weather function.
#     > After Function Call
#   > Before LLM Call. Messages: 4
#     - system: あなたは優れたアシスタントです。
#     - user: パリの気温は？
#     - assistant:
#     - tool:
#   > After LLM Call
# > After Agent Execution
# Assistant: パリの気温は摂氏約22°C（華氏72°F）で、晴れています。

from agent_framework import ContextProvider, Context

# ---------- コンテキストプロバイダー定義 ----------
class UserPrefMemory(ContextProvider):
    def __init__(self):
        self.pref_city = None  # ユーザー嗜好のストレージ（例：好きな都市）

    # 実行"前" : 記憶をLLMに注入
    async def invoking(self, messages, **kwargs) -> Context:
        if self.pref_city:
            return Context(
                instructions=f"ユーザーは {self.pref_city} が好き。回答で優先的に触れること。"
                # instructions だけでなく messages / tools も追加可能 
            )
        return Context()

    # 実行"後" : 新しい嗜好の検出 → 記憶
    async def invoked(self, request_messages, response_messages=None, **kwargs) -> None:
        text = request_messages[-1].text if request_messages else ""
        if "I like " in text:
            self.pref_city = text.split("I like ", 1)[1].split()[0]

# ---------- コンテキストプロバイダー初期化 ----------
memory = UserPrefMemory()

# ---------- エージェント・スレッドの構築 ----------
agent = ChatAgent(
    name="MemoryAgent",
    chat_client=AzureOpenAIChatClient(...),
    instructions="文脈を活用して丁寧に答える。",
    context_providers=[memory],
)

thread = agent.get_new_thread()

# ---------- 実行 ----------
async def main():

    # 1回目：学習フェーズ
    await agent.run("I like Paris", thread=thread)

    # 2回目：活用フェーズ（Paris が回答に反映される）
    await agent.run("おすすめスポットは？", thread=thread)

    # 長期記憶を確認
    user_pref_memory = thread.context_provider.providers[0]
    if user_pref_memory:
        print("\n=== 記憶された情報 ===")
        print(f"ユーザーの好みの都市: {user_pref_memory.pref_city}")

asyncio.run(main())

# ---------- 出力例 ----------
# === 記憶された情報 ===
# ユーザーの好みの都市: Paris

import asyncio

from agent_framework import (
    WorkflowBuilder,
    WorkflowContext,
    executor
)
from typing_extensions import Never

# ---------- Executor 定義 ----------
@executor(id="upper_case_executor")
async def upper_case(text: str, ctx: WorkflowContext[str]) -> None:
    """入力テキストを大文字に変換"""
    await ctx.send_message(text.upper())

@executor(id="reverse_text_executor")
async def reverse_text(text: str, ctx: WorkflowContext[Never, str]) -> None:
    """入力テキストを反転"""
    await ctx.yield_output(text[::-1])

# ---------- ワークフロー構築 ----------
# ワークフローの構築
# 1. WorkflowBuilder を初期化
# 2. 開始 Executor を指定
# 3. Edge で Executor 同士を接続（有向グラフを設計）
# 4. Build して実行モデルを確定
builder = WorkflowBuilder()
builder.set_start_executor(upper_case)
builder.add_edge(upper_case, reverse_text)
workflow = builder.build()

# ---------- ワークフロー実行 ----------
async def main() -> None:

    events = await workflow.run("hello world")
    print(events.get_outputs())   
    
if __name__ == "__main__":
    asyncio.run(main())

# ---------- 出力例 ----------
# ['DLROW OLLEH']

import asyncio
from typing import Never

from agent_framework import (
    WorkflowBuilder, 
    WorkflowContext, 
    executor, 
)

# ---------- Executor 定義 ----------
@executor(id="start")
async def start(text: str, ctx: WorkflowContext[str]) -> None:
    """開始ノード: 入力文字列をそのまま下流に流す"""
    await ctx.send_message(text)

@executor(id="handle_normal")
async def handle_normal(text: str, ctx: WorkflowContext[Never, str]) -> None:
    """通常処理: 最終出力（yield_output）を返す"""
    await ctx.yield_output(f"OK: {text}")

@executor(id="handle_spam")
async def handle_spam(text: str, ctx: WorkflowContext[Never, str]) -> None:
    """スパム処理: 最終出力（yield_output）を返す"""
    await ctx.yield_output("SPAM: blocked")

# ---------- ルーティング条件 ----------
def is_spam(msg: str) -> bool:
    """メッセージに 'spam' を含むかどうか（単純な条件判定）"""
    return isinstance(msg, str) and ("spam" in msg.lower())

def is_not_spam(msg: str) -> bool:
    """スパムではないか"""
    return isinstance(msg, str) and ("spam" not in msg.lower())

# ---------- ワークフロー構築 ----------
builder = WorkflowBuilder()
builder.set_start_executor(start)
builder.add_edge(start, handle_spam,   condition=is_spam)      # - "spam" を含む → handle_spam
builder.add_edge(start, handle_normal, condition=is_not_spam)  # - "spam" を含まない → handle_normal
workflow = builder.build()

# ---------- 実行例 ----------
async def main():
    # 例1: 通常メッセージ
    events = await workflow.run("Hello team")
    print("Output (normal):", events.get_outputs()[0])

    # 例2: スパムメッセージ
    events = await workflow.run("Buy now!! SPAM offer!!")
    print("Output (spam):  ", events.get_outputs()[0])

    # 可視化
    from agent_framework import WorkflowViz
    viz = WorkflowViz(workflow)
    print("Mermaid:")
    print(viz.to_mermaid())

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

# ---------- 出力例 ----------
# Output (normal): OK: Hello team
# Output (spam):   SPAM: blocked

import asyncio
from typing import Any, Never

from agent_framework import (
    WorkflowBuilder,      
    WorkflowContext,      
    executor,             
    Case,
    Default
)



# ---------- Executor 定義 ----------
@executor(id="start")
async def start(msg: Any, ctx: WorkflowContext[Any]) -> None:
    """開始ノード: 入力をそのまま下流へ流す（Switch の起点）"""
    await ctx.send_message(msg)

@executor(id="handle_text")
async def handle_text(text: str, ctx: WorkflowContext[Never, str]) -> None:
    """文字列: 大文字にして最終出力"""
    await ctx.yield_output(f"TEXT: {text.upper()}")

@executor(id="handle_number")
async def handle_number(n: int, ctx: WorkflowContext[Never, str]) -> None:
    """数値: 2倍して最終出力"""
    await ctx.yield_output(f"NUMBER: {n * 2}")

@executor(id="handle_other")
async def handle_other(obj: Any, ctx: WorkflowContext[Never, str]) -> None:
    """その他: 型名を返して最終出力"""
    await ctx.yield_output(f"OTHER: {type(obj).__name__}")


# ---------- ワークフロー構築 ----------
builder = WorkflowBuilder()
builder.set_start_executor(start)

# Switch-Case:
# - ソースからの入力が、"文字列" なら handle_text
# - ソースからの入力が、"数値" なら handle_number
# - ソースからの入力が、そのどれにも該当しなければ Default -> handle_other
builder.add_switch_case_edge_group(
    source=start,
    cases=[
        Case(condition=lambda m: isinstance(m, str), target=handle_text),
        Case(condition=lambda m: isinstance(m, int), target=handle_number),
        Default(target=handle_other),
    ],
)

workflow = builder.build()


# ---------- 実行例 ----------
async def main():
    samples = ["hello world", 21, {"x": 1}]
    for s in samples:
        events = await workflow.run(s)
        print("Input :", s)
        print("Output:", events.get_outputs()[0])
        print("-" * 7)

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

# ---------- 出力例 ----------
# Input : hello world
# Output: TEXT: HELLO WORLD
# -------
# Input : 21
# Output: NUMBER: 42
# -------
# Input : {'x': 1}
# Output: OTHER: dict
# -------

import asyncio
from typing import Never

from agent_framework import (
    WorkflowBuilder,      
    WorkflowContext,      
    executor,             
)

# ---------- Executors ----------

@executor(id="start")
async def start(text: str, ctx: WorkflowContext[str]) -> None:
    """開始ノード: 入力文字列をそのまま下流へ流す（Fan-Out の起点）"""
    await ctx.send_message(text)

@executor(id="count_words")
async def count_words(text: str, ctx: WorkflowContext[int]) -> None:
    """単語数を数える（スペース区切り・超シンプル）"""
    await ctx.send_message(len(text.split()))

@executor(id="count_chars")
async def count_chars(text: str, ctx: WorkflowContext[int]) -> None:
    """文字数を数える（改行含む・超シンプル）"""
    await ctx.send_message(len(text))

@executor(id="aggregate")
async def aggregate(results: list[int], ctx: WorkflowContext[Never, str]) -> None:
    """集約ノード（Fan-In の終点）"""
    words, chars = results  
    await ctx.yield_output(f"Words: {words}, Chars: {chars}")

# ---------- ワークフロー構築 ----------
builder = WorkflowBuilder()
builder.set_start_executor(start)

# Fan-Out: start から 2 系統に並列分岐
builder.add_fan_out_edges(start, [count_words, count_chars])

# Fan-In: 2 系統の結果が揃ったら aggregate に list で渡される
builder.add_fan_in_edges([count_words, count_chars], aggregate)

workflow = builder.build()

# ---------- 実行例 ----------
async def main():
    text = "Fan out and fan in are powerful workflow patterns."
    events = await workflow.run(text)
    outputs = events.get_outputs() 
    if outputs:
        print("Output:", outputs[0])

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


# ---------- 出力例 ----------
# Output: Words: 9, Chars: 50

import asyncio
from typing import Never
from agent_framework import WorkflowBuilder, WorkflowContext, executor

PREFIX_KEY = "greeting_prefix"

# ---------- Executors ----------
@executor(id="save_prefix")
async def save_prefix(name: str, ctx: WorkflowContext[str]) -> None:
    """共有状態に prefix を保存し、名前を下流へ送る"""
    await ctx.set_shared_state(key=PREFIX_KEY, value="Hello")  # 共有状態にデータを保存
    await ctx.send_message(name)

@executor(id="greet")
async def greet(name: str, ctx: WorkflowContext[Never, str]) -> None:
    """共有状態から prefix を取得し、名前を結合して出力する"""
    prefix = await ctx.get_shared_state(key=PREFIX_KEY)       # 共有状態からデータを取得
    await ctx.yield_output(f"{prefix}, {name}.")

# ---------- ワークフロー構築 ----------
builder = WorkflowBuilder()
builder.set_start_executor(save_prefix)
builder.add_edge(save_prefix, greet) 
workflow = builder.build()

# ---------- 実行例 ----------
async def main():
    events = await workflow.run("Taro")
    print(events.get_outputs()[0]) 

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

# ---------- 出力例 ----------
# Hello, Taro.

from agent_framework import (
    InMemoryCheckpointStorage,
    WorkflowBuilder,
)

...

# ---------- チェックポイントを保持するストレージ作成 ----------
checkpoint_storage = InMemoryCheckpointStorage()

# ---------- ワークフロー構築 ----------
builder = WorkflowBuilder()
builder.set_start_executor(start_executor)  
builder.add_edge(start_executor, executor_b)  
builder.add_edge(executor_b, executor_c)
builder.add_edge(executor_b, end_executor)
builder.with_checkpointing(checkpoint_storage) # チェックポイントを指定
workflow = builder.build()

# ---------- 初回実行 ----------
async for event in workflow.run_streaming(input):
...

# ---------- 保存されたチェックポイント一覧を取得 ----------
checkpoints = await checkpoint_storage.list_checkpoints()

# ---------- 最後に保存されたチェックポイントから途中再開 ----------
last_checkpoint = checkpoints[-1]
async for event in workflow.run_stream_from_checkpoint(last_checkpoint.checkpoint_id):
...

import asyncio
from dataclasses import dataclass
from typing_extensions import Never

from agent_framework import (
    WorkflowBuilder,
    WorkflowContext,
    Executor,
    RequestInfoExecutor,
    RequestInfoEvent,
    RequestInfoMessage,
    RequestResponse,
    WorkflowOutputEvent,
    WorkflowStatusEvent,
    WorkflowRunState,
    handler,
)

# ---------- 型定義 ----------
@dataclass
class ApprovalRequest(RequestInfoMessage):
    """承認リクエスト（RequestInfoMessage を継承）"""
    prompt: str = ""
    attempt_count: int = 1

# ---------- 承認管理 Executor ----------
class ApprovalManager(Executor):
    """承認プロセスを管理するExecutor"""

    def __init__(self, id: str | None = None):
        super().__init__(id=id or "approval_manager")

    @handler
    async def start(self, _: str, ctx: WorkflowContext[ApprovalRequest]) -> None:
        """最初の承認リクエストを送信"""
        print("[ApprovalManager] 承認プロセスを開始します")
        await ctx.send_message(ApprovalRequest(
            prompt="レポート公開の承認をお願いします (y/n)",
            attempt_count=1
        ))

    @handler
    async def handle_response(
        self,
        feedback: RequestResponse[ApprovalRequest, str],
        ctx: WorkflowContext[ApprovalRequest, str],
    ) -> None:
        """承認レスポンスを処理"""
        reply = (feedback.data or "").strip().lower()
        request_data = feedback.original_request
        
        print(f"[ApprovalManager] 試行 {request_data.attempt_count}: 応答='{reply}'")
        
        if reply in ("y", "yes", "1", "ok", "承認"):
            # 承認された場合
            print(f"[ApprovalManager] 承認されました！(試行回数: {request_data.attempt_count})")
            await ctx.yield_output(f"レポートが正常に公開されました (試行回数: {request_data.attempt_count})")
            return
        
        # 却下された場合、再試行
        next_attempt = request_data.attempt_count + 1
        print(f"[ApprovalManager] 却下されました。{next_attempt} 回目の承認をリクエストします...")
        
        await ctx.send_message(ApprovalRequest(
            prompt=f"レポート公開の承認をお願いします (y/n) - {next_attempt} 回目の試行",
            attempt_count=next_attempt
        ))

# ---------- ワークフロー構築 ----------
def build_workflow():
    """承認ワークフローを構築"""
    approval_manager = ApprovalManager(id="approval_manager")
    request_info = RequestInfoExecutor(id="request_info")
    
    workflow = (
        WorkflowBuilder()
        .set_start_executor(approval_manager)
        .add_edge(approval_manager, request_info)  # 承認リクエスト送信
        .add_edge(request_info, approval_manager)  # 承認レスポンス受信
        .build()
    )
    
    return workflow

# ---------- メイン実行 ----------
async def main():
    """メイン実行関数（公式サンプルのパターンに基づく）"""
    workflow = build_workflow()
    
    pending_responses: dict[str, str] | None = None
    completed = False
    workflow_output: str | None = None
    
    while not completed:
        # 初回は run_stream、2回目以降は send_responses_streaming
        stream = (
            workflow.send_responses_streaming(pending_responses) 
            if pending_responses 
            else workflow.run_stream("start")
        )
        
        # イベントを収集
        events = [event async for event in stream]
        pending_responses = None
        
        # リクエストと出力を収集
        requests: list[tuple[str, str, int]] = []  # (request_id, prompt, attempt_count)
        
        for event in events:
            if isinstance(event, RequestInfoEvent) and isinstance(event.data, ApprovalRequest):
                # 承認リクエストイベント
                requests.append((
                    event.request_id, 
                    event.data.prompt, 
                    event.data.attempt_count
                ))
            elif isinstance(event, WorkflowOutputEvent):
                # ワークフロー出力（完了）
                workflow_output = str(event.data)
                completed = True
        
        # 状態変化を表示（デバッグ用）
        pending_status = any(
            isinstance(e, WorkflowStatusEvent) and e.state == WorkflowRunState.IN_PROGRESS_PENDING_REQUESTS
            for e in events
        )
        idle_with_requests = any(
            isinstance(e, WorkflowStatusEvent) and e.state == WorkflowRunState.IDLE_WITH_PENDING_REQUESTS
            for e in events
        )
        
        if pending_status:
            print("[状態] IN_PROGRESS_PENDING_REQUESTS (リクエスト処理中)")
        if idle_with_requests:
            print("[状態] IDLE_WITH_PENDING_REQUESTS (人間の入力待ち)")
        
        # 承認リクエストがある場合、ユーザーに入力を求める
        if requests and not completed:
            responses: dict[str, str] = {}
            
            for req_id, prompt, attempt_count in requests:
                print(f"\n{'='*60}")
                print(f"承認リクエスト (試行 {attempt_count})")
                print(f"{'='*60}")
                print(f"{prompt}")
                
                answer = input("承認しますか？ (y/n/exit): ").strip().lower()
                
                if answer == "exit":
                    print("ユーザーによって処理が中断されました")
                    return
                
                responses[req_id] = answer
                print(f"入力結果: {'承認' if answer in ('y', 'yes', '1', 'ok', '承認') else '却下'}")
            
            pending_responses = responses
    
    # 最終結果を表示
    print(f"\n{'='*70}")
    print(f"最終結果: {workflow_output}")
    print(f"{'='*70}")
    
    # ワークフロー図を表示
    print(f"\n{'='*50}")
    print("ワークフロー図 (Mermaid)")
    print(f"{'='*50}")
    from agent_framework import WorkflowViz
    viz = WorkflowViz(workflow)
    print(viz.to_mermaid())

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


# ---------- 出力 ----------
# [ApprovalManager] 承認プロセスを開始します
# [状態] IN_PROGRESS_PENDING_REQUESTS (リクエスト処理中)
# [状態] IDLE_WITH_PENDING_REQUESTS (人間の入力待ち)

# ============================================================
# 承認リクエスト (試行 1)
# ============================================================
# レポート公開の承認をお願いします (y/n)
# 承認しますか？ (y/n/exit): n
# 入力結果: 却下
# [ApprovalManager] 試行 1: 応答='n'
# [ApprovalManager] 却下されました。2 回目の承認をリクエストします...
# [状態] IN_PROGRESS_PENDING_REQUESTS (リクエスト処理中)
# [状態] IDLE_WITH_PENDING_REQUESTS (人間の入力待ち)

# ============================================================
# 承認リクエスト (試行 2)
# ============================================================
# レポート公開の承認をお願いします (y/n) - 2 回目の試行
# 承認しますか？ (y/n/exit): y
# 入力結果: 承認
# [ApprovalManager] 試行 2: 応答='y'
# [ApprovalManager] 承認されました！(試行回数: 2)

# ======================================================================
# 最終結果: レポートが正常に公開されました (試行回数: 2)
# ======================================================================

from agent_framework import WorkflowBuilder
from agent_framework import WorkflowViz

# (省略)

# ---------- ワークフロー構築 ----------
builder = WorkflowBuilder()
builder.set_start_executor(start)
builder.add_fan_out_edges(start, [count_words, count_chars])
builder.add_fan_in_edges([count_words, count_chars], aggregate)
workflow = builder.build()

# ---------- 可視化 ----------
viz = WorkflowViz(workflow)
print(viz.to_mermaid())