"""
cortex/api/stats_db.py — Shared stats database for cross-user metrics and usage metering.

Separate from per-user memory databases. Used for:
- API request logging (per-user call counts, latency)
- Usage aggregation for Stripe billing
- Dashboard overview metrics
- System health tracking

Schema:
    request_log: raw per-request data (pruned after 30 days)
    usage_daily: aggregated daily totals per user
    system_metrics: periodic system health snapshots

Dependencies: aiosqlite (matches existing codebase).
"""
from __future__ import annotations

import asyncio
import logging
from typing import Any, Dict, List, Optional

import aiosqlite

logger = logging.getLogger(__name__)

_SCHEMA_SQL = """
PRAGMA journal_mode=WAL;
PRAGMA busy_timeout=5000;

CREATE TABLE IF NOT EXISTS request_log (
    id INTEGER PRIMARY KEY AUTOINCREMENT,
    owner_id TEXT NOT NULL,
    endpoint TEXT NOT NULL,
    method TEXT NOT NULL,
    status_code INTEGER NOT NULL,
    latency_ms REAL NOT NULL,
    tokens_used INTEGER DEFAULT 0,
    created_at TEXT NOT NULL DEFAULT (strftime('%Y-%m-%dT%H:%M:%fZ', 'now'))
);
CREATE INDEX IF NOT EXISTS idx_request_log_owner_date ON request_log(owner_id, created_at);
CREATE INDEX IF NOT EXISTS idx_request_log_created ON request_log(created_at);

CREATE TABLE IF NOT EXISTS usage_daily (
    id INTEGER PRIMARY KEY AUTOINCREMENT,
    owner_id TEXT NOT NULL,
    date TEXT NOT NULL,
    call_count INTEGER NOT NULL DEFAULT 0,
    tokens_total INTEGER NOT NULL DEFAULT 0,
    avg_latency_ms REAL,
    p95_latency_ms REAL,
    UNIQUE(owner_id, date)
);
CREATE INDEX IF NOT EXISTS idx_usage_daily_owner ON usage_daily(owner_id, date);

CREATE TABLE IF NOT EXISTS system_metrics (
    id INTEGER PRIMARY KEY AUTOINCREMENT,
    metric_name TEXT NOT NULL,
    metric_value REAL NOT NULL,
    recorded_at TEXT NOT NULL DEFAULT (strftime('%Y-%m-%dT%H:%M:%fZ', 'now'))
);
"""


class StatsDB:
    """Manages the shared stats.db SQLite database."""

    def __init__(self, db_path: str = "/data/stats.db") -> None:
        self._db_path = db_path
        self._conn: Optional[aiosqlite.Connection] = None
        self._lock = asyncio.Lock()

    async def initialize(self) -> None:
        """Create tables if not exist. WAL mode enabled."""
        self._conn = await aiosqlite.connect(self._db_path)
        self._conn.row_factory = aiosqlite.Row
        await self._conn.executescript(_SCHEMA_SQL)
        await self._conn.commit()
        logger.info("StatsDB initialized at %s", self._db_path)

    async def record_request(
        self,
        owner_id: str,
        endpoint: str,
        method: str,
        status_code: int,
        latency_ms: float,
        tokens_used: int = 0,
    ) -> None:
        """Record a single API request. Called by middleware."""
        if self._conn is None:
            return
        try:
            await self._conn.execute(
                """
                INSERT INTO request_log (owner_id, endpoint, method, status_code, latency_ms, tokens_used)
                VALUES (?, ?, ?, ?, ?, ?)
                """,
                (owner_id, endpoint, method, status_code, latency_ms, tokens_used),
            )
            await self._conn.commit()
        except Exception:
            logger.exception("StatsDB.record_request failed")

    async def get_usage_daily(self, owner_id: str, days: int = 30) -> List[Dict[str, Any]]:
        """Get daily usage for a specific user (for billing)."""
        if self._conn is None:
            return []
        cursor = await self._conn.execute(
            """
            SELECT date, call_count, tokens_total, avg_latency_ms, p95_latency_ms
            FROM usage_daily
            WHERE owner_id = ?
              AND date >= date('now', ? || ' days')
            ORDER BY date DESC
            """,
            (owner_id, f"-{days}"),
        )
        rows = await cursor.fetchall()
        return [dict(r) for r in rows]

    async def get_usage_summary(self, owner_id: str, month: Optional[str] = None) -> Dict[str, Any]:
        """Monthly usage summary for Stripe reporting.

        Args:
            owner_id: User identifier.
            month: ISO month string like '2026-03'. Defaults to current month.

        Returns:
            dict with call_count, tokens_total, avg_latency_ms, period.
        """
        if self._conn is None:
            return {"call_count": 0, "tokens_total": 0, "avg_latency_ms": 0.0, "period": month or ""}

        if month is None:
            # current month via SQL
            cursor = await self._conn.execute("SELECT strftime('%Y-%m', 'now')")
            row = await cursor.fetchone()
            month = row[0] if row else ""

        cursor = await self._conn.execute(
            """
            SELECT
                SUM(call_count)     AS call_count,
                SUM(tokens_total)   AS tokens_total,
                AVG(avg_latency_ms) AS avg_latency_ms
            FROM usage_daily
            WHERE owner_id = ?
              AND strftime('%Y-%m', date) = ?
            """,
            (owner_id, month),
        )
        row = await cursor.fetchone()
        return {
            "call_count": int(row["call_count"] or 0),
            "tokens_total": int(row["tokens_total"] or 0),
            "avg_latency_ms": float(row["avg_latency_ms"] or 0.0),
            "period": month,
        }

    async def get_system_overview(self) -> Dict[str, Any]:
        """System-wide stats for dashboard overview.

        Returns:
            dict with total_users, total_requests_24h, avg_latency_ms, active_users_24h.
        """
        if self._conn is None:
            return {"total_users": 0, "total_requests_24h": 0, "avg_latency_ms": 0.0, "active_users_24h": 0}

        cursor = await self._conn.execute(
            """
            SELECT
                COUNT(DISTINCT owner_id)  AS total_users,
                COUNT(*)                  AS total_requests_24h,
                AVG(latency_ms)           AS avg_latency_ms,
                COUNT(DISTINCT CASE
                    WHEN created_at >= strftime('%Y-%m-%dT%H:%M:%fZ', 'now', '-1 day')
                    THEN owner_id END)    AS active_users_24h
            FROM request_log
            WHERE created_at >= strftime('%Y-%m-%dT%H:%M:%fZ', 'now', '-1 day')
            """
        )
        row = await cursor.fetchone()
        # total_users should be across all time, not just 24h — do a second query
        cursor2 = await self._conn.execute("SELECT COUNT(DISTINCT owner_id) FROM request_log")
        row2 = await cursor2.fetchone()
        return {
            "total_users": int(row2[0] or 0),
            "total_requests_24h": int(row["total_requests_24h"] or 0),
            "avg_latency_ms": float(row["avg_latency_ms"] or 0.0),
            "active_users_24h": int(row["active_users_24h"] or 0),
        }

    async def aggregate_daily(self) -> None:
        """Roll up request_log into usage_daily. Run nightly via cron/circadian cycle.

        p95 latency is computed per (owner_id, date) by fetching sorted latencies
        and selecting the 95th-percentile entry in Python (SQLite <3.43 lacks
        PERCENTILE_CONT; correlated-subquery COUNT() is disallowed in SQLite).
        """
        if self._conn is None:
            return
        async with self._lock:
            # Step 1: upsert call_count / tokens_total / avg_latency from SQL aggregate
            await self._conn.execute(
                """
                INSERT INTO usage_daily (owner_id, date, call_count, tokens_total, avg_latency_ms)
                SELECT
                    owner_id,
                    strftime('%Y-%m-%d', created_at) AS date,
                    COUNT(*)                          AS call_count,
                    SUM(tokens_used)                  AS tokens_total,
                    AVG(latency_ms)                   AS avg_latency_ms
                FROM request_log
                GROUP BY owner_id, strftime('%Y-%m-%d', created_at)
                ON CONFLICT(owner_id, date) DO UPDATE SET
                    call_count     = excluded.call_count,
                    tokens_total   = excluded.tokens_total,
                    avg_latency_ms = excluded.avg_latency_ms
                """
            )
            await self._conn.commit()

            # Step 2: compute p95 per (owner_id, date) in Python and update
            cursor = await self._conn.execute(
                """
                SELECT owner_id, strftime('%Y-%m-%d', created_at) AS date,
                       GROUP_CONCAT(latency_ms ORDER BY latency_ms) AS latencies_csv
                FROM request_log
                GROUP BY owner_id, date
                """
            )
            rows = await cursor.fetchall()
            updates = []
            for row in rows:
                latencies = [float(v) for v in row["latencies_csv"].split(",") if v]
                if latencies:
                    idx = max(0, int(len(latencies) * 0.95) - 1)
                    p95 = sorted(latencies)[idx]
                    updates.append((p95, row["owner_id"], row["date"]))
            if updates:
                await self._conn.executemany(
                    "UPDATE usage_daily SET p95_latency_ms = ? WHERE owner_id = ? AND date = ?",
                    updates,
                )
                await self._conn.commit()

            logger.info("StatsDB.aggregate_daily complete")

    async def prune_old_requests(self, days: int = 30) -> int:
        """Delete request_log entries older than N days. Returns row count deleted."""
        if self._conn is None:
            return 0
        cursor = await self._conn.execute(
            """
            DELETE FROM request_log
            WHERE created_at < strftime('%Y-%m-%dT%H:%M:%fZ', 'now', ? || ' days')
            """,
            (f"-{days}",),
        )
        await self._conn.commit()
        deleted = cursor.rowcount
        logger.info("StatsDB.prune_old_requests: deleted %d rows older than %d days", deleted, days)
        return deleted

    async def close(self) -> None:
        """Close the database connection."""
        if self._conn is not None:
            await self._conn.close()
            self._conn = None
            logger.info("StatsDB closed")
