from __future__ import annotations

import json
from collections.abc import Sequence
from decimal import Decimal, ROUND_HALF_UP

from sqlalchemy.orm import Session

from app.models.ai_config import ModelConfig
from app.models.logs import AiRequestLog
from app.services.rag_service import RetrievedChunk


class AiRequestLogService:
    @staticmethod
    def write_success(
        db: Session,
        *,
        session_id: int | None,
        message_id: int | None,
        user_id: int | None,
        model_name: str,
        prompt: str,
        knowledge_ids: str,
        retrieve_count: int,
        input_token: int,
        output_token: int,
        cost_ms: int,
        retrieved_chunks: Sequence[RetrievedChunk] | None = None,
        model_id: int | None = None,
        route_reason: str | None = None,
        question_type: str | None = None,
    ) -> None:
        model = db.get(ModelConfig, model_id) if model_id else None
        estimated_cost, currency = _estimate_cost(model, input_token, output_token)
        db.add(
            AiRequestLog(
                session_id=session_id,
                message_id=message_id,
                user_id=user_id,
                model_id=model_id,
                model_name=model_name,
                prompt=prompt,
                knowledge_ids=knowledge_ids,
                retrieve_count=retrieve_count,
                retrieved_chunks=_dump_retrieved_chunks(retrieved_chunks),
                input_token=input_token,
                output_token=output_token,
                total_token=input_token + output_token,
                cost_ms=cost_ms,
                estimated_cost=estimated_cost,
                currency=currency,
                route_reason=route_reason or _default_route_reason(retrieve_count),
                question_type=question_type or _question_type(retrieve_count),
                knowledge_hit=1 if retrieve_count > 0 else 0,
                status="SUCCESS",
            )
        )

    @staticmethod
    def write_failed(
        db: Session,
        *,
        session_id: int | None,
        message_id: int | None,
        user_id: int | None,
        model_name: str | None,
        prompt: str | None,
        knowledge_ids: str | None,
        retrieve_count: int,
        cost_ms: int,
        error_message: str,
        retrieved_chunks: Sequence[RetrievedChunk] | None = None,
        model_id: int | None = None,
        route_reason: str | None = None,
        question_type: str | None = None,
    ) -> None:
        db.add(
            AiRequestLog(
                session_id=session_id,
                message_id=message_id,
                user_id=user_id,
                model_id=model_id,
                model_name=model_name,
                prompt=prompt,
                knowledge_ids=knowledge_ids,
                retrieve_count=retrieve_count,
                retrieved_chunks=_dump_retrieved_chunks(retrieved_chunks),
                cost_ms=cost_ms,
                route_reason=route_reason or _default_route_reason(retrieve_count),
                question_type=question_type or _question_type(retrieve_count),
                knowledge_hit=1 if retrieve_count > 0 else 0,
                status="FAILED",
                error_message=error_message,
            )
        )


def _dump_retrieved_chunks(chunks: Sequence[RetrievedChunk] | None) -> str | None:
    if not chunks:
        return None
    payload = [
        {
            "index": index,
            "knowledgeId": chunk.knowledge_id,
            "knowledgeName": chunk.knowledge_name,
            "title": chunk.title,
            "content": chunk.content,
            "sourceUrl": chunk.source_url,
        }
        for index, chunk in enumerate(chunks, start=1)
    ]
    return json.dumps(payload, ensure_ascii=False)


def _estimate_cost(model: ModelConfig | None, input_token: int, output_token: int) -> tuple[Decimal | None, str | None]:
    if model is None:
        return None, None
    input_price = Decimal(model.input_price_per_1k or 0)
    output_price = Decimal(model.output_price_per_1k or 0)
    if input_price <= 0 and output_price <= 0:
        return None, model.currency
    value = (Decimal(input_token) / Decimal(1000) * input_price) + (Decimal(output_token) / Decimal(1000) * output_price)
    return value.quantize(Decimal("0.000001"), rounding=ROUND_HALF_UP), model.currency


def _default_route_reason(retrieve_count: int) -> str:
    return "当前默认正式模型；本轮命中知识库" if retrieve_count > 0 else "当前默认正式模型；本轮未命中知识库"


def _question_type(retrieve_count: int) -> str:
    return "knowledge_grounded" if retrieve_count > 0 else "general_chat"
