"""Usage Store — Persistent usage tracking with Redis Hash backend.

Provides UsageStore Protocol with InMemoryUsageStore and RedisUsageStore
backends. Replaces the in-memory list in UsageTracker with a pluggable
store that survives restarts and supports multi-instance deployment.

Key schema (Redis):
    agentkit:usage:v2:{date}:{user_id}:{department_id}  → Hash: {agent_name:model → JSON(UsageBucket)}
    agentkit:usage_records:v2:{date}:{user_id}:{department_id}  → List: JSON(UsageRecord) with LTRIM

Legacy v1 keys (still readable for backward compat):
    agentkit:usage:{date}                                → Hash
    agentkit:usage_records:{date}                         → List
"""

import asyncio
import json
import logging
from dataclasses import dataclass, field
from datetime import datetime, timedelta, timezone
from typing import Protocol, runtime_checkable

from agentkit.llm.protocol import TokenUsage

logger = logging.getLogger(__name__)


class UsageStoreUnavailableError(Exception):
    """Raised when the usage store is unavailable for quota-critical queries.

    The gateway treats this as a fail-closed condition (HTTP 503) —
    refusing the request rather than allowing untracked usage.
    """


@dataclass
class UsageRecord:
    """使用量记录"""

    agent_name: str
    model: str
    prompt_tokens: int
    completion_tokens: int
    total_tokens: int
    cost: float
    latency_ms: float
    timestamp: str = ""  # ISO 8601 string for JSON serialization
    user_id: str | None = None
    department_id: str | None = None

    def __post_init__(self):
        if not self.timestamp:
            self.timestamp = datetime.now(timezone.utc).isoformat()


@dataclass
class UsageBucket:
    """Aggregated usage for an agent+model pair on a given date."""

    prompt_tokens: int = 0
    completion_tokens: int = 0
    total_tokens: int = 0
    cost: float = 0.0
    count: int = 0


@dataclass
class UsageSummary:
    """使用量汇总"""

    total_tokens: int = 0
    total_cost: float = 0.0
    by_model: dict[str, dict[str, int | float]] = field(default_factory=dict)
    records: list[UsageRecord] = field(default_factory=list)
    by_user: dict[str, dict[str, int | float]] = field(default_factory=dict)
    by_department: dict[str, dict[str, int | float]] = field(default_factory=dict)


# ---------------------------------------------------------------------------
# UsageStore Protocol
# ---------------------------------------------------------------------------


@runtime_checkable
class UsageStore(Protocol):
    """Persistent usage store interface."""

    def record(
        self,
        agent_name: str,
        model: str,
        usage: TokenUsage,
        cost: float,
        latency_ms: float,
        user_id: str | None = None,
        department_id: str | None = None,
    ) -> None:
        """Record a usage event (sync — may block in async contexts)."""
        ...

    def get_usage(
        self,
        agent_name: str | None = None,
        start_time: datetime | None = None,
        end_time: datetime | None = None,
        user_id: str | None = None,
        department_id: str | None = None,
    ) -> UsageSummary:
        """Query usage summary.

        Raises:
            UsageStoreUnavailableError: When the store is degraded and
                cannot answer quota-critical queries (fail-closed).
        """
        ...


# ---------------------------------------------------------------------------
# InMemoryUsageStore
# ---------------------------------------------------------------------------


class InMemoryUsageStore:
    """In-memory usage store (drop-in replacement for old UsageTracker)."""

    MAX_RECORDS = 10000

    def __init__(self):
        self._records: list[UsageRecord] = []

    def record(
        self,
        agent_name: str,
        model: str,
        usage: TokenUsage,
        cost: float,
        latency_ms: float,
        user_id: str | None = None,
        department_id: str | None = None,
    ) -> None:
        rec = UsageRecord(
            agent_name=agent_name,
            model=model,
            prompt_tokens=usage.prompt_tokens,
            completion_tokens=usage.completion_tokens,
            total_tokens=usage.total_tokens,
            cost=cost,
            latency_ms=latency_ms,
            user_id=user_id,
            department_id=department_id,
        )
        self._records.append(rec)
        if len(self._records) > self.MAX_RECORDS:
            self._records = self._records[-self.MAX_RECORDS :]

    async def record_async(
        self,
        agent_name: str,
        model: str,
        usage: TokenUsage,
        cost: float,
        latency_ms: float,
        user_id: str | None = None,
        department_id: str | None = None,
    ) -> None:
        """Async wrapper — InMemory store has no I/O, so just delegates to sync."""
        self.record(
            agent_name,
            model,
            usage,
            cost,
            latency_ms,
            user_id=user_id,
            department_id=department_id,
        )

    def get_usage(
        self,
        agent_name: str | None = None,
        start_time: datetime | None = None,
        end_time: datetime | None = None,
        user_id: str | None = None,
        department_id: str | None = None,
    ) -> UsageSummary:
        filtered = self._records

        if agent_name is not None:
            filtered = [r for r in filtered if r.agent_name == agent_name]
        if start_time is not None:
            filtered = [r for r in filtered if datetime.fromisoformat(r.timestamp) >= start_time]
        if end_time is not None:
            filtered = [r for r in filtered if datetime.fromisoformat(r.timestamp) <= end_time]
        if user_id is not None:
            filtered = [r for r in filtered if r.user_id == user_id]
        if department_id is not None:
            filtered = [r for r in filtered if r.department_id == department_id]

        if not filtered:
            return UsageSummary()

        return self._aggregate(filtered)

    def get_usage_by_user(
        self,
        user_id: str,
        start_time: datetime | None = None,
        end_time: datetime | None = None,
    ) -> UsageSummary:
        """Aggregate usage for a specific user."""
        return self.get_usage(user_id=user_id, start_time=start_time, end_time=end_time)

    def get_usage_by_department(
        self,
        department_id: str,
        start_time: datetime | None = None,
        end_time: datetime | None = None,
    ) -> UsageSummary:
        """Aggregate usage for a specific department."""
        return self.get_usage(department_id=department_id, start_time=start_time, end_time=end_time)

    @staticmethod
    def _aggregate(records: list[UsageRecord]) -> UsageSummary:
        """Build a :class:`UsageSummary` from a list of records."""
        total_tokens = sum(r.total_tokens for r in records)
        total_cost = sum(r.cost for r in records)

        by_model: dict[str, dict[str, int | float]] = {}
        by_user: dict[str, dict[str, int | float]] = {}
        by_department: dict[str, dict[str, int | float]] = {}

        def _bump(bucket_map: dict[str, dict[str, int | float]], key: str, r: UsageRecord) -> None:
            if key not in bucket_map:
                bucket_map[key] = {"total_tokens": 0, "total_cost": 0.0, "count": 0}
            bucket_map[key]["total_tokens"] += r.total_tokens
            bucket_map[key]["total_cost"] += r.cost
            bucket_map[key]["count"] += 1

        for r in records:
            _bump(by_model, r.model, r)
            if r.user_id is not None:
                _bump(by_user, r.user_id, r)
            if r.department_id is not None:
                _bump(by_department, r.department_id, r)

        return UsageSummary(
            total_tokens=total_tokens,
            total_cost=total_cost,
            by_model=by_model,
            by_user=by_user,
            by_department=by_department,
            records=records,
        )


# ---------------------------------------------------------------------------
# RedisUsageStore
# ---------------------------------------------------------------------------


class RedisUsageStore:
    """Redis-backed usage store using Hash per date for O(1) writes.

    Key schema (v2 — includes user_id/department_id in key):
        agentkit:usage:v2:{YYYY-MM-DD}:{user_id or 'none'}:{department_id or 'none'}
            → Hash: {agent:model → JSON(UsageBucket)}
        agentkit:usage_records:v2:{YYYY-MM-DD}:{user_id or 'none'}:{department_id or 'none'}
            → List: JSON(UsageRecord) with LTRIM

    Legacy v1 keys (still readable for backward compat):
        agentkit:usage:{YYYY-MM-DD}             → Hash
        agentkit:usage_records:{YYYY-MM-DD}      → List

    Fail-closed semantics:
        When Redis is unreachable and the store degrades to the in-memory
        fallback, ``get_usage`` raises :class:`UsageStoreUnavailableError`.
        The gateway translates this to HTTP 503, refusing the request
        rather than allowing untracked usage (KTD-1). The fallback store
        is used only for ``record`` (best-effort persistence) — never for
        quota queries.

    Degradation recovery:
        A background health-check task pings Redis every 30 seconds. On
        success, the degraded flag is cleared and the fallback store is
        discarded (KTD-5).
    """

    USAGE_PREFIX = "agentkit:usage:"
    RECORDS_PREFIX = "agentkit:usage_records:"
    USAGE_PREFIX_V2 = "agentkit:usage:v2:"
    RECORDS_PREFIX_V2 = "agentkit:usage_records:v2:"
    MAX_RECORDS_PER_DAY = 50000
    TTL_DAYS = 90  # Auto-expire after 90 days
    HEALTH_CHECK_INTERVAL = 30.0  # seconds

    def __init__(self, redis_url: str = "redis://localhost:6379"):
        self._redis_url = redis_url
        self._redis: object | None = None
        self._sync_redis: object | None = None
        self._fallback: InMemoryUsageStore | None = None
        self._degraded = False
        self._health_check_task: asyncio.Task[None] | None = None

    async def _get_redis(self):
        if self._redis is None:
            import redis.asyncio as aioredis

            self._redis = aioredis.from_url(self._redis_url, decode_responses=True)
        return self._redis

    def _get_sync_redis(self):
        """Get or create a persistent sync Redis client (connection pool backed)."""
        if self._sync_redis is None:
            import redis as sync_redis

            self._sync_redis = sync_redis.from_url(self._redis_url, decode_responses=True)
        return self._sync_redis

    async def aclose(self) -> None:
        # Cancel the health-check task first to avoid races.
        if self._health_check_task is not None:
            self._health_check_task.cancel()
            try:
                await self._health_check_task
            except (asyncio.CancelledError, Exception):
                pass
            self._health_check_task = None
        if self._redis is not None:
            await self._redis.aclose()
            self._redis = None
        if self._sync_redis is not None:
            self._sync_redis.close()
            self._sync_redis = None

    def _degrade_to_fallback(self) -> None:
        if not self._degraded:
            self._degraded = True
            if self._fallback is None:
                self._fallback = InMemoryUsageStore()
            logger.warning("Redis usage store unreachable, degraded to in-memory")
            self._start_health_check()

    def _start_health_check(self) -> None:
        """Start the background health-check task (idempotent).

        The task pings Redis every ``HEALTH_CHECK_INTERVAL`` seconds. On
        success, the degraded flag is cleared and the fallback store is
        discarded (KTD-5).
        """
        if self._health_check_task is not None and not self._health_check_task.done():
            return
        try:
            loop = asyncio.get_running_loop()
        except RuntimeError:
            # No running event loop (e.g. sync-only context) — skip.
            logger.debug("No event loop available for health-check task")
            return
        self._health_check_task = loop.create_task(self._health_check_loop())

    async def _health_check_loop(self) -> None:
        """Periodically ping Redis; clear degraded state on success."""
        while self._degraded:
            try:
                await asyncio.sleep(self.HEALTH_CHECK_INTERVAL)
            except asyncio.CancelledError:
                return
            try:
                redis_client = await self._get_redis()
                await redis_client.ping()
                # Redis is back — clear degraded state.
                self._degraded = False
                self._fallback = None
                logger.info("Redis usage store recovered, cleared degraded state")
                return
            except asyncio.CancelledError:
                return
            except Exception:
                logger.debug("Redis health check still failing")
                # Keep looping until Redis recovers.

    def _today_key(self) -> str:
        return datetime.now(timezone.utc).strftime("%Y-%m-%d")

    @staticmethod
    def _scope_key(part: str | None) -> str:
        """Normalize a user_id/department_id for use in a Redis key."""
        return part if part else "none"

    def _v2_keys(
        self, date_key: str, user_id: str | None, department_id: str | None
    ) -> tuple[str, str]:
        """Return (hash_key, list_key) for v2 schema."""
        u = self._scope_key(user_id)
        d = self._scope_key(department_id)
        return (
            f"{self.USAGE_PREFIX_V2}{date_key}:{u}:{d}",
            f"{self.RECORDS_PREFIX_V2}{date_key}:{u}:{d}",
        )

    def _build_record_json(
        self,
        agent_name: str,
        model: str,
        usage: TokenUsage,
        cost: float,
        latency_ms: float,
        user_id: str | None,
        department_id: str | None,
    ) -> tuple[str, str, str, str]:
        """Build (hash_key, list_key, bucket_field, record_json) for a write."""
        date_key = self._today_key()
        hash_key, list_key = self._v2_keys(date_key, user_id, department_id)
        bucket_field = f"{agent_name}:{model}"
        rec = UsageRecord(
            agent_name=agent_name,
            model=model,
            prompt_tokens=usage.prompt_tokens,
            completion_tokens=usage.completion_tokens,
            total_tokens=usage.total_tokens,
            cost=cost,
            latency_ms=latency_ms,
            user_id=user_id,
            department_id=department_id,
        )
        record_json = json.dumps(
            {
                "agent_name": rec.agent_name,
                "model": rec.model,
                "prompt_tokens": rec.prompt_tokens,
                "completion_tokens": rec.completion_tokens,
                "total_tokens": rec.total_tokens,
                "cost": rec.cost,
                "latency_ms": rec.latency_ms,
                "timestamp": rec.timestamp,
                "user_id": rec.user_id,
                "department_id": rec.department_id,
            }
        )
        return hash_key, list_key, bucket_field, record_json

    def record(
        self,
        agent_name: str,
        model: str,
        usage: TokenUsage,
        cost: float,
        latency_ms: float,
        user_id: str | None = None,
        department_id: str | None = None,
    ) -> None:
        """Record usage — sync wrapper using sync Redis client.

        Note: This is a sync method because UsageTracker.record() is sync.
        For Redis, we use a sync Redis client for writes to avoid
        needing an event loop in the caller. In async contexts prefer
        :meth:`record_async` to avoid blocking the event loop (KTD-6).
        """
        if self._degraded and self._fallback is not None:
            self._fallback.record(
                agent_name,
                model,
                usage,
                cost,
                latency_ms,
                user_id=user_id,
                department_id=department_id,
            )
            return

        try:
            r = self._get_sync_redis()
            hash_key, list_key, bucket_field, record_json = self._build_record_json(
                agent_name, model, usage, cost, latency_ms, user_id, department_id
            )

            # Atomic HINCRBYFLOAT for bucket aggregation
            pipe = r.pipeline()
            pipe.hincrbyfloat(hash_key, f"{bucket_field}:cost", cost)
            pipe.hincrby(hash_key, f"{bucket_field}:prompt_tokens", usage.prompt_tokens)
            pipe.hincrby(hash_key, f"{bucket_field}:completion_tokens", usage.completion_tokens)
            pipe.hincrby(hash_key, f"{bucket_field}:total_tokens", usage.total_tokens)
            pipe.hincrby(hash_key, f"{bucket_field}:count", 1)

            # Append record
            pipe.rpush(list_key, record_json)
            pipe.ltrim(list_key, -self.MAX_RECORDS_PER_DAY, -1)

            # Set TTL on first write of the day
            pipe.expire(hash_key, self.TTL_DAYS * 86400)
            pipe.expire(list_key, self.TTL_DAYS * 86400)

            pipe.execute()
        except Exception as e:
            logger.warning(f"Redis usage record failed: {e}")
            self._degrade_to_fallback()
            if self._fallback is not None:
                self._fallback.record(
                    agent_name,
                    model,
                    usage,
                    cost,
                    latency_ms,
                    user_id=user_id,
                    department_id=department_id,
                )

    async def record_async(
        self,
        agent_name: str,
        model: str,
        usage: TokenUsage,
        cost: float,
        latency_ms: float,
        user_id: str | None = None,
        department_id: str | None = None,
    ) -> None:
        """Async record using ``redis.asyncio`` — does not block the event loop (KTD-6).

        Preferred over :meth:`record` in async contexts (gateway, routes).
        Falls back to the in-memory store on Redis failure, same as the
        sync version.
        """
        if self._degraded and self._fallback is not None:
            await self._fallback.record_async(
                agent_name,
                model,
                usage,
                cost,
                latency_ms,
                user_id=user_id,
                department_id=department_id,
            )
            return

        try:
            r = await self._get_redis()
            hash_key, list_key, bucket_field, record_json = self._build_record_json(
                agent_name, model, usage, cost, latency_ms, user_id, department_id
            )

            pipe = r.pipeline()
            pipe.hincrbyfloat(hash_key, f"{bucket_field}:cost", cost)
            pipe.hincrby(hash_key, f"{bucket_field}:prompt_tokens", usage.prompt_tokens)
            pipe.hincrby(hash_key, f"{bucket_field}:completion_tokens", usage.completion_tokens)
            pipe.hincrby(hash_key, f"{bucket_field}:total_tokens", usage.total_tokens)
            pipe.hincrby(hash_key, f"{bucket_field}:count", 1)
            pipe.rpush(list_key, record_json)
            pipe.ltrim(list_key, -self.MAX_RECORDS_PER_DAY, -1)
            pipe.expire(hash_key, self.TTL_DAYS * 86400)
            pipe.expire(list_key, self.TTL_DAYS * 86400)
            await pipe.execute()
        except Exception as e:
            logger.warning(f"Redis async usage record failed: {e}")
            self._degrade_to_fallback()
            if self._fallback is not None:
                await self._fallback.record_async(
                    agent_name,
                    model,
                    usage,
                    cost,
                    latency_ms,
                    user_id=user_id,
                    department_id=department_id,
                )

    def get_usage(
        self,
        agent_name: str | None = None,
        start_time: datetime | None = None,
        end_time: datetime | None = None,
        user_id: str | None = None,
        department_id: str | None = None,
    ) -> UsageSummary:
        """Query usage summary from Redis.

        Scans v2 keys (filtered by user_id/department_id when provided)
        and legacy v1 keys (no per-user/department scoping). Records
        from both schemas are merged.

        Key construction (U1 fix):
            - Both user_id and department_id provided → direct key lookup
              (exact match).
            - Only user_id provided → SCAN ``...:{user_id}:*`` to aggregate
              across all departments the user belongs to.
            - Only department_id provided → SCAN ``...:*:{department_id}``
              to aggregate across all users in that department.
            - Neither provided → SCAN ``...:*`` (all records for the date).

        Fail-closed (KTD-1):
            When the store is degraded (Redis unreachable), raises
            :class:`UsageStoreUnavailableError`. The fallback store is
            used only for ``record`` — never for quota queries.
        """
        # Fail-closed: when degraded, refuse quota-critical queries.
        # The in-memory fallback is only for recording (best-effort),
        # not for quota checks — returning an empty summary would
        # make quota checks pass (fail-open), which is a security bug.
        if self._degraded:
            raise UsageStoreUnavailableError(
                "Redis usage store is degraded — cannot answer quota-critical query"
            )

        try:
            r = self._get_sync_redis()

            # Determine date range to scan
            start = start_time or datetime(2020, 1, 1, tzinfo=timezone.utc)
            end = end_time or datetime.now(timezone.utc)

            all_records: list[UsageRecord] = []

            # Scan v2 keys.
            current = start.date()
            end_date = end.date()
            while current <= end_date:
                date_key = current.isoformat()
                # Determine the SCAN pattern based on which scope dims are provided.
                # Key format: {RECORDS_PREFIX_V2}{date}:{user_or_none}:{dept_or_none}
                if user_id is not None and department_id is not None:
                    # Both provided → direct key lookup (exact match).
                    list_key = (
                        f"{self.RECORDS_PREFIX_V2}{date_key}:"
                        f"{self._scope_key(user_id)}:{self._scope_key(department_id)}"
                    )
                    all_records.extend(self._read_list(r, list_key, start, end, agent_name))
                elif user_id is not None:
                    # Only user_id → aggregate across all departments.
                    pattern = f"{self.RECORDS_PREFIX_V2}{date_key}:{self._scope_key(user_id)}:*"
                    for key in r.scan_iter(match=pattern, count=200):
                        all_records.extend(self._read_list(r, key, start, end, agent_name))
                elif department_id is not None:
                    # Only department_id → aggregate across all users.
                    pattern = (
                        f"{self.RECORDS_PREFIX_V2}{date_key}:*:{self._scope_key(department_id)}"
                    )
                    for key in r.scan_iter(match=pattern, count=200):
                        all_records.extend(self._read_list(r, key, start, end, agent_name))
                else:
                    # Neither provided → scan all v2 list keys for this date.
                    pattern = f"{self.RECORDS_PREFIX_V2}{date_key}:*"
                    for key in r.scan_iter(match=pattern, count=200):
                        all_records.extend(self._read_list(r, key, start, end, agent_name))
                current = current + timedelta(days=1)

            # Also scan legacy v1 keys (no user/department scoping).
            current = start.date()
            while current <= end_date:
                list_key = f"{self.RECORDS_PREFIX}{current.isoformat()}"
                all_records.extend(self._read_list(r, list_key, start, end, agent_name))
                current = current + timedelta(days=1)

            # Apply user_id/department_id filters to records from legacy
            # v1 keys (which don't carry these fields — they'll be None).
            if user_id is not None:
                all_records = [r for r in all_records if r.user_id == user_id]
            if department_id is not None:
                all_records = [r for r in all_records if r.department_id == department_id]

            if not all_records:
                return UsageSummary()

            return InMemoryUsageStore._aggregate(all_records)
        except UsageStoreUnavailableError:
            raise
        except Exception as e:
            logger.warning(f"Redis usage query failed: {e}")
            # Degrade for future calls, but fail-closed for this one.
            self._degrade_to_fallback()
            raise UsageStoreUnavailableError(f"Redis usage query failed: {e}") from e

    def get_usage_by_user(
        self,
        user_id: str,
        start_time: datetime | None = None,
        end_time: datetime | None = None,
    ) -> UsageSummary:
        """Aggregate usage for a specific user."""
        return self.get_usage(user_id=user_id, start_time=start_time, end_time=end_time)

    def get_usage_by_department(
        self,
        department_id: str,
        start_time: datetime | None = None,
        end_time: datetime | None = None,
    ) -> UsageSummary:
        """Aggregate usage for a specific department."""
        return self.get_usage(department_id=department_id, start_time=start_time, end_time=end_time)

    @staticmethod
    def _read_list(
        r: object,
        list_key: str,
        start: datetime,
        end: datetime,
        agent_name: str | None,
    ) -> list[UsageRecord]:
        """Read all records from a Redis list, filtered by time range and agent."""
        out: list[UsageRecord] = []
        raw_records = r.lrange(list_key, 0, -1)
        for raw in raw_records:
            try:
                data = json.loads(raw)
            except json.JSONDecodeError:
                continue
            # Build record, tolerating legacy records without user_id/department_id.
            rec = UsageRecord(
                agent_name=data["agent_name"],
                model=data["model"],
                prompt_tokens=data["prompt_tokens"],
                completion_tokens=data["completion_tokens"],
                total_tokens=data["total_tokens"],
                cost=data["cost"],
                latency_ms=data["latency_ms"],
                timestamp=data.get("timestamp", ""),
                user_id=data.get("user_id"),
                department_id=data.get("department_id"),
            )
            if not rec.timestamp:
                continue
            try:
                rec_ts = datetime.fromisoformat(rec.timestamp)
            except ValueError:
                continue
            if rec_ts < start or rec_ts > end:
                continue
            if agent_name is not None and rec.agent_name != agent_name:
                continue
            out.append(rec)
        return out


# ---------------------------------------------------------------------------
# Factory
# ---------------------------------------------------------------------------


def create_usage_store(
    backend: str = "auto",
    redis_url: str = "redis://localhost:6379",
) -> UsageStore:
    """Create a usage store backend.

    Args:
        backend: "auto" (try Redis, fallback to memory), "redis", "memory".
        redis_url: Redis connection URL.

    Returns:
        A UsageStore instance.
    """
    if backend in ("auto", "redis"):
        try:
            import redis  # noqa: F401

            return RedisUsageStore(redis_url=redis_url)
        except ImportError:
            logger.warning("redis package not available, falling back to in-memory usage store")
            return InMemoryUsageStore()
    return InMemoryUsageStore()
