from __future__ import annotations

from time import perf_counter

from sqlalchemy.orm import Session

from app.services.ai_request_log_service import AiRequestLogService
from app.services.model_routing_service import ModelRoutingService, ModelScenario
from app.services.model_service import ModelClientService, ModelCompletion


class TrackedGenerationService:
    """Runs non-chat generation and records its actual model, tokens and cost."""

    @staticmethod
    def generate(
        db: Session,
        *,
        prompt: str,
        scenario: ModelScenario,
        user_id: int | None,
    ) -> ModelCompletion:
        started = perf_counter()
        try:
            completion = ModelClientService.generate_text_for_scenario(db, prompt, scenario)
            AiRequestLogService.write_success(
                db,
                session_id=None,
                message_id=None,
                user_id=user_id,
                model_id=completion.model_id,
                model_name=completion.model_name,
                prompt=prompt,
                knowledge_ids="",
                retrieve_count=0,
                input_token=completion.input_token,
                output_token=completion.output_token,
                cost_ms=_elapsed_ms(started),
                route_reason=completion.route_reason,
                question_type=f"background_{scenario}",
            )
            return completion
        except Exception as exc:
            route = ModelRoutingService.resolve(db, scenario)
            AiRequestLogService.write_failed(
                db,
                session_id=None,
                message_id=None,
                user_id=user_id,
                model_id=route.model.id if route.model is not None else None,
                model_name=route.model.model_name if route.model is not None else None,
                prompt=prompt,
                knowledge_ids="",
                retrieve_count=0,
                cost_ms=_elapsed_ms(started),
                route_reason=route.reason,
                question_type=f"background_{scenario}",
                error_message=str(exc)[:2000],
            )
            raise


def _elapsed_ms(started: float) -> int:
    return max(0, int((perf_counter() - started) * 1000))
