"""一次大屏图调用的 Langfuse 追踪

不属于处理管道的任一步。把图调用收成可过滤的 trace，并把 Jev 记成当前节点下的 generation
每轮调用单独记录，首次请求与澄清恢复通过 session 关联，不修改业务状态
"""

from __future__ import annotations

import os
from asyncio import CancelledError
from collections.abc import Iterator
from contextlib import contextmanager
from typing import Any

from loguru import logger

from app.config import settings

# 面板 Trace Name 列用的稳定名字，不要拼 thread_id 或用户原文
TRACE_NAME = "screen-instruct"
# 按功能过滤用的 tag，不是每一次请求的动态标签
FEATURE_TAG = "screen-agent"
# 客户端复用连接；callback 按调用创建，不保留跨请求的恢复关联
_client: Any = None
_unavailable = False


def tracing_enabled() -> bool:
    """三项连接配置都有，且没有显式关掉，才发送 trace

    LANGFUSE_TRACING_ENABLED=false 用来本地临时关闭，缺任何一项连接配置也视为关闭
    """
    flag = os.getenv("LANGFUSE_TRACING_ENABLED", "").strip().lower()
    if flag in {"0", "false", "no"}:
        return False
    return bool(
        os.getenv("LANGFUSE_PUBLIC_KEY", "").strip()
        and os.getenv("LANGFUSE_SECRET_KEY", "").strip()
        and os.getenv("LANGFUSE_BASE_URL", "").strip()
    )


def flush_tracing() -> None:
    """把还在内存里的 span 立刻发出去

    Langfuse 默认攒一批再发。演示脚本马上退出时不调用，这条 trace 会丢
    追踪没启用或客户端没建成时直接返回，不报错
    """
    client = _client
    if client is None or _unavailable:
        return
    try:
        client.flush()
    except Exception as exc:  # noqa: BLE001 追踪故障不阻断业务
        logger.opt(exception=exc).warning("Langfuse flush 失败")


async def ainvoke_screen(
    graph: Any,
    payload: Any,
    *,
    thread_id: str,
    session_id: str | None = None,
    user_id: int | None = None,
) -> Any:
    """每轮调用记录一个根 span，返回、等待用户或异常时结束。

    session_id 由调用方在首次与恢复时一致传入（通常为 conversation_id）。
    未传时始终使用 thread_id，不从仅首次请求才有的 payload 推断。
    user_id 为用户中心的用户 ID，接口调用时传入，本地脚本不传。
    """
    with _screen_trace(payload, thread_id, session_id, user_id) as (config, root):
        result = await graph.ainvoke(payload, config=config)
        _finish_trace(root, result)
        return result


async def astream_screen(
    graph, payload, *, thread_id, session_id=None, user_id=None
):
    """同一轮流保持同一个追踪上下文，关闭流时结束追踪"""
    from contextlib import aclosing

    result = {}
    with _screen_trace(payload, thread_id, session_id, user_id) as (config, root):
        async with aclosing(
            graph.astream(payload, config=config, stream_mode=["custom", "values"])
        ) as events:
            async for mode, data in events:
                if mode == "values":
                    result = data
                yield mode, data
        _finish_trace(root, result)


def _finish_trace(root, result):
    if root is not None:
        output = _run_output(result)
        root.update(
            output=output, level="ERROR" if output["outcome"] == "failed" else "DEFAULT"
        )


@contextmanager
def _screen_trace(payload, thread_id, session_id, user_id=None):
    config = {"configurable": {"thread_id": thread_id}}
    client = _ensure_client() if tracing_enabled() else None
    if client is None:
        yield config, None
        return

    from langfuse import propagate_attributes
    from langfuse.langchain import CallbackHandler
    from langgraph.types import Command

    # 明确回调所属项目，避免进程中存在多个客户端时跳过追踪
    handler = CallbackHandler(public_key=os.getenv("LANGFUSE_PUBLIC_KEY"))
    # 在当前异步上下文执行回调，手动记录的 Jev generation 才会挂到所属节点。
    handler.run_inline = True
    resume = isinstance(payload, Command) and payload.resume is not None
    trace_input = (
        {"reply": payload.resume}
        if resume
        else {"prompt": payload.get("request", {}).get("prompt")}
    )
    with (
        client.start_as_current_observation(
            name=TRACE_NAME,
            as_type="span",
            input=trace_input,
        ) as root,
        propagate_attributes(
            trace_name=TRACE_NAME,
            session_id=session_id or thread_id,
            # 用户中心的 user_id，面板 Users 页按用户聚合；脚本调用没有用户时不设
            user_id=str(user_id) if user_id is not None else None,
            tags=[FEATURE_TAG],
            metadata={
                "thread_id": thread_id,
                "phase": "resume" if resume else "request",
            },
        ),
    ):
        try:
            yield (
                {**config, "callbacks": [handler], "run_name": "run-screen-graph"},
                root,
            )
        except (CancelledError, GeneratorExit):
            root.update(output={"outcome": "cancelled"})
            raise
        except Exception as error:
            root.update(
                output={"outcome": "failed", "error_type": type(error).__name__},
                level="ERROR",
                status_message=str(error),
            )
            raise


def _run_output(state: dict[str, Any]) -> dict[str, Any]:
    """记录本轮结束原因；completed 表示图已返回，不代表前端应用了动作。"""
    pending = state.get("__interrupt__")
    outcome = "completed"
    if pending:
        outcome = "waiting_for_user"
    elif state.get("status") in {"failed", "cancelled"}:
        outcome = state["status"]
    return {
        "outcome": outcome,
        "intent": state.get("intent"),
        "answer": state.get("answer"),
        "offer": state.get("offer"),
        "clarification": [item.value for item in pending] if pending else None,
        "actions": state.get("actions"),
        "validation_passed": state.get("validation_passed"),
        "failure_reason": state.get("failure_reason"),
    }


@contextmanager
def observe_decision(request: Any) -> Iterator[Any]:
    """在当前节点下面记一次 Jev 调用，yield 的是 generation，未追踪时是 None

    TypeSafe 不走 LangChain callback，所以要手记。类型必须是 generation，面板才能算模型和 token
    只有图回调已经挂上当前 span 时才记。没有父 span 就跳过，避免 Jev 单独变成一条 trace
    这里不吞异常：失败时标 ERROR 后原样抛出，调用方的 DecisionError 类型不变
    """
    client = _ensure_client() if tracing_enabled() else None
    if client is None or not _has_active_span():
        yield None
        return

    model = request.model or settings.typesafe_model
    # 意图判断可能附带对上一轮提议的回应判断，仍属于同一次意图识别
    name = (
        "classify-intent"
        if set(request.decisions) <= {"intent", "offer_response"}
        and "intent" in request.decisions
        else "run-decision"
    )
    # 复用本应用已经选定的客户端
    with client.start_as_current_observation(
        as_type="generation",
        name=name,
        model=model,
        input={
            "questions": {
                name: question.model_dump(mode="json")
                for name, question in request.decisions.items()
            },
            "state": request.state,
        },
    ) as generation:
        try:
            yield generation
        except Exception as error:
            generation.update(
                level="ERROR",
                status_message=str(error),
                metadata=_error_metadata(error),
            )
            raise


def record_decision(generation: Any, response: Any) -> None:
    """把成功的 Jev 结果写回刚才那条 generation

    generation 为 None 表示这次没追踪，直接返回，不影响判断结果
    只写选项、置信度和 token。不写原始 HTTP 响应
    """
    if generation is None:
        return
    usage: dict[str, int] = {}
    if response.usage.input_tokens is not None:
        usage["input_tokens"] = response.usage.input_tokens
    if response.usage.output_tokens is not None:
        usage["output_tokens"] = response.usage.output_tokens

    generation.update(
        model=response.model,
        output={
            name: _result_summary(result) for name, result in response.results.items()
        },
        usage_details=usage or None,
        metadata={
            "request_id": response.request_id or "",
        },
    )


def _ensure_client() -> Any:
    """第一次使用时建客户端

    必须在 load_dotenv 之后调用。初始化失败就记一次警告，本进程后续不再重试
    不要在业务失败时调用，追踪不可用不能挡住图执行
    """
    global _client, _unavailable

    if _unavailable:
        return None
    if _client is not None:
        return _client
    os.environ.setdefault("LANGFUSE_TRACING_ENVIRONMENT", "development")
    from langfuse import Langfuse

    try:
        client = Langfuse()
    except Exception as exc:  # noqa: BLE001 追踪故障不阻断业务
        _unavailable = True
        logger.opt(exception=exc).warning("Langfuse 不可用，本进程跳过追踪")
        return None
    _client = client
    return client


def _has_active_span() -> bool:
    """当前调用栈上有没有 Langfuse 已经打开的 span

    图回调会在节点执行期间挂上 span。没有的话，Jev 不该自己开一条 trace
    """
    from opentelemetry import trace

    return trace.get_current_span().get_span_context().is_valid


def _result_summary(result: Any) -> dict[str, Any]:
    """按决策题类型抽出面板上要看的字段，不回写原始协议对象"""
    summary: dict[str, Any] = {"type": result.type}

    if result.type == "choice":
        summary["choice"] = result.choice
        summary["confidence"] = result.confidence
        summary["probabilities"] = result.probabilities
    elif result.type == "score":
        summary["score"] = result.score
        summary["confidence"] = result.confidence
    elif result.type == "noul":
        summary["noul"] = result.noul
    return summary


def _error_metadata(error: BaseException) -> dict[str, str]:
    """失败时附上错误种类和 request_id，方便和 Jev 服务端对日志

    只取已经整理过的安全字段，不把请求体或密钥写进 metadata
    """
    metadata = {"error_type": type(error).__name__}
    kind = getattr(error, "kind", None)
    request_id = getattr(error, "request_id", None)
    status = getattr(error, "status", None)

    if isinstance(kind, str):
        metadata["kind"] = kind
    if isinstance(request_id, str):
        metadata["request_id"] = request_id
    if isinstance(status, int):
        metadata["status"] = str(status)
    return metadata