from __future__ import annotations

import sqlite3


class LLMRepository:
    @staticmethod
    def usage_today(connection: sqlite3.Connection, user_id: int, usage_date: str) -> int:
        row = connection.execute(
            "SELECT successful_calls FROM llm_usage_daily WHERE user_id = ? AND usage_date = ?",
            (user_id, usage_date),
        ).fetchone()
        return int(row["successful_calls"]) if row else 0

    @staticmethod
    def active_today(connection: sqlite3.Connection, user_id: int, day_prefix: str) -> int:
        row = connection.execute(
            """
            SELECT COUNT(*) AS total FROM llm_requests
            WHERE user_id = ? AND started_at LIKE ? AND status IN ('reserved', 'streaming')
            """,
            (user_id, f"{day_prefix}%"),
        ).fetchone()
        return int(row["total"])

    @staticmethod
    def reserve(
        connection: sqlite3.Connection,
        *,
        request_id: str,
        user_id: int,
        feature: str,
        business_id: str,
        prompt_version: str,
        started_at: str,
        input_chars: int,
    ) -> None:
        connection.execute(
            """
            INSERT INTO llm_requests (
                id, user_id, feature, business_id, prompt_version,
                status, started_at, input_chars
            ) VALUES (?, ?, ?, ?, ?, 'reserved', ?, ?)
            """,
            (
                request_id,
                user_id,
                feature,
                business_id,
                prompt_version,
                started_at,
                input_chars,
            ),
        )

    @staticmethod
    def set_request_status(
        connection: sqlite3.Connection,
        request_id: str,
        status: str,
        *,
        completed_at: str | None = None,
        duration_ms: int = 0,
        error_type: str = "",
        output_chars: int = 0,
    ) -> None:
        connection.execute(
            """
            UPDATE llm_requests SET status = ?, completed_at = ?, duration_ms = ?,
                error_type = ?, output_chars = ? WHERE id = ?
            """,
            (status, completed_at, duration_ms, error_type, output_chars, request_id),
        )

    @staticmethod
    def increment_usage(
        connection: sqlite3.Connection, user_id: int, usage_date: str, updated_at: str
    ) -> None:
        connection.execute(
            """
            INSERT INTO llm_usage_daily (user_id, usage_date, successful_calls, updated_at)
            VALUES (?, ?, 1, ?)
            ON CONFLICT(user_id, usage_date) DO UPDATE SET
                successful_calls = successful_calls + 1,
                updated_at = excluded.updated_at
            """,
            (user_id, usage_date, updated_at),
        )

    @staticmethod
    def start_attempt(
        connection: sqlite3.Connection,
        request_id: str,
        model_id: int,
        role: str,
        started_at: str,
        input_chars: int,
    ) -> int:
        cursor = connection.execute(
            """
            INSERT INTO llm_attempts (
                request_id, model_id, role, status, started_at, input_chars
            ) VALUES (?, ?, ?, 'streaming', ?, ?)
            """,
            (request_id, model_id, role, started_at, input_chars),
        )
        return int(cursor.lastrowid)

    @staticmethod
    def finish_attempt(
        connection: sqlite3.Connection,
        attempt_id: int,
        status: str,
        completed_at: str,
        duration_ms: int,
        error_type: str,
        output_chars: int,
    ) -> None:
        connection.execute(
            """
            UPDATE llm_attempts SET status = ?, completed_at = ?, duration_ms = ?,
                error_type = ?, output_chars = ? WHERE id = ?
            """,
            (status, completed_at, duration_ms, error_type, output_chars, attempt_id),
        )