from __future__ import annotations

import asyncio
import os
import sys
from typing import Any

from langgraph_sdk import get_client


BASE_URL = os.getenv("LANGGRAPH_URL", "http://192.168.0.100:2026")
ASSISTANT_ID = "student_score"
QUESTION = "帮我看看白若耶同学的信息"
CONTEXT = {
    "team_id": "19",
    "env": "dev",
}


if hasattr(sys.stdout, "reconfigure"):
    sys.stdout.reconfigure(encoding="utf-8")


def event_data(event: Any) -> Any:
    if isinstance(event, dict):
        return event.get("data")
    return getattr(event, "data", None)


def event_name(event: Any) -> str:
    if isinstance(event, dict):
        return str(event.get("event") or "")
    return str(getattr(event, "event", ""))


def pick_final_answer(data: Any) -> str:
    if not isinstance(data, dict):
        return ""
    render_answer = data.get("render_answer")
    if isinstance(render_answer, dict) and isinstance(render_answer.get("answer"), str):
        return render_answer["answer"]
    answer = data.get("answer")
    return answer if isinstance(answer, str) else ""


def pick_answer_source(data: Any) -> str:
    if not isinstance(data, dict):
        return ""
    render_answer = data.get("render_answer")
    if isinstance(render_answer, dict) and isinstance(render_answer.get("answer_source"), str):
        return render_answer["answer_source"]
    answer_source = data.get("answer_source")
    return answer_source if isinstance(answer_source, str) else ""


async def main() -> None:
    client = get_client(url=BASE_URL)
    thread = await client.threads.create()
    answer_parts: list[str] = []
    final_answers: list[str] = []

    print("=== LangGraph updates + messages-tuple 流式调用 ===", flush=True)
    print("[updates] 节点进度；[messages] 最终 Markdown token。", flush=True)
    print("", flush=True)

    async for event in client.runs.stream(
        thread_id=thread["thread_id"],
        assistant_id=ASSISTANT_ID,
        input={"question": QUESTION},
        context=CONTEXT,
        stream_mode=["updates", "messages-tuple"],
    ):
        name = event_name(event)
        data = event_data(event)

        if name == "updates":
            final_answer = pick_final_answer(data)
            answer_source = pick_answer_source(data)
            if isinstance(data, dict):
                print(f"\n[updates] {', '.join(data.keys())}", flush=True)
                render_answer = data.get("render_answer")
                if isinstance(render_answer, dict):
                    print(f"[render_answer.keys] {', '.join(render_answer.keys())}", flush=True)
            else:
                print(f"\n[updates] {data}", flush=True)
            if answer_source:
                print(f"[answer_source] {answer_source}", flush=True)
            if final_answer:
                final_answers.append(final_answer)
                print("\n[final.answer]", flush=True)
                print(final_answer, flush=True)
            continue

        if name == "messages":
            if isinstance(data, list) and data:
                message_chunk = data[0]
                content = message_chunk.get("content") if isinstance(message_chunk, dict) else ""
                if isinstance(content, str) and content:
                    answer_parts.append(content)
                    print(content, end="", flush=True)
            continue

    print("\n\n=== 最终 Markdown 长度 ===", flush=True)
    print(len(final_answers[-1] if final_answers else "".join(answer_parts)), flush=True)


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