"""
Quota enforcement middleware with grace band support.
Protects both your infra and user's SFDC org limits.
"""

from typing import Dict, Any, Optional
from datetime import datetime, timezone
from fastapi import HTTPException
from sqlalchemy.ext.asyncio import AsyncSession
from sqlalchemy import select, update
import logging

log = logging.getLogger(__name__)


# =============================================================================
# Tier limits configuration - NOW LOADED FROM CENTRALIZED CONFIG
# =============================================================================
# This imports from config/tier_limits.yaml via the config loader.
# If you need to change tier allocations, update the YAML file instead of
# hardcoding values here. This ensures backend, website, and docs stay in sync.

try:
    from ..config.tier_limits import get_tier_limits

    TIER_LIMITS = get_tier_limits()
    log.info(
        f"Loaded tier limits from centralized config: {list(TIER_LIMITS.keys())}"
    )
except Exception as e:
    # Fallback to minimal config if YAML is missing (prevents startup failure)
    log.error(
        f"Failed to load tier limits from config/tier_limits.yaml: {e}\n"
        f"Using fallback minimal configuration. This should only happen in dev!"
    )
    # Fallback values synced with config/tier_limits.yaml v1.2.0 (2026-01-19)
    # Changes: users=1 for all tiers, max_schedules=5/25/100, enrichments shelved
    TIER_LIMITS = {
        "free": {
            "monthly_rows": 5000,
            "rows_per_run": 5000,
            "grace_percentage": 0.10,
            "enrichment_credits": 0,  # Shelved for V1
            "max_schedules": 5,
            "allowed_frequencies": ["daily", "weekly"],
            "concurrent_jobs": 1,
            "candidate_cap": 50,
            "require_b2c": True,
            "allow_advanced_algos": False,
            "soql_runs_per_month": 15,
            "salesforce_orgs": 1,
            "users": 1,
            # AI Limits
            "gremlin_enabled": False,
            "gremlin_messages_per_day": 0,
            "ai_report_row_cap": 5000,
        },
        "pro": {
            "monthly_rows": 100000,
            "rows_per_run": 50000,
            "grace_percentage": 0.05,
            "enrichment_credits": 0,  # Shelved for V1
            "max_schedules": 25,
            "allowed_frequencies": ["hourly", "daily", "weekly"],
            "concurrent_jobs": 1,
            "candidate_cap": 2000,
            "require_b2c": False,
            "allow_advanced_algos": True,
            "soql_runs_per_month": 100,
            "salesforce_orgs": 1,
            "users": 1,  # Seats sold as add-ons at $34/mo
            # AI Limits
            "gremlin_enabled": False,              # Gremlin Chat shelved
            "gremlin_messages_per_day": 0,
            "ai_report_row_cap": 50000,
        },
        "scale": {
            "monthly_rows": 500000,
            "rows_per_run": 100000,
            "grace_percentage": 0,
            "enrichment_credits": 0,  # Shelved for V1
            "max_schedules": 100,
            "allowed_frequencies": ["15min", "hourly", "daily", "weekly", "cron"],
            "concurrent_jobs": 4,
            "candidate_cap": 5000,
            "require_b2c": False,
            "allow_advanced_algos": True,
            "soql_runs_per_month": None,
            "salesforce_orgs": 1,  # Multi-org postponed to enterprise
            "users": 1,  # Seats sold as add-ons at $34/mo
            # AI Limits
            "gremlin_enabled": False,              # Gremlin Chat shelved
            "gremlin_messages_per_day": 0,
            "ai_report_row_cap": 100000,
        },
        "unleashed": {
            "monthly_rows": 2000000,
            "rows_per_run": 500000,
            "grace_percentage": 0,
            "enrichment_credits": 0,  # Shelved for V1
            "max_schedules": 500,
            "allowed_frequencies": ["15min", "hourly", "daily", "weekly", "cron"],
            "concurrent_jobs": 8,
            "candidate_cap": 50000,
            "require_b2c": False,
            "allow_advanced_algos": True,
            "soql_runs_per_month": None,
            "salesforce_orgs": 5,
            "users": 1,  # Seats sold as add-ons at $34/mo
            # AI Limits
            "gremlin_enabled": False,              # Gremlin Chat shelved
            "gremlin_messages_per_day": 0,
            "ai_report_row_cap": 500000,
            # BYO Enrichment - BUNDLED with unlimited limits
            "byo_enrichment_enabled": True,
            "byo_bundled": True,
            "byo_providers": ["foundrygraph", "fullenrich", "pdl"],
            "byo_max_rows_per_job": 1000,
            "byo_concurrent_enrichment_jobs": 10,
            "byo_daily_enrichment_rows": 50000,
            "byo_monthly_enrichment_rows": None,  # Unlimited
            "byo_scheduled_enrichment": True,
            "byo_priority_queue": True,
        },
    }
    TIER_LIMITS["business"] = TIER_LIMITS["scale"].copy()

    # BYO Provider Locker add-on limits (for Pro/Scale with add-on purchased)
    # Synced with config/enrichment_pricing.yaml addons.byo_provider_locker
    BYO_ADDON_LIMITS = {
        "pro": {
            "max_rows_per_job": 250,
            "concurrent_enrichment_jobs": 1,
            "daily_enrichment_rows": 1000,
            "monthly_enrichment_rows": 5000,
            "scheduled_enrichment": False,
            "providers": ["foundrygraph", "fullenrich", "pdl"],
        },
        "scale": {
            "max_rows_per_job": 500,
            "concurrent_enrichment_jobs": 2,
            "daily_enrichment_rows": 5000,
            "monthly_enrichment_rows": 25000,
            "scheduled_enrichment": True,
            "scheduled_frequencies": ["hourly", "daily"],
            "providers": ["foundrygraph", "fullenrich", "pdl"],
        },
    }


class QuotaExceeded(HTTPException):
    """Quota limit exceeded with upgrade prompt."""

    def __init__(
        self, message: str, limit_type: str, current: int, limit: int, tier: str
    ):
        super().__init__(
            status_code=429,
            detail={
                "error": "quota_exceeded",
                "message": message,
                "limit_type": limit_type,
                "current": current,
                "limit": limit,
                "tier": tier,
                "upgrade_url": "/pricing",
                "upgrade_message": f"Upgrade to Pro for {TIER_LIMITS['pro'][limit_type]} {limit_type}/day",
            },
        )


class UsageQuota:
    """Usage quota tracking with grace band."""

    def __init__(
        self,
        account_id: str,
        tier: str = "free",
        monthly_rows_used: int = 0,
        enrichment_credits_used: int = 0,
        active_schedules: int = 0,
        grace_uses_remaining: int = 1,
        last_grace_at: Optional[datetime] = None,
        reset_at: Optional[datetime] = None,
        timezone: str = "UTC",
        # Legacy fields for backwards compatibility
        daily_runs_used: int = 0,
        daily_rows_used: int = 0,
    ):
        self.account_id = account_id
        self.tier = tier
        self.monthly_rows_used = monthly_rows_used
        self.enrichment_credits_used = enrichment_credits_used
        self.active_schedules = active_schedules
        self.grace_uses_remaining = grace_uses_remaining
        self.last_grace_at = last_grace_at
        self.reset_at = reset_at or self._next_reset(timezone)
        self.timezone = timezone
        # Legacy compatibility
        self.daily_runs_used = daily_runs_used
        self.daily_rows_used = daily_rows_used

    @staticmethod
    def _next_reset(tz_name: str) -> datetime:
        """Calculate next quota reset (1st of next month)."""
        now = datetime.now(timezone.utc)
        if now.month == 12:
            next_month = datetime(now.year + 1, 1, 1, tzinfo=timezone.utc)
        else:
            next_month = datetime(now.year, now.month + 1, 1, tzinfo=timezone.utc)
        return next_month

    def get_limits(self) -> Dict[str, Any]:
        """Get limits for current tier."""
        return TIER_LIMITS.get(self.tier, TIER_LIMITS["free"])

    def check_can_run(
        self, estimated_rows: int = 0, credits_needed: int = 0
    ) -> Dict[str, Any]:
        """
        Check if a run is allowed based on monthly quotas.
        Returns: {allowed: bool, used_grace: bool, message: str}
        """
        limits = self.get_limits()

        # Check if quota reset is needed
        if datetime.now(timezone.utc) >= self.reset_at:
            return {
                "allowed": True,
                "used_grace": False,
                "needs_reset": True,
                "message": "Quota reset due",
            }

        # Check monthly rows quota
        if estimated_rows > 0:
            total_rows_after = self.monthly_rows_used + estimated_rows
            monthly_limit = limits["monthly_rows"]

            if total_rows_after > monthly_limit:
                # Check grace band
                grace_limit = monthly_limit * (1 + limits["grace_percentage"])
                if total_rows_after <= grace_limit and self.grace_uses_remaining > 0:
                    return {
                        "allowed": True,
                        "used_grace": True,
                        "message": f"Using grace allowance. {total_rows_after:,}/{int(grace_limit):,} rows this month. Upgrade for more.",
                        "show_upgrade": True,
                    }
                else:
                    remaining = max(0, monthly_limit - self.monthly_rows_used)
                    next_tier_limit = (
                        TIER_LIMITS["pro"]["monthly_rows"]
                        if self.tier == "free"
                        else TIER_LIMITS["scale"]["monthly_rows"]
                    )
                    return {
                        "allowed": False,
                        "message": f"Monthly row limit reached ({self.monthly_rows_used:,}/{monthly_limit:,}). You have {remaining:,} rows remaining. Resets {self.reset_at.strftime('%b %d')}. Upgrade for {next_tier_limit:,} rows/month.",
                        "show_upgrade": True,
                        "remaining": remaining,
                    }

        # Check enrichment credits
        if credits_needed > 0:
            total_credits_after = self.enrichment_credits_used + credits_needed
            credits_limit = limits["enrichment_credits"]

            if total_credits_after > credits_limit:
                remaining = max(0, credits_limit - self.enrichment_credits_used)
                return {
                    "allowed": False,
                    "message": f"Enrichment credit limit reached ({self.enrichment_credits_used}/{credits_limit}). Buy credit packs or upgrade.",
                    "show_upgrade": True,
                    "remaining_credits": remaining,
                }

        return {"allowed": True, "used_grace": False, "message": "Within quota limits"}

    def record_run(
        self, rows_returned: int, credits_consumed: int = 0, used_grace: bool = False
    ) -> None:
        """Record a completed run with row and credit usage."""
        self.monthly_rows_used += rows_returned
        self.enrichment_credits_used += credits_consumed
        # Legacy tracking
        self.daily_runs_used += 1
        self.daily_rows_used += rows_returned

        if used_grace:
            self.grace_uses_remaining = max(0, self.grace_uses_remaining - 1)
            self.last_grace_at = datetime.now(timezone.utc)

    def reset(self) -> None:
        """Reset monthly counters."""
        self.monthly_rows_used = 0
        self.enrichment_credits_used = 0
        # Legacy fields
        self.daily_runs_used = 0
        self.daily_rows_used = 0
        self.grace_uses_remaining = 1  # Reset grace
        self.reset_at = self._next_reset(self.timezone)


async def get_quota(account_id: str, db: AsyncSession) -> UsageQuota:
    """Get or create quota for account."""
    from ..models import UsageQuotaModel

    result = await db.execute(
        select(UsageQuotaModel).where(UsageQuotaModel.account_id == account_id)
    )
    quota_model = result.scalar_one_or_none()

    if not quota_model:
        # Create new quota
        quota_model = UsageQuotaModel(
            account_id=account_id,
            tier="free",
            monthly_rows_used=0,
            enrichment_credits_used=0,
            daily_runs_used=0,  # Legacy
            daily_rows_used=0,  # Legacy
            active_schedules=0,
            grace_uses_remaining=1,
            reset_at=UsageQuota._next_reset("UTC"),
            created_at=datetime.now(timezone.utc),
            updated_at=datetime.now(timezone.utc),
        )
        db.add(quota_model)
        await db.commit()

    # Handle models that don't have new fields yet (backwards compatibility)
    monthly_rows = getattr(
        quota_model, "monthly_rows_used", quota_model.daily_rows_used
    )
    enrichment_credits = getattr(quota_model, "enrichment_credits_used", 0)

    quota = UsageQuota(
        account_id=quota_model.account_id,
        tier=quota_model.tier,
        monthly_rows_used=monthly_rows,
        enrichment_credits_used=enrichment_credits,
        active_schedules=quota_model.active_schedules,
        grace_uses_remaining=quota_model.grace_uses_remaining,
        last_grace_at=quota_model.last_grace_at,
        reset_at=quota_model.reset_at if (quota_model.reset_at is None or quota_model.reset_at.tzinfo is not None) else quota_model.reset_at.replace(tzinfo=timezone.utc),
        timezone=getattr(quota_model, "timezone", "UTC"),
        # Legacy compatibility
        daily_runs_used=quota_model.daily_runs_used,
        daily_rows_used=quota_model.daily_rows_used,
    )

    # Auto-reset if needed
    if datetime.now(timezone.utc) >= quota.reset_at:
        quota.reset()
        update_values = {
            "monthly_rows_used": 0,
            "enrichment_credits_used": 0,
            "daily_runs_used": 0,
            "daily_rows_used": 0,
            "grace_uses_remaining": 1,
            "reset_at": quota.reset_at,
            "updated_at": datetime.now(timezone.utc),
        }
        await db.execute(
            update(UsageQuotaModel)
            .where(UsageQuotaModel.account_id == account_id)
            .values(**update_values)
        )
        await db.commit()

    return quota


async def record_quota_usage(
    account_id: str,
    db: AsyncSession,
    rows_returned: int,
    credits_consumed: int = 0,
    used_grace: bool = False,
) -> None:
    """Record quota usage after successful run."""
    from ..models import UsageQuotaModel

    update_values = {
        "monthly_rows_used": UsageQuotaModel.monthly_rows_used + rows_returned,
        "enrichment_credits_used": UsageQuotaModel.enrichment_credits_used
        + credits_consumed,
        # Legacy fields
        "daily_runs_used": UsageQuotaModel.daily_runs_used + 1,
        "daily_rows_used": UsageQuotaModel.daily_rows_used + rows_returned,
        "updated_at": datetime.now(timezone.utc),
    }

    if used_grace:
        update_values["grace_uses_remaining"] = UsageQuotaModel.grace_uses_remaining - 1
        update_values["last_grace_at"] = datetime.now(timezone.utc)

    await db.execute(
        update(UsageQuotaModel)
        .where(UsageQuotaModel.account_id == account_id)
        .values(**update_values)
    )
    await db.commit()
