from __future__ import annotations

import asyncio
import hashlib
import json
import os
import time
from dataclasses import dataclass
import math
import re
from typing import Any, Dict, Optional, Tuple

import redis.asyncio as aioredis
from fastapi import Depends, Header, HTTPException, Request
from fastapi.security import HTTPAuthorizationCredentials
from sqlalchemy.ext.asyncio import AsyncSession

from ...db import get_session
from ...settings import (
    API_QPS_DEFAULT,
    REDIS_URL,
    DEV_ACCOUNT_ID,
)
from ...repos.api_repo import ApiRepo, ApiContext
from ...model_defs.api_models import ApiUsage
from sqlalchemy import select


# Simple in-proc token bucket per key and per IP
class _TokenBucket:
    def __init__(self, rate_per_sec: float, capacity: int):
        self.rate = float(rate_per_sec)
        self.capacity = int(max(1, capacity))
        self.tokens = float(self.capacity)
        self.last = time.monotonic()

    def allow(self) -> bool:
        now = time.monotonic()
        elapsed = now - self.last
        self.last = now
        self.tokens = min(self.capacity, self.tokens + elapsed * self.rate)
        if self.tokens >= 1.0:
            self.tokens -= 1.0
            return True
        return False

    def retry_after_seconds(self) -> int:
        """Approximate seconds until 1 token is available."""
        if self.tokens >= 1.0:
            return 0
        need = 1.0 - max(0.0, self.tokens)
        if self.rate <= 0:
            return 1
        return max(1, math.ceil(need / self.rate))

    def remaining_int(self) -> int:
        return max(0, int(self.tokens))


_buckets_key: Dict[str, _TokenBucket] = {}
_buckets_ip: Dict[str, _TokenBucket] = {}


def _qps_for_tier(tier: str) -> int:
    tier = (tier or "free").upper()
    return int(API_QPS_DEFAULT.get(tier, API_QPS_DEFAULT.get("FREE", 2)))


async def _get_redis() -> Optional[aioredis.Redis]:
    url = os.getenv("REDIS_URL", REDIS_URL or "")
    if not url:
        return None
    try:
        return aioredis.from_url(url, decode_responses=True)
    except Exception:
        return None


# Idempotency store (local + optional Redis)
class IdempotencyStore:
    def __init__(self):
        self._local: Dict[str, Tuple[float, Dict[str, Any]]] = {}
        self._ttl = int(os.getenv("API_IDEMPOTENCY_TTL_SEC", "600") or 600)
        self._lock = asyncio.Lock()

    def _key(self, api_key_prefix: str, idem_key: str) -> str:
        return f"idem:{api_key_prefix}:{hashlib.md5((idem_key or '').encode()).hexdigest()}"

    async def get(
        self, api_key_prefix: str, idem_key: Optional[str]
    ) -> Optional[Dict[str, Any]]:
        if not idem_key:
            return None
        # Try Redis first
        r = await _get_redis()
        cache_key = self._key(api_key_prefix, idem_key)
        if r:
            try:
                raw = await r.get(cache_key)
                if raw:
                    return json.loads(raw)
            except Exception:
                pass
        # Fallback local
        async with self._lock:
            item = self._local.get(cache_key)
            if not item:
                return None
            ts, payload = item
            if time.time() - ts > self._ttl:
                self._local.pop(cache_key, None)
                return None
            return payload

    async def set(
        self, api_key_prefix: str, idem_key: Optional[str], payload: Dict[str, Any]
    ) -> None:
        if not idem_key:
            return
        cache_key = self._key(api_key_prefix, idem_key)
        ttl = self._ttl
        # Redis
        r = await _get_redis()
        if r:
            try:
                await r.set(cache_key, json.dumps(payload), ex=ttl)
                return
            except Exception:
                pass
        # Local
        async with self._lock:
            self._local[cache_key] = (time.time(), payload)


idem_store = IdempotencyStore()


@dataclass
class ApiRequestContext:
    api_key_prefix: str
    account_id: str
    tier: str
    qps_limit: int


_AUTH_RX = re.compile(r"^\s*(Bearer|Api-Key)\s+(.+)$", re.I)


async def require_api_key(
    request: Request,
    authorization: Optional[str] = Header(default=None, alias="Authorization"),
    x_api_key: Optional[str] = Header(default=None, alias="X-API-Key"),
    db: AsyncSession = Depends(get_session),
) -> ApiRequestContext:
    # Dev bypass: allow requests without DB/key only in explicit dev environments,
    # and only when explicitly enabled (opt-in).
    env = (os.getenv("ENV") or "").strip().lower()
    if env in ("dev", "development") and os.getenv("FM_BYPASS_API_KEY", "0") == "1":
        # Populate request.state with reasonable defaults for downstream
        request.state.account_id = DEV_ACCOUNT_ID
        request.state.api_key_prefix = "dev"
        request.state.tier = "DEV"
        try:
            request.state.x_ratelimit_limit = 1000
            request.state.x_ratelimit_remaining = 999
        except Exception:
            pass
        return ApiRequestContext(
            api_key_prefix="dev",
            account_id=DEV_ACCOUNT_ID,
            tier="DEV",
            qps_limit=1000,
        )

    # Accept either X-API-Key or Authorization (Bearer|Api-Key <token>)
    token: str = ""
    scheme: Optional[str] = None
    if x_api_key:
        token = (x_api_key or "").strip()
        scheme = "api-key"
    elif authorization:
        m = _AUTH_RX.match(authorization or "")
        scheme = (m.group(1).lower() if m else None) if authorization else None
        token = (m.group(2) if m else authorization).strip()
    if not token:
        raise HTTPException(
            status_code=401,
            detail={"error": "UNAUTHORIZED", "message": "Missing API key"},
        )

    # Bearer JWT support (Sheets/Clerk tokens) for v1 endpoints
    if scheme == "bearer" and not token.startswith("fm_"):
        from ...auth_core import get_current_identity

        ident = await get_current_identity(
            HTTPAuthorizationCredentials(scheme="Bearer", credentials=token)
        )
        account_id = ident.get("org_id")
        if not account_id:
            raise HTTPException(
                status_code=401,
                detail={"error": "UNAUTHORIZED", "message": "No organization on token"},
            )

        api_key_prefix = f"jwt:{account_id}"
        tier = "JWT"
        qps_limit = _qps_for_tier(tier)

        request.state.account_id = account_id
        request.state.api_key_prefix = api_key_prefix
        request.state.tier = tier
        try:
            request.state.x_ratelimit_limit = qps_limit
            request.state.x_ratelimit_remaining = qps_limit
        except Exception:
            pass

        return ApiRequestContext(
            api_key_prefix=api_key_prefix,
            account_id=account_id,
            tier=tier,
            qps_limit=qps_limit,
        )

    repo = ApiRepo(db)
    ctx: Optional[ApiContext] = await repo.get_key(token)
    if not ctx:
        raise HTTPException(
            status_code=401,
            detail={"error": "UNAUTHORIZED", "message": "Invalid API key"},
        )
    if ctx.status != "active":
        raise HTTPException(
            status_code=401,
            detail={"error": "UNAUTHORIZED", "message": "API key disabled"},
        )

    # Quota enforcement (monthly)
    if ctx.monthly_quota is not None and ctx.monthly_quota >= 0:
        now = time.gmtime()
        period = now.tm_year * 100 + now.tm_mon
        stmt = select(ApiUsage).where(
            ApiUsage.api_key == ctx.api_key_prefix, ApiUsage.period_yyyymm == period
        )
        r = await db.execute(stmt)
        usage = r.scalar_one_or_none()
        if usage and usage.calls is not None and usage.calls >= int(ctx.monthly_quota):
            raise HTTPException(
                status_code=429,
                detail={"error": "RATE_LIMIT", "message": "Monthly quota exceeded"},
            )

    # Per-key rate limiting (token bucket)
    key = f"key:{ctx.api_key_prefix}"
    bucket = _buckets_key.get(key)
    key_qps = int(ctx.qps_limit or _qps_for_tier(ctx.tier))
    if not bucket:
        bucket = _TokenBucket(rate_per_sec=key_qps, capacity=max(1, key_qps))
        _buckets_key[key] = bucket
    if not bucket.allow():
        headers = {
            "Retry-After": str(bucket.retry_after_seconds()),
            "X-RateLimit-Limit": str(key_qps),
            "X-RateLimit-Remaining": "0",
        }
        raise HTTPException(
            status_code=429,
            detail={"error": "RATE_LIMIT", "message": "Per-key rate limit exceeded"},
            headers=headers,
        )

    # Per-IP rate limiting with tier defaults
    ip = request.client.host if request.client else "unknown"
    ipq = _qps_for_tier(ctx.tier)
    ip_bucket = _buckets_ip.get(ip)
    if not ip_bucket:
        ip_bucket = _TokenBucket(rate_per_sec=ipq, capacity=max(1, ipq))
        _buckets_ip[ip] = ip_bucket
    if not ip_bucket.allow():
        headers = {
            "Retry-After": str(ip_bucket.retry_after_seconds()),
            "X-RateLimit-Limit": str(ipq),
            "X-RateLimit-Remaining": "0",
        }
        raise HTTPException(
            status_code=429,
            detail={"error": "RATE_LIMIT", "message": "Per-IP rate limit exceeded"},
            headers=headers,
        )

    # Attach to request.state for downstream
    request.state.account_id = ctx.account_id
    request.state.api_key_prefix = ctx.api_key_prefix
    request.state.tier = ctx.tier

    # Expose rate-limit hints for middleware headers (per key)
    try:
        request.state.x_ratelimit_limit = key_qps
        request.state.x_ratelimit_remaining = bucket.remaining_int()
    except Exception:
        pass

    return ApiRequestContext(
        api_key_prefix=ctx.api_key_prefix,
        account_id=ctx.account_id,
        tier=ctx.tier,
        qps_limit=key_qps,
    )
