#!/usr/bin/env python3
"""Measure DeerFlow model-factory latency without exposing prompts or secrets."""

from __future__ import annotations

import argparse
import json
import sys
import time
from pathlib import Path
from typing import Any

from langchain_core.messages import HumanMessage

from deerflow.config import get_app_config
from deerflow.models import create_chat_model


def _args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(description="Diagnose DeerFlow LLM first-chunk and completion latency")
    parser.add_argument("--model", required=True, help="DeerFlow config model name")
    parser.add_argument("--prompt", default="请先思考，然后只回答：17乘以23等于多少？")
    parser.add_argument("--prompt-file", type=Path, help="UTF-8 prompt file; content is never printed")
    parser.add_argument("--repeat", type=int, default=1, help="Repeat the selected prompt N times")
    parser.add_argument("--max-tokens", type=int, default=256)
    parser.add_argument("--thinking", action=argparse.BooleanOptionalAction, default=True)
    return parser.parse_args()


def _reasoning_text(chunk: Any) -> str:
    kwargs = getattr(chunk, "additional_kwargs", None) or {}
    value = kwargs.get("reasoning_content") or kwargs.get("reasoning")
    return value if isinstance(value, str) else ""


def _content_text(chunk: Any) -> str:
    content = getattr(chunk, "content", "")
    if isinstance(content, str):
        return content
    if isinstance(content, list):
        return "".join(part.get("text", "") for part in content if isinstance(part, dict))
    return str(content or "")


def main() -> int:
    args = _args()
    prompt = args.prompt_file.read_text(encoding="utf-8") if args.prompt_file else args.prompt
    prompt *= max(args.repeat, 1)

    config = get_app_config()
    model_config = config.get_model_config(args.model)
    if model_config is None:
        print(f"ERROR model_not_found={args.model}", file=sys.stderr)
        return 2

    print(
        "CONFIG "
        + json.dumps(
            {
                "name": model_config.name,
                "use": model_config.use,
                "base_url": str(model_config.base_url or ""),
                "request_timeout": model_config.request_timeout,
                "max_retries": model_config.max_retries,
                "supports_thinking": model_config.supports_thinking,
                "when_thinking_enabled": model_config.when_thinking_enabled,
                "prompt_chars": len(prompt),
            },
            ensure_ascii=False,
        ),
        flush=True,
    )

    build_started = time.monotonic()
    model = create_chat_model(args.model, thinking_enabled=args.thinking)
    build_ms = int((time.monotonic() - build_started) * 1000)
    print(
        "MODEL "
        + json.dumps(
            {
                "runtime_class": f"{type(model).__module__}.{type(model).__name__}",
                "build_ms": build_ms,
            }
        ),
        flush=True,
    )

    messages = [HumanMessage(content=prompt)]
    if hasattr(model, "_get_request_payload"):
        payload = model._get_request_payload(messages, max_tokens=args.max_tokens)
        safe_payload = dict(payload)
        safe_payload["messages"] = [
            {
                "role": item.get("role"),
                "content_chars": len(str(item.get("content") or "")),
                "has_tools": bool(item.get("tool_calls")),
            }
            for item in payload.get("messages", [])
        ]
        print("PAYLOAD " + json.dumps(safe_payload, ensure_ascii=False, default=str), flush=True)

    started = time.monotonic()
    first_chunk_ms: int | None = None
    first_reasoning_ms: int | None = None
    first_content_ms: int | None = None
    chunks = 0
    reasoning_chars = 0
    content_chars = 0
    usage: dict[str, Any] = {}

    try:
        for chunk in model.stream(messages, max_tokens=args.max_tokens):
            chunks += 1
            now_ms = int((time.monotonic() - started) * 1000)
            if first_chunk_ms is None:
                first_chunk_ms = now_ms
                print(f"FIRST_CHUNK elapsed_ms={now_ms}", flush=True)

            reasoning = _reasoning_text(chunk)
            content = _content_text(chunk)
            if reasoning:
                reasoning_chars += len(reasoning)
                if first_reasoning_ms is None:
                    first_reasoning_ms = now_ms
                    print(f"FIRST_REASONING elapsed_ms={now_ms}", flush=True)
            if content:
                content_chars += len(content)
                if first_content_ms is None:
                    first_content_ms = now_ms
                    print(f"FIRST_CONTENT elapsed_ms={now_ms}", flush=True)

            chunk_usage = getattr(chunk, "usage_metadata", None)
            if chunk_usage:
                usage = dict(chunk_usage)
    except Exception as exc:
        total_ms = int((time.monotonic() - started) * 1000)
        print(
            "FAILED "
            + json.dumps(
                {"total_ms": total_ms, "error_type": type(exc).__name__, "error": str(exc)},
                ensure_ascii=False,
            ),
            file=sys.stderr,
            flush=True,
        )
        return 1

    total_ms = int((time.monotonic() - started) * 1000)
    print(
        "RESULT "
        + json.dumps(
            {
                "status": "success",
                "first_chunk_ms": first_chunk_ms,
                "first_reasoning_ms": first_reasoning_ms,
                "first_content_ms": first_content_ms,
                "total_ms": total_ms,
                "chunks": chunks,
                "reasoning_chars": reasoning_chars,
                "content_chars": content_chars,
                "usage": usage,
            },
            ensure_ascii=False,
        ),
        flush=True,
    )
    return 0


if __name__ == "__main__":
    raise SystemExit(main())