from __future__ import annotations

import asyncio
import json
import time
from ast import literal_eval
from typing import Any, Dict, Optional
from dataclasses import dataclass, field
from enum import Enum
import re

from .logging import get_logger

logger = get_logger(__name__)


class RateLimitType(str, Enum):
    """Types of rate limits."""
    PER_USER = "per_user"
    PER_APP = "per_app"
    PER_ENDPOINT = "per_endpoint"


@dataclass
class RateLimit:
    """Rate limit configuration."""
    requests_per_hour: int
    requests_per_minute: int = 0
    burst_limit: int = 0
    window_size: int = 3600  # seconds
    
    def __post_init__(self):
        if self.requests_per_minute == 0:
            self.requests_per_minute = self.requests_per_hour // 60
        if self.burst_limit == 0:
            self.burst_limit = self.requests_per_minute * 2


@dataclass
class TokenBucket:
    """Token bucket implementation for rate limiting."""
    capacity: int
    tokens: float = 0
    last_refill: float = field(default_factory=time.time)
    refill_rate: float = 0  # tokens per second
    
    def __post_init__(self):
        if self.tokens == 0:
            self.tokens = self.capacity
        if self.refill_rate == 0:
            self.refill_rate = self.capacity / 60  # refill over 1 minute


class RateLimitInfo:
    """Information about current rate limit status."""
    
    def __init__(self, limit: int, remaining: int, reset_time: int):
        self.limit = limit
        self.remaining = remaining
        self.reset_time = reset_time
    
    @property
    def is_exceeded(self) -> bool:
        return self.remaining <= 0
    
    @property
    def reset_in_seconds(self) -> int:
        return max(0, self.reset_time - int(time.time()))


class RateLimiter:
    """Advanced rate limiter with token bucket algorithm and Meta API integration."""
    
    def __init__(self):
        self._buckets: Dict[str, TokenBucket] = {}
        self._endpoint_limits: Dict[str, RateLimit] = {}
        self._lock = asyncio.Lock()
        
        # Meta API rate limits (per app)
        self._meta_limits = {
            "insights": RateLimit(requests_per_hour=200, requests_per_minute=5),
            "campaigns": RateLimit(requests_per_hour=1000, requests_per_minute=20),
            "adsets": RateLimit(requests_per_hour=1000, requests_per_minute=20),
            "ads": RateLimit(requests_per_hour=1000, requests_per_minute=20),
            "audiences": RateLimit(requests_per_hour=500, requests_per_minute=10),
            "pixels": RateLimit(requests_per_hour=200, requests_per_minute=5),
            "default": RateLimit(requests_per_hour=200, requests_per_minute=5),
        }
        # Baseline limits used to recover from temporary throttling adjustments.
        self._baseline_limits = {
            key: RateLimit(
                requests_per_hour=value.requests_per_hour,
                requests_per_minute=value.requests_per_minute,
                burst_limit=value.burst_limit,
                window_size=value.window_size,
            )
            for key, value in self._meta_limits.items()
        }
    
    def _get_endpoint_type(self, path: str) -> str:
        """Extract endpoint type from API path."""
        # Remove version and account ID
        clean_path = re.sub(r'^v\d+\.\d+/', '', path)
        clean_path = re.sub(r'^act_\d+/', '', clean_path)
        
        if 'insights' in clean_path:
            return 'insights'
        elif 'campaigns' in clean_path:
            return 'campaigns'
        elif 'adsets' in clean_path:
            return 'adsets'
        elif 'ads' in clean_path and 'adsets' not in clean_path:
            return 'ads'
        elif 'audience' in clean_path or 'customaudience' in clean_path:
            return 'audiences'
        elif 'pixel' in clean_path:
            return 'pixels'
        else:
            return 'default'
    
    def _get_bucket_key(self, endpoint_type: str, user_id: Optional[str] = None) -> str:
        """Generate bucket key for rate limiting."""
        if user_id:
            return f"user:{user_id}:{endpoint_type}"
        return f"app:{endpoint_type}"
    
    def _get_or_create_bucket(self, key: str, limit: RateLimit) -> TokenBucket:
        """Get or create a token bucket for the given key."""
        target_capacity = max(1, int(limit.burst_limit))
        target_refill_rate = max(1 / 60.0, float(limit.requests_per_minute) / 60.0)

        if key not in self._buckets:
            self._buckets[key] = TokenBucket(
                capacity=target_capacity,
                refill_rate=target_refill_rate,
            )
            return self._buckets[key]

        bucket = self._buckets[key]
        if bucket.capacity != target_capacity:
            bucket.capacity = target_capacity
            bucket.tokens = min(bucket.tokens, target_capacity)
        bucket.refill_rate = target_refill_rate
        return bucket
    
    def _refill_bucket(self, bucket: TokenBucket) -> None:
        """Refill tokens in the bucket based on elapsed time."""
        now = time.time()
        elapsed = now - bucket.last_refill
        
        # Add tokens based on elapsed time
        tokens_to_add = elapsed * bucket.refill_rate
        bucket.tokens = min(bucket.capacity, bucket.tokens + tokens_to_add)
        bucket.last_refill = now
    
    def _seconds_until_next_token(self, bucket: TokenBucket) -> float:
        """Compute the wait required until at least one token is available."""
        if bucket.tokens >= 1:
            return 0.0
        if bucket.refill_rate <= 0:
            return 60.0
        return max(0.01, (1 - bucket.tokens) / bucket.refill_rate)
    
    async def acquire(
        self,
        path: str,
        user_id: Optional[str] = None,
        timeout: float = 30.0
    ) -> RateLimitInfo:
        """Acquire permission to make a request."""
        endpoint_type = self._get_endpoint_type(path)
        bucket_key = self._get_bucket_key(endpoint_type, user_id)

        deadline = time.monotonic() + timeout
        while True:
            async with self._lock:
                limit = self._endpoint_limits.get(endpoint_type, self._meta_limits[endpoint_type])
                bucket = self._get_or_create_bucket(bucket_key, limit)
                self._refill_bucket(bucket)

                if bucket.tokens >= 1:
                    bucket.tokens -= 1
                    remaining = int(bucket.tokens)
                    reset_time = int(time.time() + 60)  # Approximate reset window

                    logger.debug(
                        "Rate limit token acquired",
                        endpoint=endpoint_type,
                        remaining=remaining,
                        bucket=bucket_key
                    )

                    return RateLimitInfo(limit.burst_limit, remaining, reset_time)

                wait_seconds = self._seconds_until_next_token(bucket)

            remaining_timeout = deadline - time.monotonic()
            if remaining_timeout <= 0:
                logger.warning(
                    "Rate limit wait timeout",
                    endpoint=endpoint_type,
                    bucket=bucket_key,
                    timeout=timeout,
                )
                raise RateLimitExceededError(
                    f"Rate limit exceeded for {endpoint_type}",
                    reset_time=int(time.time() + 60)
                )

            sleep_for = min(wait_seconds, remaining_timeout)
            logger.debug(
                "Rate limit exceeded, sleeping for refill",
                endpoint=endpoint_type,
                bucket=bucket_key,
                wait_seconds=round(sleep_for, 3),
            )
            await asyncio.sleep(sleep_for)
    
    async def release_token(self, path: str, user_id: Optional[str] = None) -> None:
        """Release a token back to the bucket (for failed requests)."""
        endpoint_type = self._get_endpoint_type(path)
        bucket_key = self._get_bucket_key(endpoint_type, user_id)
        
        async with self._lock:
            if bucket_key in self._buckets:
                bucket = self._buckets[bucket_key]
                bucket.tokens = min(bucket.capacity, bucket.tokens + 1)
    
    def update_limits_from_headers(self, headers: Dict[str, str]) -> None:
        """Update rate limits based on Meta API response headers."""
        # Meta API rate limit headers
        if 'x-app-usage' in headers:
            try:
                raw_usage = headers['x-app-usage']
                try:
                    usage = json.loads(raw_usage)
                except json.JSONDecodeError:
                    usage = literal_eval(raw_usage)

                if isinstance(usage, dict):
                    usage_values: list[float] = []
                    for value in usage.values():
                        try:
                            usage_values.append(float(value))
                        except (TypeError, ValueError):
                            continue

                    if not usage_values:
                        return

                    # x-app-usage values are percentages (0-100)
                    max_usage = max(usage_values)
                    for endpoint, limit_config in self._meta_limits.items():
                        baseline = self._baseline_limits[endpoint]
                        if max_usage >= 80:
                            new_rpm = max(1, int(baseline.requests_per_minute * 0.6))
                        elif max_usage <= 30:
                            new_rpm = baseline.requests_per_minute
                        else:
                            new_rpm = max(1, int(baseline.requests_per_minute * 0.8))

                        limit_config.requests_per_minute = new_rpm
                        limit_config.burst_limit = max(2, new_rpm * 2)
            except Exception as e:
                logger.warning("Failed to parse rate limit headers: %s", str(e))
    
    def get_status(self, endpoint_type: str, user_id: Optional[str] = None) -> Dict[str, Any]:
        """Get current rate limit status."""
        bucket_key = self._get_bucket_key(endpoint_type, user_id)
        
        if bucket_key in self._buckets:
            bucket = self._buckets[bucket_key]
            limit = self._endpoint_limits.get(endpoint_type, self._meta_limits[endpoint_type])
            
            return {
                "endpoint": endpoint_type,
                "tokens_available": int(bucket.tokens),
                "capacity": bucket.capacity,
                "refill_rate": bucket.refill_rate,
                "requests_per_hour": limit.requests_per_hour,
                "requests_per_minute": limit.requests_per_minute,
                "last_refill": bucket.last_refill,
            }
        
        return {"endpoint": endpoint_type, "status": "no_bucket"}


class RateLimitExceededError(Exception):
    """Raised when rate limit is exceeded."""
    
    def __init__(self, message: str, reset_time: int):
        super().__init__(message)
        self.reset_time = reset_time
        self.retry_after = max(0, reset_time - int(time.time()))


# Global rate limiter instance
_rate_limiter: Optional[RateLimiter] = None


def get_rate_limiter() -> RateLimiter:
    """Get the global rate limiter instance."""
    global _rate_limiter
    if _rate_limiter is None:
        _rate_limiter = RateLimiter()
    return _rate_limiter


async def with_rate_limit(
    path: str,
    user_id: Optional[str] = None,
    timeout: float = 30.0
):
    """Context manager for rate limiting."""
    limiter = get_rate_limiter()
    rate_info = await limiter.acquire(path, user_id, timeout)
    return rate_info


class RateLimitMiddleware:
    """Middleware for automatic rate limiting."""
    
    def __init__(self, limiter: Optional[RateLimiter] = None):
        self.limiter = limiter or get_rate_limiter()
    
    async def __call__(self, request_func, path: str, user_id: Optional[str] = None, **kwargs):
        """Execute request with rate limiting."""
        await self.limiter.acquire(path, user_id)
        
        try:
            # Execute the request
            response = await request_func(path=path, **kwargs)
            
            # Update limits from response headers if available
            if hasattr(response, 'headers'):
                self.limiter.update_limits_from_headers(response.headers)
            
            return response
            
        except Exception as e:
            # Release token for failed requests
            await self.limiter.release_token(path, user_id)
            raise e
