"""Quota enforcement for the Company List Builder feature."""

from __future__ import annotations

import logging
from datetime import datetime, timedelta, timezone
from typing import Any, Dict, Optional, Tuple

from sqlalchemy import select, update
from sqlalchemy.ext.asyncio import AsyncSession

from ..config.tier_limits import get_tier_config
from ..models import UsageEvent, UsageEventType, UsageQuotaModel, utcnow
from ..services.usage_gateway import UsageGateway

logger = logging.getLogger(__name__)


class ListBuilderQuota:
    """Org-level quota checks for list builder exports."""

    def __init__(self, db: AsyncSession, redis=None, job_ttl_seconds: int = 7200) -> None:
        self.db = db
        self.redis = redis
        self.job_ttl_seconds = job_ttl_seconds

    @staticmethod
    def _next_daily_reset(now: Optional[datetime] = None) -> datetime:
        now = now or datetime.now(timezone.utc)
        return datetime(now.year, now.month, now.day, tzinfo=timezone.utc) + timedelta(days=1)

    @staticmethod
    def _next_monthly_reset(now: Optional[datetime] = None) -> datetime:
        now = now or datetime.now(timezone.utc)
        if now.month == 12:
            return datetime(now.year + 1, 1, 1, tzinfo=timezone.utc)
        return datetime(now.year, now.month + 1, 1, tzinfo=timezone.utc)

    def _active_jobs_key(self, account_id: str) -> str:
        return f"listbuilder:account:{account_id}:jobs"

    async def _get_or_create_quota(self, account_id: str) -> UsageQuotaModel:
        result = await self.db.execute(
            select(UsageQuotaModel).where(UsageQuotaModel.account_id == account_id)
        )
        quota = result.scalar_one_or_none()
        if quota:
            return quota

        quota = UsageQuotaModel(
            account_id=account_id,
            tier="free",
            monthly_rows_used=0,
            enrichment_credits_used=0,
            daily_runs_used=0,
            daily_rows_used=0,
            list_builder_monthly_exports=0,
            list_builder_daily_exports=0,
            list_builder_daily_reset_at=self._next_daily_reset(),
            reset_at=self._next_monthly_reset(),
            created_at=utcnow(),
            updated_at=utcnow(),
        )
        self.db.add(quota)
        await self.db.commit()
        return quota

    async def _refresh_resets(self, quota: UsageQuotaModel) -> UsageQuotaModel:
        now = datetime.now(timezone.utc)
        updated = False

        if quota.list_builder_daily_reset_at is None or now >= quota.list_builder_daily_reset_at:
            quota.list_builder_daily_exports = 0
            quota.list_builder_daily_reset_at = self._next_daily_reset(now)
            updated = True

        if quota.reset_at is None or now >= quota.reset_at:
            quota.list_builder_monthly_exports = 0
            quota.reset_at = self._next_monthly_reset(now)
            updated = True

        if updated:
            quota.updated_at = utcnow()
            self.db.add(quota)
            await self.db.commit()

        return quota

    async def get_limits(
        self, account_id: str
    ) -> Tuple[str, Any, UsageQuotaModel]:
        quota = await self._get_or_create_quota(account_id)
        quota = await self._refresh_resets(quota)

        tier_name = (quota.tier or "free").lower()
        try:
            tier_config = get_tier_config(tier_name)
        except Exception as exc:  # noqa: BLE001 - fall back defensively
            logger.warning("Falling back to free tier limits for %s: %s", account_id, exc)
            tier_config = get_tier_config("free")
            tier_name = "free"

        return tier_name, tier_config, quota

    async def check_can_export(self, account_id: str, requested_rows: int) -> Dict[str, Any]:
        tier_name, tier_config, quota = await self.get_limits(account_id)

        monthly_limit = tier_config.list_builder_monthly_rows
        daily_limit = tier_config.list_builder_daily_rows

        monthly_remaining = max(
            0, monthly_limit - int(quota.list_builder_monthly_exports or 0)
        )
        daily_remaining = max(
            0, daily_limit - int(quota.list_builder_daily_exports or 0)
        )

        allowed_rows = min(max(requested_rows, 0), monthly_remaining, daily_remaining)

        return {
            "allowed": allowed_rows > 0,
            "allowed_rows": allowed_rows,
            "daily_remaining": daily_remaining,
            "monthly_remaining": monthly_remaining,
            "tier": tier_name,
        }

    async def record_export(
        self,
        account_id: str,
        rows_exported: int,
        *,
        idempotency_key: Optional[str] = None,
    ) -> None:
        if rows_exported <= 0:
            return

        await self._refresh_resets(await self._get_or_create_quota(account_id))

        await self.db.execute(
            update(UsageQuotaModel)
            .where(UsageQuotaModel.account_id == account_id)
            .values(
                list_builder_monthly_exports=UsageQuotaModel.list_builder_monthly_exports
                + rows_exported,
                list_builder_daily_exports=UsageQuotaModel.list_builder_daily_exports
                + rows_exported,
                updated_at=utcnow(),
            )
        )
        await self.db.commit()

        try:
            gateway = UsageGateway(self.db)
            await gateway.check_and_consume(
                account_id=account_id,
                resource="list_builder",
                amount=rows_exported,
                source="list_builder_export",
                idempotency_key=idempotency_key,
            )
        except Exception as exc:  # noqa: BLE001 - best effort for audit
            logger.warning(
                "Failed to record list builder usage in gateway for %s: %s",
                account_id,
                exc,
            )

        try:
            event = UsageEvent(
                account_id=account_id,
                event_type=UsageEventType.LIST_BUILDER_EXPORT,
                quantity=rows_exported,
                event_metadata={"feature": "company_list_builder"},
                created_at=utcnow(),
            )
            self.db.add(event)
            await self.db.commit()
        except Exception as exc:  # noqa: BLE001 - usage event is best-effort
            logger.warning("Failed to record list builder usage event: %s", exc)

    async def get_active_job_count(self, account_id: str) -> int:
        if not self.redis:
            return 0
        try:
            return int(await self.redis.scard(self._active_jobs_key(account_id)) or 0)
        except Exception as exc:  # noqa: BLE001 - Redis is non-critical
            logger.warning("Could not read active job count for %s: %s", account_id, exc)
            return 0

    async def mark_job_started(self, account_id: str, job_id: str) -> None:
        if not self.redis:
            return
        key = self._active_jobs_key(account_id)
        try:
            await self.redis.sadd(key, job_id)
            await self.redis.expire(key, self.job_ttl_seconds)
        except Exception as exc:  # noqa: BLE001
            logger.warning("Failed to register active job %s for %s: %s", job_id, account_id, exc)

    async def clear_job(self, account_id: str, job_id: str) -> None:
        if not self.redis:
            return
        try:
            await self.redis.srem(self._active_jobs_key(account_id), job_id)
        except Exception as exc:  # noqa: BLE001
            logger.warning("Failed to clear active job %s for %s: %s", job_id, account_id, exc)
