"""
Metrics persistence store for the kvict Proxy.

Stores per-request token usage and savings data in SQLite for the
dashboard and badge APIs. Supports async operations via aiosqlite.

Production deployments can swap to Redis via PROXY_STORE_BACKEND=redis.
"""

from __future__ import annotations

import logging
import sqlite3
import time
import uuid
from contextlib import contextmanager
from typing import Any, Dict, List, Optional

from proxy.config import ProxyConfig, get_config
from proxy.schemas import SavingsHistoryEntry, SavingsSummary, UsageByModel

logger = logging.getLogger(__name__)


class MetricsStore:
    """SQLite-backed metrics store for proxy request data.

    Tracks every proxied request and maintains daily rollups for
    efficient dashboard queries.
    """

    def __init__(self, config: ProxyConfig | None = None):
        self._config = config or get_config()
        self._db_path = self._config.db_path
        self._init_db()

    def _init_db(self) -> None:
        """Create tables if they don't exist."""
        with self._connect() as conn:
            conn.executescript("""
                CREATE TABLE IF NOT EXISTS requests (
                    id TEXT PRIMARY KEY,
                    timestamp REAL NOT NULL,
                    provider TEXT NOT NULL,
                    model TEXT NOT NULL,
                    prompt_tokens_actual INTEGER NOT NULL,
                    prompt_tokens_baseline INTEGER NOT NULL,
                    completion_tokens INTEGER NOT NULL,
                    cached_tokens INTEGER DEFAULT 0,
                    cache_hit INTEGER DEFAULT 0,
                    latency_ms REAL DEFAULT 0.0,
                    cost_actual_usd REAL DEFAULT 0.0,
                    cost_baseline_usd REAL DEFAULT 0.0,
                    user_id TEXT
                );

                CREATE TABLE IF NOT EXISTS savings_daily (
                    date TEXT PRIMARY KEY,
                    total_prompt_tokens_actual INTEGER DEFAULT 0,
                    total_prompt_tokens_baseline INTEGER DEFAULT 0,
                    total_completion_tokens INTEGER DEFAULT 0,
                    total_cached_tokens INTEGER DEFAULT 0,
                    total_requests INTEGER DEFAULT 0,
                    cache_hit_count INTEGER DEFAULT 0,
                    cost_actual_usd REAL DEFAULT 0.0,
                    cost_baseline_usd REAL DEFAULT 0.0
                );

                CREATE INDEX IF NOT EXISTS idx_requests_timestamp
                    ON requests(timestamp);
                CREATE INDEX IF NOT EXISTS idx_requests_model
                    ON requests(model);
                """)

    @contextmanager
    def _connect(self):
        conn = sqlite3.connect(self._db_path)
        conn.row_factory = sqlite3.Row
        try:
            yield conn
            conn.commit()
        finally:
            conn.close()

    def record_request(
        self,
        provider: str,
        model: str,
        prompt_tokens_actual: int,
        prompt_tokens_baseline: int,
        completion_tokens: int,
        cached_tokens: int,
        cache_hit: bool,
        latency_ms: float,
        cost_actual_usd: float,
        cost_baseline_usd: float,
        user_id: str = "",
    ) -> str:
        """Record a single proxied request.

        Also updates the daily rollup table.

        Returns:
            The request ID.
        """
        request_id = uuid.uuid4().hex
        now = time.time()
        date_str = time.strftime("%Y-%m-%d", time.gmtime(now))

        with self._connect() as conn:
            conn.execute(
                """
                INSERT INTO requests
                    (id, timestamp, provider, model, prompt_tokens_actual,
                     prompt_tokens_baseline, completion_tokens, cached_tokens,
                     cache_hit, latency_ms, cost_actual_usd, cost_baseline_usd, user_id)
                VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)
                """,
                (
                    request_id,
                    now,
                    provider,
                    model,
                    prompt_tokens_actual,
                    prompt_tokens_baseline,
                    completion_tokens,
                    cached_tokens,
                    1 if cache_hit else 0,
                    latency_ms,
                    cost_actual_usd,
                    cost_baseline_usd,
                    user_id,
                ),
            )

            # Upsert daily rollup
            conn.execute(
                """
                INSERT INTO savings_daily
                    (date, total_prompt_tokens_actual, total_prompt_tokens_baseline,
                     total_completion_tokens, total_cached_tokens,
                     total_requests, cache_hit_count,
                     cost_actual_usd, cost_baseline_usd)
                VALUES (?, ?, ?, ?, ?, 1, ?, ?, ?)
                ON CONFLICT(date) DO UPDATE SET
                    total_prompt_tokens_actual = total_prompt_tokens_actual + excluded.total_prompt_tokens_actual,
                    total_prompt_tokens_baseline = total_prompt_tokens_baseline + excluded.total_prompt_tokens_baseline,
                    total_completion_tokens = total_completion_tokens + excluded.total_completion_tokens,
                    total_cached_tokens = total_cached_tokens + excluded.total_cached_tokens,
                    total_requests = total_requests + 1,
                    cache_hit_count = cache_hit_count + excluded.cache_hit_count,
                    cost_actual_usd = cost_actual_usd + excluded.cost_actual_usd,
                    cost_baseline_usd = cost_baseline_usd + excluded.cost_baseline_usd
                """,
                (
                    date_str,
                    prompt_tokens_actual,
                    prompt_tokens_baseline,
                    completion_tokens,
                    cached_tokens,
                    1 if cache_hit else 0,
                    cost_actual_usd,
                    cost_baseline_usd,
                ),
            )

        return request_id

    def get_savings_summary(self) -> SavingsSummary:
        """Get aggregate savings across all time.

        Used by the dashboard overview and badge API.
        """
        with self._connect() as conn:
            row = conn.execute("""
                SELECT
                    COALESCE(SUM(total_cached_tokens), 0) as tokens_saved,
                    COALESCE(SUM(cost_baseline_usd - cost_actual_usd), 0) as dollars_saved,
                    COALESCE(SUM(total_requests), 0) as total_requests,
                    COALESCE(SUM(cache_hit_count), 0) as cache_hits
                FROM savings_daily
                """).fetchone()

        total_requests = row["total_requests"]
        cache_hits = row["cache_hits"]
        tokens_saved = row["tokens_saved"]
        avg_per_msg = self._config.avg_tokens_per_message

        return SavingsSummary(
            tokens_saved_total=tokens_saved,
            messages_gained=tokens_saved // avg_per_msg if avg_per_msg > 0 else 0,
            dollars_saved=round(row["dollars_saved"], 4),
            cache_hit_rate=cache_hits / total_requests if total_requests > 0 else 0.0,
            total_requests=total_requests,
            cache_hits=cache_hits,
            cache_misses=total_requests - cache_hits,
        )

    def get_tokens_saved_today(self) -> int:
        """Get tokens saved today only.

        Used by the badge API for the "tokens saved today" metric.
        """
        today = time.strftime("%Y-%m-%d", time.gmtime())
        with self._connect() as conn:
            row = conn.execute(
                """
                SELECT COALESCE(total_cached_tokens, 0) as tokens_saved
                FROM savings_daily
                WHERE date = ?
                """,
                (today,),
            ).fetchone()
        return row["tokens_saved"] if row else 0

    def get_savings_history(self, days: int = 30) -> List[SavingsHistoryEntry]:
        """Get daily savings time series.

        Args:
            days: Number of days to look back.

        Returns:
            List of SavingsHistoryEntry sorted by date.
        """
        cutoff = time.time() - (days * 86400)
        cutoff_date = time.strftime("%Y-%m-%d", time.gmtime(cutoff))

        with self._connect() as conn:
            rows = conn.execute(
                """
                SELECT date, total_cached_tokens, total_requests,
                       cache_hit_count, cost_baseline_usd, cost_actual_usd
                FROM savings_daily
                WHERE date >= ?
                ORDER BY date
                """,
                (cutoff_date,),
            ).fetchall()

        result = []
        for r in rows:
            total = r["total_requests"]
            hits = r["cache_hit_count"]
            result.append(
                SavingsHistoryEntry(
                    date=r["date"],
                    tokens_saved=r["total_cached_tokens"],
                    dollars_saved=round(
                        r["cost_baseline_usd"] - r["cost_actual_usd"], 4
                    ),
                    requests=total,
                    cache_hit_rate=hits / total if total > 0 else 0.0,
                )
            )
        return result

    def get_usage_by_model(self) -> List[UsageByModel]:
        """Get per-model usage breakdown."""
        with self._connect() as conn:
            rows = conn.execute("""
                SELECT
                    model,
                    SUM(prompt_tokens_actual + completion_tokens) as total_tokens,
                    SUM(cached_tokens) as tokens_saved,
                    SUM(cost_baseline_usd - cost_actual_usd) as dollars_saved,
                    COUNT(*) as requests,
                    SUM(cache_hit) as cache_hits
                FROM requests
                GROUP BY model
                ORDER BY total_tokens DESC
                """).fetchall()

        result = []
        for r in rows:
            total = r["requests"]
            hits = r["cache_hits"]
            result.append(
                UsageByModel(
                    model=r["model"],
                    total_tokens=r["total_tokens"],
                    tokens_saved=r["tokens_saved"],
                    dollars_saved=round(r["dollars_saved"], 4),
                    requests=total,
                    cache_hit_rate=hits / total if total > 0 else 0.0,
                )
            )
        return result
