from __future__ import annotations

from sqlalchemy.orm import Session

from app.models.logs import AiRequestLog


class AiRequestLogService:
    @staticmethod
    def write_success(
        db: Session,
        *,
        session_id: int,
        message_id: int | None,
        user_id: int,
        model_name: str,
        prompt: str,
        knowledge_ids: str,
        retrieve_count: int,
        input_token: int,
        output_token: int,
        cost_ms: int,
    ) -> None:
        db.add(
            AiRequestLog(
                session_id=session_id,
                message_id=message_id,
                user_id=user_id,
                model_name=model_name,
                prompt=prompt,
                knowledge_ids=knowledge_ids,
                retrieve_count=retrieve_count,
                input_token=input_token,
                output_token=output_token,
                total_token=input_token + output_token,
                cost_ms=cost_ms,
                status="SUCCESS",
            )
        )

    @staticmethod
    def write_failed(
        db: Session,
        *,
        session_id: int,
        message_id: int | None,
        user_id: int,
        model_name: str | None,
        prompt: str | None,
        knowledge_ids: str | None,
        retrieve_count: int,
        cost_ms: int,
        error_message: str,
    ) -> None:
        db.add(
            AiRequestLog(
                session_id=session_id,
                message_id=message_id,
                user_id=user_id,
                model_name=model_name,
                prompt=prompt,
                knowledge_ids=knowledge_ids,
                retrieve_count=retrieve_count,
                cost_ms=cost_ms,
                status="FAILED",
                error_message=error_message,
            )
        )
