"""
Advanced Rate Limiting with Tiered Quotas

Implements sliding window rate limiting with:
- Per-minute, per-hour, per-day limits
- Concurrent session tracking
- Priority queuing for higher tiers
- Burst allowance
- Usage metering for billing
"""
from __future__ import annotations

import time
from dataclasses import dataclass
from typing import Optional

from app.auth.api_keys import ApiKeyInfo, PlanLimits
from app.observability.logging import get_logger


class RateLimitError(Exception):
    """Raised when rate limit is exceeded."""
    def __init__(
        self,
        message: str,
        limit_type: str,
        limit: int,
        current: int,
        retry_after: int,
    ):
        super().__init__(message)
        self.limit_type = limit_type
        self.limit = limit
        self.current = current
        self.retry_after = retry_after


class ConcurrencyLimitError(Exception):
    """Raised when concurrent session limit is exceeded."""
    def __init__(self, message: str, limit: int, current: int):
        super().__init__(message)
        self.limit = limit
        self.current = current


@dataclass
class UsageStats:
    """Current usage statistics for an organization."""
    requests_minute: int
    requests_hour: int
    requests_day: int
    concurrent_sessions: int
    total_voice_seconds: float
    total_requests: int


_logger = get_logger("rate_limit")


async def check_rate_limit(
    redis,
    org_id: str,
    limits: PlanLimits,
    endpoint: str = "default",
) -> UsageStats:
    """
    Check if request is within rate limits using sliding window.
    
    Raises RateLimitError if any limit is exceeded.
    Returns current usage stats if within limits.
    """
    now = int(time.time())
    minute_key = f"ratelimit:{org_id}:minute:{now // 60}"
    hour_key = f"ratelimit:{org_id}:hour:{now // 3600}"
    day_key = f"ratelimit:{org_id}:day:{now // 86400}"
    
    pipe = redis.pipeline()
    
    # Get current counts
    pipe.get(minute_key)
    pipe.get(hour_key)
    pipe.get(day_key)
    pipe.get(f"sessions:{org_id}:count")
    
    results = await pipe.execute()
    
    minute_count = int(results[0] or 0)
    hour_count = int(results[1] or 0)
    day_count = int(results[2] or 0)
    session_count = int(results[3] or 0)
    
    # Check limits
    if minute_count >= limits.requests_per_minute:
        retry_after = 60 - (now % 60)
        _logger.warning(
            "rate_limit_minute",
            org_id=org_id,
            count=minute_count,
            limit=limits.requests_per_minute,
        )
        raise RateLimitError(
            f"Rate limit exceeded: {minute_count}/{limits.requests_per_minute} requests per minute",
            limit_type="minute",
            limit=limits.requests_per_minute,
            current=minute_count,
            retry_after=retry_after,
        )
    
    if hour_count >= limits.requests_per_hour:
        retry_after = 3600 - (now % 3600)
        _logger.warning(
            "rate_limit_hour",
            org_id=org_id,
            count=hour_count,
            limit=limits.requests_per_hour,
        )
        raise RateLimitError(
            f"Rate limit exceeded: {hour_count}/{limits.requests_per_hour} requests per hour",
            limit_type="hour",
            limit=limits.requests_per_hour,
            current=hour_count,
            retry_after=retry_after,
        )
    
    if day_count >= limits.requests_per_day:
        retry_after = 86400 - (now % 86400)
        _logger.warning(
            "rate_limit_day",
            org_id=org_id,
            count=day_count,
            limit=limits.requests_per_day,
        )
        raise RateLimitError(
            f"Rate limit exceeded: {day_count}/{limits.requests_per_day} requests per day",
            limit_type="day",
            limit=limits.requests_per_day,
            current=day_count,
            retry_after=retry_after,
        )
    
    # Increment counters
    pipe = redis.pipeline()
    pipe.incr(minute_key)
    pipe.expire(minute_key, 120)  # 2 min TTL
    pipe.incr(hour_key)
    pipe.expire(hour_key, 7200)  # 2 hour TTL
    pipe.incr(day_key)
    pipe.expire(day_key, 172800)  # 2 day TTL
    await pipe.execute()
    
    return UsageStats(
        requests_minute=minute_count + 1,
        requests_hour=hour_count + 1,
        requests_day=day_count + 1,
        concurrent_sessions=session_count,
        total_voice_seconds=0,
        total_requests=day_count + 1,
    )


async def check_concurrent_sessions(
    redis,
    org_id: str,
    limits: PlanLimits,
) -> int:
    """
    Check if organization can start a new concurrent session.
    
    Raises ConcurrencyLimitError if at limit.
    Returns current session count if within limits.
    """
    count_key = f"sessions:{org_id}:count"
    current = int(await redis.get(count_key) or 0)
    
    if current >= limits.concurrent_sessions:
        _logger.warning(
            "concurrent_limit",
            org_id=org_id,
            count=current,
            limit=limits.concurrent_sessions,
        )
        raise ConcurrencyLimitError(
            f"Concurrent session limit reached: {current}/{limits.concurrent_sessions}",
            limit=limits.concurrent_sessions,
            current=current,
        )
    
    return current


async def reserve_session(
    redis,
    org_id: str,
    session_id: str,
    limits: PlanLimits,
    ttl_seconds: Optional[int] = None,
) -> None:
    """
    Reserve a session slot for an organization.
    
    Atomically increments session count and tracks individual session.
    """
    if ttl_seconds is None:
        ttl_seconds = limits.max_session_duration_seconds
    
    count_key = f"sessions:{org_id}:count"
    session_key = f"sessions:{org_id}:{session_id}"
    
    # Check limit first
    await check_concurrent_sessions(redis, org_id, limits)
    
    pipe = redis.pipeline()
    pipe.incr(count_key)
    pipe.expire(count_key, ttl_seconds + 60)
    pipe.set(session_key, str(int(time.time())), ex=ttl_seconds)
    await pipe.execute()
    
    _logger.info(
        "session_reserved",
        org_id=org_id,
        session_id=session_id,
        ttl=ttl_seconds,
    )


async def release_session(
    redis,
    org_id: str,
    session_id: str,
) -> None:
    """Release a session slot."""
    count_key = f"sessions:{org_id}:count"
    session_key = f"sessions:{org_id}:{session_id}"
    
    # Only decrement if session exists
    existed = await redis.delete(session_key)
    if existed:
        await redis.decr(count_key)
        _logger.info(
            "session_released",
            org_id=org_id,
            session_id=session_id,
        )


async def extend_session(
    redis,
    org_id: str,
    session_id: str,
    extension_seconds: int,
    limits: PlanLimits,
) -> bool:
    """
    Extend a session's TTL.
    
    Returns True if extended, False if session doesn't exist or max duration reached.
    """
    session_key = f"sessions:{org_id}:{session_id}"
    
    # Get session start time
    start_time = await redis.get(session_key)
    if not start_time:
        return False
    
    elapsed = int(time.time()) - int(start_time)
    if elapsed + extension_seconds > limits.max_session_duration_seconds:
        _logger.warning(
            "session_max_duration",
            org_id=org_id,
            session_id=session_id,
            elapsed=elapsed,
            max=limits.max_session_duration_seconds,
        )
        return False
    
    await redis.expire(session_key, extension_seconds)
    return True


async def record_voice_usage(
    redis,
    org_id: str,
    session_id: str,
    duration_seconds: float,
    usage_type: str = "voice",
) -> None:
    """
    Record voice usage for billing.
    
    Tracks usage by day for billing aggregation.
    """
    now = int(time.time())
    day_key = f"usage:{org_id}:{usage_type}:{now // 86400}"
    total_key = f"usage:{org_id}:{usage_type}:total"
    
    pipe = redis.pipeline()
    pipe.incrbyfloat(day_key, duration_seconds)
    pipe.expire(day_key, 86400 * 90)  # Keep 90 days
    pipe.incrbyfloat(total_key, duration_seconds)
    await pipe.execute()
    
    _logger.debug(
        "voice_usage_recorded",
        org_id=org_id,
        session_id=session_id,
        duration=duration_seconds,
    )


async def get_usage_summary(
    redis,
    org_id: str,
    limits: PlanLimits,
) -> UsageStats:
    """Get current usage summary for an organization."""
    now = int(time.time())
    
    pipe = redis.pipeline()
    pipe.get(f"ratelimit:{org_id}:minute:{now // 60}")
    pipe.get(f"ratelimit:{org_id}:hour:{now // 3600}")
    pipe.get(f"ratelimit:{org_id}:day:{now // 86400}")
    pipe.get(f"sessions:{org_id}:count")
    pipe.get(f"usage:{org_id}:voice:total")
    
    results = await pipe.execute()
    
    return UsageStats(
        requests_minute=int(results[0] or 0),
        requests_hour=int(results[1] or 0),
        requests_day=int(results[2] or 0),
        concurrent_sessions=int(results[3] or 0),
        total_voice_seconds=float(results[4] or 0),
        total_requests=int(results[2] or 0),
    )


async def get_remaining_quota(
    redis,
    org_id: str,
    limits: PlanLimits,
) -> dict:
    """Get remaining quota for rate limit headers."""
    stats = await get_usage_summary(redis, org_id, limits)
    
    return {
        "X-RateLimit-Limit-Minute": limits.requests_per_minute,
        "X-RateLimit-Remaining-Minute": max(0, limits.requests_per_minute - stats.requests_minute),
        "X-RateLimit-Limit-Hour": limits.requests_per_hour,
        "X-RateLimit-Remaining-Hour": max(0, limits.requests_per_hour - stats.requests_hour),
        "X-RateLimit-Limit-Day": limits.requests_per_day,
        "X-RateLimit-Remaining-Day": max(0, limits.requests_per_day - stats.requests_day),
        "X-Concurrent-Sessions": stats.concurrent_sessions,
        "X-Concurrent-Sessions-Limit": limits.concurrent_sessions,
    }
