from __future__ import annotations

from datetime import datetime, timezone


class LLMAuditRepositoryMixin:
    def record_llm_usage(
        self,
        user_id: int,
        feature: str,
        source: str,
        model: str,
        status: str,
        latency_ms: int = 0,
        *,
        role: str = "",
        prompt_version: str = "",
        error_code: str = "",
        input_tokens: int = 0,
        output_tokens: int = 0,
    ) -> None:
        now = datetime.now(timezone.utc).isoformat(timespec="seconds")
        with self.connect() as connection:
            connection.execute(
                """
                INSERT INTO llm_usage
                    (user_id, feature, source, model, status, latency_ms, created_at,
                     role, prompt_version, error_code, input_tokens, output_tokens)
                VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)
                """,
                (
                    user_id, feature, source, model, status, int(latency_ms), now,
                    role, prompt_version, error_code, int(input_tokens), int(output_tokens),
                ),
            )

    def count_llm_usage_since(self, user_id: int, source: str, since: str) -> int:
        with self.connect() as connection:
            row = connection.execute(
                """
                SELECT COUNT(*) AS total FROM llm_usage
                WHERE user_id = ? AND source = ? AND created_at >= ?
                """,
                (user_id, source, since),
            ).fetchone()
        return int(row["total"] if row else 0)