from __future__ import annotations

import logging
from typing import TYPE_CHECKING, Annotated

from fastapi import Depends, HTTPException, Request

from codeflash_api._config import settings
from codeflash_api.auth._keys import (
    extract_bearer_token,
    hash_api_key,
)
from codeflash_api.auth.models import (
    SUBSCRIPTION_PLANS,
    AuthenticatedUser,
    SubscriptionInfo,
    is_codeflash_employee,
)

if TYPE_CHECKING:
    from codeflash_api.db._queries import Queries

log = logging.getLogger(__name__)

ENDPOINT_TOKEN_COST: dict[str, int] = {
    "optimize": 10,
    "optimize-line-profiler": 10,
    "testgen": 20,
    "refinement": 20,
    "explain": 10,
}


async def require_auth(
    request: Request,
) -> AuthenticatedUser:
    """
    FastAPI dependency that authenticates via Bearer token.
    """
    authorization = request.headers.get("Authorization")
    if not authorization:
        raise HTTPException(status_code=403, detail="Invalid API key")

    token = extract_bearer_token(authorization)
    hashed = hash_api_key(token)

    queries: Queries = request.app.state.queries
    api_key = await queries.get_api_key_by_hash(hashed)
    if api_key is None:
        raise HTTPException(status_code=403, detail="Invalid API key")

    should_log = await _check_should_log_features(
        queries,
        api_key.user_id or "",
        api_key.tier,
        api_key.organization_id,
    )

    user = AuthenticatedUser(
        user_id=api_key.user_id or "",
        tier=api_key.tier,
        api_key_id=api_key.id,
        organization_id=api_key.organization_id,
        should_log_features=should_log,
    )
    request.state.user = user
    return user


async def check_rate_limit(
    request: Request,
    user: Annotated[AuthenticatedUser, Depends(require_auth)],
) -> None:
    """
    FastAPI dependency that enforces per-user per-endpoint rate limits.
    """
    if is_codeflash_employee(user.user_id):
        return

    path = request.url.path
    if not path.startswith("/ai/"):
        return

    cache = request.app.state.rate_limit_cache
    key = f"ratelimit:user:{user.user_id}:{path}"
    current: int = cache.get(key, 0)

    if current >= settings.rate_limit_max:
        log.warning(
            "Rate limit exceeded for user=%s path=%s",
            user.user_id,
            path,
        )
        raise HTTPException(status_code=429, detail="Rate limit exceeded")

    cache[key] = current + 1


async def track_usage(
    request: Request,
    user: Annotated[AuthenticatedUser, Depends(require_auth)],
) -> None:
    """
    FastAPI dependency that tracks and enforces subscription usage.
    """
    if is_codeflash_employee(user.user_id):
        request.state.subscription_info = SubscriptionInfo(
            user_id=user.user_id,
            tier="EMPLOYEE",
            used=0,
            limit=None,
        )
        return

    path = request.url.path.removeprefix("/ai/").split("?")[0]
    cost = ENDPOINT_TOKEN_COST.get(path, 0)
    if cost == 0:
        return

    queries: Queries = request.app.state.queries

    if user.organization_id:
        org = await queries.get_organization(user.organization_id)
        if org and org.subscription:
            request.state.subscription_info = SubscriptionInfo(
                user_id=user.user_id,
                tier="ENTERPRISE",
                used=0,
                limit=None,
                organization_id=user.organization_id,
            )
            return

    subscription = await queries.get_subscription(user.user_id)
    if subscription is None:
        plan = "FREE"
        limit = int(SUBSCRIPTION_PLANS["FREE"]["optimizations"])  # type: ignore[call-overload]
        subscription = await queries.create_subscription(
            user_id=user.user_id,
            plan_type=plan,
            optimizations_limit=limit,
        )
        log.warning(
            "Created backup subscription for user=%s",
            user.user_id,
        )

    if subscription.subscription_status != "active":
        raise HTTPException(
            status_code=403,
            detail={
                "error": "Subscription is not active",
                "status": subscription.subscription_status,
            },
        )

    new_used = subscription.optimizations_used + cost
    if new_used > subscription.optimizations_limit:
        raise HTTPException(
            status_code=403,
            detail={
                "error": "Usage limit exceeded",
                "used": subscription.optimizations_used,
                "limit": subscription.optimizations_limit,
                "tier": subscription.plan_type,
            },
        )

    await queries.increment_usage(user_id=user.user_id, cost=cost)

    request.state.subscription_info = SubscriptionInfo(
        user_id=user.user_id,
        tier=subscription.plan_type or "FREE",
        used=new_used,
        limit=subscription.optimizations_limit,
        organization_id=user.organization_id,
    )


async def _check_should_log_features(
    queries: Queries,
    user_id: str,
    tier: str | None,
    organization_id: str | None,
) -> bool:
    """
    Return True if the user's code should be logged for feature extraction.
    """
    if tier is not None:
        return False

    if organization_id:
        org = await queries.get_organization(organization_id)
        if org and (org.name == "codeflash-ai" or org.subscription):
            return False

    sub = await queries.get_subscription(user_id)
    return not (sub and sub.plan_type in ("PRO", "ENTERPRISE"))
