# src/fmatch/saas/services/rate_limiter.py
"""
Production-ready rate limiting for Salesforce API.
Handles daily limits, concurrent requests, and backoff.
"""

import asyncio
from typing import Any, Callable, Optional, Dict, List
from datetime import datetime, timedelta
import logging
from dataclasses import dataclass
import random

logger = logging.getLogger(__name__)


@dataclass
class RateLimitConfig:
    """Rate limit configuration."""

    daily_limit: int = 15000  # Standard Salesforce limit
    concurrent_limit: int = 25  # Max concurrent requests
    requests_per_second: int = 10  # Throttle rate
    backoff_base: float = 1.0  # Base backoff in seconds
    backoff_max: float = 60.0  # Max backoff
    retry_count: int = 5  # Max retries


class SmartRateLimiter:
    """Intelligent rate limiter with backoff and usage tracking."""

    def __init__(self, config: Optional[RateLimitConfig] = None):
        self.config = config or RateLimitConfig()
        self.daily_usage = 0
        self.daily_reset = datetime.utcnow().replace(hour=0, minute=0, second=0)
        self.semaphore = asyncio.Semaphore(self.config.concurrent_limit)
        self.request_times = []
        self.lock = asyncio.Lock()

    async def check_daily_limit(self) -> bool:
        """Check if daily limit is exceeded."""
        async with self.lock:
            # Reset daily counter if new day
            now = datetime.utcnow()
            if now >= self.daily_reset + timedelta(days=1):
                self.daily_usage = 0
                self.daily_reset = now.replace(hour=0, minute=0, second=0)

            return self.daily_usage < self.config.daily_limit

    async def wait_for_rate_limit(self):
        """Wait if we're exceeding requests per second."""
        async with self.lock:
            now = datetime.utcnow()

            # Clean old request times
            cutoff = now - timedelta(seconds=1)
            self.request_times = [t for t in self.request_times if t > cutoff]

            # Wait if necessary
            if len(self.request_times) >= self.config.requests_per_second:
                wait_time = 1.0 - (now - self.request_times[0]).total_seconds()
                if wait_time > 0:
                    await asyncio.sleep(wait_time)

            self.request_times.append(now)

    async def execute_with_backoff(self, func: Callable, *args, **kwargs) -> Any:
        """Execute function with rate limiting and exponential backoff."""

        # Check daily limit
        if not await self.check_daily_limit():
            raise Exception(
                f"Daily API limit reached: {self.daily_usage}/{self.config.daily_limit}"
            )

        retry = 0
        backoff = self.config.backoff_base

        while retry < self.config.retry_count:
            try:
                # Wait for rate limit
                await self.wait_for_rate_limit()

                # Execute with concurrent limit
                async with self.semaphore:
                    result = await func(*args, **kwargs)

                    # Update usage counter
                    async with self.lock:
                        self.daily_usage += 1

                    return result

            except Exception as e:
                error_msg = str(e).lower()

                # Check for rate limit errors
                if any(
                    x in error_msg for x in ["rate limit", "429", "too many requests"]
                ):
                    retry += 1

                    if retry >= self.config.retry_count:
                        raise

                    # Exponential backoff with jitter
                    wait_time = min(
                        backoff * (2**retry) + random.uniform(0, 1),
                        self.config.backoff_max,
                    )

                    logger.warning(
                        f"Rate limit hit, retry {retry}/{self.config.retry_count} "
                        f"after {wait_time:.1f}s"
                    )

                    await asyncio.sleep(wait_time)

                else:
                    # Not a rate limit error, re-raise
                    raise

        raise Exception(f"Max retries ({self.config.retry_count}) exceeded")

    def get_usage_stats(self) -> Dict[str, Any]:
        """Get current usage statistics."""
        return {
            "daily_usage": self.daily_usage,
            "daily_limit": self.config.daily_limit,
            "daily_remaining": self.config.daily_limit - self.daily_usage,
            "daily_reset": self.daily_reset.isoformat(),
            "current_concurrent": self.config.concurrent_limit - self.semaphore._value,
            "requests_last_second": len(self.request_times),
        }


class AdaptiveSampler:
    """Adaptive sampling for large orgs."""

    @staticmethod
    def get_sampling_strategy(total_records: int) -> Dict[str, Any]:
        """Determine optimal sampling strategy based on org size."""

        if total_records < 5000:
            return {
                "strategy": "full",
                "sample_size": total_records,
                "confidence": 1.0,
                "message": "Analyzing all records",
            }

        elif total_records < 50000:
            return {
                "strategy": "stratified",
                "sample_size": min(10000, total_records // 2),
                "confidence": 0.95,
                "strata": ["LastModifiedDate", "CreatedDate"],
                "message": f"Sampling {min(10000, total_records // 2)} records across time periods",
            }

        elif total_records < 500000:
            return {
                "strategy": "progressive",
                "initial_sample": 5000,
                "additional_samples": 3,
                "sample_size": 20000,
                "confidence": 0.90,
                "message": "Using progressive sampling for large dataset",
            }

        else:
            # Very large org (500k+ records)
            return {
                "strategy": "intelligent",
                "initial_sample": 10000,
                "focused_samples": True,
                "sample_size": 50000,
                "confidence": 0.85,
                "focus_areas": [
                    "Recently modified (last 90 days)",
                    "High value records (by revenue/size)",
                    "Frequently accessed records",
                ],
                "message": "Using intelligent sampling for enterprise dataset",
            }

    @staticmethod
    async def execute_progressive_sampling(
        sf_client,
        object_type: str,
        strategy: Dict[str, Any],
        callback: Optional[Callable] = None,
    ) -> List[Dict]:
        """Execute progressive sampling based on strategy."""

        all_samples = []

        if strategy["strategy"] == "full":
            # Get all records
            return await sf_client.get_all_records(
                object_type, limit=strategy["sample_size"]
            )

        elif strategy["strategy"] == "stratified":
            # Sample across time periods
            for stratum in strategy.get("strata", ["LastModifiedDate"]):
                sample = await sf_client.get_sample_records(
                    object_type,
                    ["Id", "Name", "Email", "Phone", "Website"],
                    n=strategy["sample_size"] // len(strategy.get("strata", [1])),
                    order_by=stratum,
                )
                all_samples.extend(sample)

                if callback:
                    await callback(
                        {
                            "phase": "sampling",
                            "message": f"Sampled {len(all_samples)} records",
                        }
                    )

        elif strategy["strategy"] == "progressive":
            # Start with initial sample
            initial = await sf_client.get_sample_records(
                object_type,
                ["Id", "Name", "Email", "Phone", "Website"],
                n=strategy["initial_sample"],
            )
            all_samples.extend(initial)

            # Additional targeted samples
            for i in range(strategy.get("additional_samples", 3)):
                additional = await sf_client.get_sample_records(
                    object_type,
                    ["Id", "Name", "Email", "Phone", "Website"],
                    n=strategy["initial_sample"],
                    offset=len(all_samples),
                )
                all_samples.extend(additional)

                if callback:
                    await callback(
                        {
                            "phase": "sampling",
                            "message": f"Progressive sampling: {len(all_samples)} records",
                        }
                    )

        return all_samples
