"""把 RuntimeChat provider attempt 接到独立事务商业计量。"""

from __future__ import annotations

from datetime import UTC, datetime
from typing import Any

from sqlalchemy import select
from sqlalchemy.orm import Session, sessionmaker

from app.core.logging import get_logger
from app.models.agent_run import AgentRun
from app.services.commercial_direct_operation import (
    CommercialDirectOperationBridge,
    DirectOperationIdentity,
)
from app.services.commercial_runtime_policy import tenant_from_agent_route
from app.services.runtime_chat_attempts import (
    RuntimeChatAttemptCompletedEvent,
    RuntimeChatAttemptObserver,
    RuntimeChatAttemptPermit,
    RuntimeChatAttemptPermitEvent,
    RuntimeChatObserverFailure,
    RuntimeChatOperationContext,
)

logger = get_logger("app.services.runtime_chat_commercial")


class CommercialDirectReconciliationRequired(RuntimeError):
    """补偿状态已持久化；由 RuntimeChat trace 显式标记而不吞掉回答。"""


class CommercialRuntimeChatAttemptObserver(RuntimeChatAttemptObserver):
    def __init__(self, bridge: CommercialDirectOperationBridge) -> None:
        self.bridge = bridge

    def on_permit(
        self,
        event: RuntimeChatAttemptPermitEvent,
    ) -> RuntimeChatAttemptPermit:
        permit = self.bridge.permit(_identity(event.identity, event.started_at))
        return RuntimeChatAttemptPermit(
            allowed=permit.allowed,
            reason="" if permit.allowed else permit.reason,
        )

    def on_completed(self, event: RuntimeChatAttemptCompletedEvent) -> None:
        result = self.bridge.complete(
            _identity(event.identity, event.started_at),
            outcome=event.outcome,
            authoritative_quantities=_quantities(event),
            completed_at=event.completed_at,
            usage_source=event.usage.source,
            usage_availability=event.usage.availability,
        )
        if result.requires_reconciliation:
            raise CommercialDirectReconciliationRequired(
                f"{result.reason_code}: {result.reason}"
            )

    def on_observer_failure(self, failure: RuntimeChatObserverFailure) -> None:
        if failure.phase != "completion":
            return
        if failure.error_type == CommercialDirectReconciliationRequired.__name__:
            return
        result = self.bridge.complete(
            _identity(failure.identity, failure.occurred_at),
            outcome="outcome_unknown",
            authoritative_quantities={},
            completed_at=failure.occurred_at,
            usage_source="observer_failure",
            usage_availability="unavailable",
        )
        if result.requires_reconciliation:
            logger.warning(
                "RuntimeChat commercial observer persisted compensation "
                "operation_call_id=%s reason_code=%s",
                result.operation_call_id,
                result.reason_code,
            )


def build_runtime_chat_commercial_observer(
    db: Session,
) -> CommercialRuntimeChatAttemptObserver:
    """只复用 Engine；permit/completion 从不提交传入的业务 Session。"""

    factory = sessionmaker(bind=db.get_bind(), expire_on_commit=False)
    return CommercialRuntimeChatAttemptObserver(
        CommercialDirectOperationBridge(factory, lookup_session=db),
    )


def trusted_runtime_chat_operation_context(
    db: Session,
    *,
    run_id: str,
    attempt_scope: str,
    invocation_seq: int = 1,
) -> RuntimeChatOperationContext | None:
    """只从 AgentRun.route_json 解析租户，不相信 prompt/context_json 身份。"""

    normalized_run_id = str(run_id or "").strip()
    if not normalized_run_id:
        return None
    run = db.scalar(select(AgentRun).where(AgentRun.run_id == normalized_run_id))
    if run is None:
        return None
    tenant_id = tenant_from_agent_route(run.route_json)
    if tenant_id is None:
        return None
    return RuntimeChatOperationContext(
        tenant_id=tenant_id,
        operation_id=f"agent-run:{run.run_id}",
        invocation_seq=invocation_seq,
        attempt_scope=str(attempt_scope or "").strip(),
        run_id=run.run_id,
    )


def _identity(raw: Any, started_at: datetime) -> DirectOperationIdentity:
    return DirectOperationIdentity(
        tenant_id=raw.tenant_id,
        operation_key=raw.attempt_key,
        run_key=raw.run_id or raw.operation_id,
        tool_type="llm",
        tool_name="chat.completions",
        provider=raw.provider,
        model_name=raw.model,
        started_at=_utc(started_at),
    )


def _quantities(event: RuntimeChatAttemptCompletedEvent) -> dict[str, int | None]:
    usage = event.usage
    input_tokens = (
        usage.prompt_tokens
        if usage.prompt_tokens is not None
        else usage.prompt_eval_count
    )
    output_tokens = (
        usage.completion_tokens
        if usage.completion_tokens is not None
        else usage.eval_count
    )
    total_tokens = usage.total_tokens
    if total_tokens is None and input_tokens is not None and output_tokens is not None:
        total_tokens = input_tokens + output_tokens
    return {
        "input_tokens": input_tokens,
        "output_tokens": output_tokens,
        "total_tokens": total_tokens,
    }


def _utc(value: datetime) -> datetime:
    return value.replace(tzinfo=UTC) if value.tzinfo is None else value.astimezone(UTC)
