"""
usage_store.py
--------------
Pluggable backend for usage tracking.

Supports SQLite for single-node setups and Redis for distributed setups.
"""

from __future__ import annotations

import sqlite3
from abc import ABC, abstractmethod
from datetime import datetime
from pathlib import Path
from threading import Lock
from typing import Any, Dict

try:
    import redis
except ImportError:
    redis = None


class UsageStore(ABC):
    """Abstract interface for usage stats storage."""

    @abstractmethod
    def increment_counter(
        self, tenant_id: str, operation: str, period: str, period_type: str
    ) -> None:
        """Increment usage counter atomically."""
        pass

    @abstractmethod
    def get_count(
        self, tenant_id: str, operation: str, period: str, period_type: str
    ) -> int:
        """Get current usage count."""
        pass

    @abstractmethod
    def get_all_tenant_usage(self) -> Dict[str, Any]:
        """Get raw usage data (implementation specific structure)."""
        pass


class SQLiteUsageStore(UsageStore):
    """Local usage tracking using SQLite."""

    def __init__(self, db_path: str):
        self._path = Path(db_path)
        self._path.parent.mkdir(parents=True, exist_ok=True)
        self._lock = Lock()
        self._init_db()

    def _connect(self) -> sqlite3.Connection:
        return sqlite3.connect(self._path)

    def _init_db(self) -> None:
        with self._connect() as conn:
            conn.executescript(
                """
                CREATE TABLE IF NOT EXISTS usage_stats (
                    id INTEGER PRIMARY KEY AUTOINCREMENT,
                    tenant_id TEXT NOT NULL,
                    operation TEXT NOT NULL,
                    count INTEGER NOT NULL DEFAULT 0,
                    period TEXT NOT NULL,
                    period_type TEXT NOT NULL,
                    created_at TEXT NOT NULL,
                    updated_at TEXT NOT NULL,
                    UNIQUE(tenant_id, operation, period, period_type)
                );
                CREATE INDEX IF NOT EXISTS idx_usage_tenant_period
                    ON usage_stats(tenant_id, period, period_type);
                """
            )
            conn.commit()

    def increment_counter(
        self, tenant_id: str, operation: str, period: str, period_type: str
    ) -> None:
        now = datetime.utcnow().isoformat()
        with self._lock:
            with self._connect() as conn:
                conn.execute(
                    """
                    INSERT INTO usage_stats (
                        tenant_id, operation, count, period, period_type,
                        created_at, updated_at
                    ) VALUES (?, ?, 1, ?, ?, ?, ?)
                    ON CONFLICT(tenant_id, operation, period, period_type) DO UPDATE SET
                        count = count + 1,
                        updated_at = excluded.updated_at
                    """,
                    (tenant_id, operation, period, period_type, now, now),
                )
                conn.commit()

    def get_count(
        self, tenant_id: str, operation: str, period: str, period_type: str
    ) -> int:
        with self._connect() as conn:
            row = conn.execute(
                """
                SELECT count FROM usage_stats
                WHERE tenant_id = ? AND operation = ? AND period = ? AND period_type = ?
                """,
                (tenant_id, operation, period, period_type),
            ).fetchone()
            return row[0] if row else 0

    def get_all_tenant_usage(self) -> Dict[str, Any]:
        # Minimal implementation for admin view
        with self._connect() as conn:
            # Just return a list of tenants for simplicity in this abstract view
            rows = conn.execute("SELECT DISTINCT tenant_id FROM usage_stats").fetchall()
        return {row[0]: {} for row in rows}


class RedisUsageStore(UsageStore):
    """Distributed usage tracking using Redis."""

    def __init__(self, redis_url: str):
        if redis is None:
            raise ImportError("Redis package is required for RedisUsageStore")
        self._redis = redis.from_url(redis_url, decode_responses=True)
        self._ttl = 86400 * 60  # 60 days retention for stats

    def _get_key(self, tenant_id: str, operation: str, period: str) -> str:
        return f"usage:{tenant_id}:{operation}:{period}"

    def increment_counter(
        self, tenant_id: str, operation: str, period: str, period_type: str
    ) -> None:
        key = self._get_key(tenant_id, operation, period)
        pipe = self._redis.pipeline()
        pipe.incr(key)
        pipe.expire(key, self._ttl)
        pipe.sadd("tenants:active", tenant_id)
        pipe.execute()

    def get_count(
        self, tenant_id: str, operation: str, period: str, period_type: str
    ) -> int:
        key = self._get_key(tenant_id, operation, period)
        val = self._redis.get(key)
        return int(val) if val else 0

    def get_all_tenant_usage(self) -> Dict[str, Any]:
        tenants = self._redis.smembers("tenants:active")
        return {t: {} for t in tenants}
