from typing import Protocol

import sentry_sdk
from django.db.models.functions import Now
from ninja.errors import HttpError
from ninja.security import HttpBearer

from aiservice.background import fire_and_forget
from authapp.auth_utils import hash_api_key
from authapp.models import CFAPIKeys, Organizations, Subscriptions


class AuthenticatedRequest(Protocol):
    """Protocol for requests authenticated via API key middleware.

    These attributes are added by the ApiKey authentication class.
    """

    user: str  # user_id from the API key
    tier: str  # subscription tier
    api_key_id: int  # API key instance ID
    should_log_features: bool  # whether to log optimization features


async def check_subscription_status(request, user_id, tier, organization_id=None) -> bool:
    """Check if a user has a premium subscription that doesn't require feature logging.

    Attaches fetched organization and subscription objects to the request
    so downstream middleware can reuse them without re-querying.
    """
    try:
        org = None
        if organization_id:
            org = await Organizations.objects.filter(id=organization_id).afirst()
            request.organization = org

        subscription = await Subscriptions.objects.filter(user_id=user_id).afirst()
        request.subscription = subscription

        if tier is not None:
            return False

        if org and org.name == "codeflash-ai":
            return True
        if org and org.subscription:
            return False

        if subscription and subscription.plan_type.lower() in ["pro", "enterprise"]:
            return False
    except Exception as e:
        sentry_sdk.capture_exception(e)
        return False
    return True


class AuthBearer(HttpBearer):
    async def authenticate(self, request, token):
        hashed_token = hash_api_key(token)
        api_key_instance = await CFAPIKeys.objects.filter(key=hashed_token).afirst()
        if api_key_instance is None:
            raise HttpError(403, "Invalid API key")
        fire_and_forget(CFAPIKeys.objects.filter(id=api_key_instance.id).aupdate(last_used=Now()))
        request.user = api_key_instance.user_id
        request.tier = api_key_instance.tier
        request.api_key_id = api_key_instance.id
        request.organization_id = api_key_instance.organization_id
        request.should_log_features = await check_subscription_status(
            request, user_id=request.user, tier=request.tier, organization_id=request.organization_id
        )
        return token
