import datetime
import logging

import sentry_sdk
from asgiref.sync import iscoroutinefunction, markcoroutinefunction
from django.db import IntegrityError
from django.db.models import F
from django.http import JsonResponse
from django.utils.decorators import async_only_middleware
from django.utils.timezone import now

from aiservice.common_utils import is_codeflash_employee
from authapp.models import Organizations, Subscriptions

SUBSCRIPTION_PLANS = {
    "FREE": {"name": "Free", "optimizations": 4000, "price": 0},
    "PRO": {
        "name": "Pro",
        "optimizations": 100000,
        "monthlyPrice": 3000,  # $30.00
        "yearlyPrice": 30000,  # $300.00
    },
    "ENTERPRISE": {"name": "Enterprise", "optimizations": "Custom", "price": "Custom"},
}

PRICE_ID_KEYS = {"PRO_MONTHLY": "STRIPE_PRO_PRICE_MONTHLY_ID", "PRO_YEARLY": "STRIPE_PRO_PRICE_YEARLY_ID"}
# subscription_utils.py


def add_months_safe(date: datetime.datetime | datetime.date, months: int) -> datetime.date:
    """Safely add months to a date, preserving the day if possible."""
    year = date.year + (date.month + months - 1) // 12
    month = (date.month + months - 1) % 12 + 1
    day = date.day

    # Get the last valid day of the new month
    last_day = (datetime.date(year, month % 12 + 1, 1) - datetime.timedelta(days=1)).day
    return datetime.date(year, month, min(day, last_day))


def get_next_subscription_period(period_end: datetime.datetime | datetime.date):
    start = add_months_safe(period_end, 1)
    end = add_months_safe(start, 1)
    return {"start": start, "end": end}


async def check_and_reset_subscription_period(user_id: int, subscription: Subscriptions | None = None):
    """Check if the subscription period has ended.
    If yes, reset usage and move to the next billing cycle.
    """
    if subscription is None:
        subscription = await Subscriptions.objects.filter(user_id=user_id).afirst()
    if not subscription:
        return None

    current_time = now()
    period_end = subscription.current_period_end
    # Pro User Stripe handles billing cycle, so we only need to manage FREE plan here
    if subscription.plan_type == "FREE":
        if period_end and current_time > period_end:
            # Reset usage and roll period
            next_period = get_next_subscription_period(period_end)  # ty: ignore[invalid-argument-type]

            subscription.optimizations_used = 0
            subscription.current_period_start = next_period["start"]
            subscription.current_period_end = next_period["end"]
            await subscription.asave(update_fields=["optimizations_used", "current_period_start", "current_period_end"])
        return subscription

    return subscription


ENDPOINT_TOKEN_COST = {"optimize": 10, "optimize-line-profiler": 10, "testgen": 20, "refinement": 20, "explain": 10}


@async_only_middleware
class TrackUsageMiddleware:
    def __init__(self, get_response) -> None:
        self.get_response = get_response
        if iscoroutinefunction(self.get_response):
            logging.debug("TrackUsageMiddleware is async.")
            markcoroutinefunction(self)  # ty:ignore[invalid-argument-type]

    async def __call__(self, request):
        endpoint = request.path.replace("/ai/", "").split("?")[0]
        cost = ENDPOINT_TOKEN_COST.get(endpoint, 0)

        user_id = getattr(request, "user", None)
        if not user_id:
            return JsonResponse({"error": "Authentication required for rate-limited endpoint"}, status=401)
        if is_codeflash_employee(user_id):
            request.subscription_info = {"userId": user_id, "tier": "EMPLOYEE", "used": 0, "limit": float("inf")}
            return await self.get_response(request)

        # Check if API key is linked to an organization with subscription
        organization_id = getattr(request, "organization_id", None)
        if organization_id:
            org = getattr(request, "organization", None)
            if org is None:
                org = await Organizations.objects.filter(id=organization_id).afirst()
            if org and org.subscription:
                request.subscription_info = {
                    "userId": user_id,
                    "tier": "ENTERPRISE",
                    "used": 0,
                    "limit": float("inf"),
                    "organizationId": organization_id,
                }
                return await self.get_response(request)

        try:
            # Get or create subscription (reuse from auth if available)
            subscription = getattr(request, "subscription", None)
            if subscription is None:
                subscription = await Subscriptions.objects.filter(user_id=user_id).afirst()

            if not subscription:
                # Subscription is now created during login in cf-webapp
                # As a backup, creating one here (this should be rare only now and for old users)
                logging.warning("No subscription found for user %s. Creating backup subscription.", user_id)
                try:
                    subscription = await Subscriptions.objects.acreate(
                        user_id=user_id,
                        plan_type="FREE",
                        optimizations_limit=SUBSCRIPTION_PLANS["FREE"]["optimizations"],
                        subscription_status="active",
                        optimizations_used=cost,  # Charge for this request
                    )
                    logging.info("Created backup subscription for user %s", user_id)
                    request.subscription_info = {
                        "userId": user_id,
                        "tier": subscription.plan_type,
                        "used": cost,
                        "limit": subscription.optimizations_limit,
                    }
                    return await self.get_response(request)
                except IntegrityError:
                    subscription = await Subscriptions.objects.filter(user_id=user_id).afirst()
                    if not subscription:
                        logging.exception("Failed to create or fetch subscription for user %s", user_id)
                        return JsonResponse({"error": "Failed to initialize user subscription"}, status=500)

            if subscription.subscription_status != "active":
                logging.warning(
                    "403 subscription inactive: user_id=%s, status=%s, endpoint=%s",
                    user_id,
                    subscription.subscription_status,
                    endpoint,
                )
                return JsonResponse(
                    {"error": "Subscription is not active", "status": subscription.subscription_status}, status=403
                )

            # Lazy reset monthly usage
            subscription = await check_and_reset_subscription_period(user_id, subscription=subscription)
            current_used = subscription.optimizations_used or 0

            if current_used + cost > subscription.optimizations_limit:
                logging.warning(
                    "403 usage limit exceeded: user_id=%s, used=%s, limit=%s, tier=%s, endpoint=%s",
                    user_id,
                    current_used,
                    subscription.optimizations_limit,
                    subscription.plan_type,
                    endpoint,
                )
                return JsonResponse(
                    {
                        "error": "Usage limit exceeded",
                        "used": current_used,
                        "limit": subscription.optimizations_limit,
                        "tier": subscription.plan_type,
                    },
                    status=403,
                )

            # Atomically increment usage using F() expressions to prevent race conditions.
            # This ensures concurrent requests each increment the counter correctly
            # instead of overwriting each other's updates (lost update problem).
            await Subscriptions.objects.filter(user_id=user_id).aupdate(
                optimizations_used=F("optimizations_used") + cost,
                total_lifetime_optimizations=F("total_lifetime_optimizations") + cost,
            )

            # Compute new usage from the known current value instead of re-reading from DB
            new_used = current_used + cost

            logging.debug(
                "track_usage_middleware.py|__call__|Atomic update completed: user_id=%s, endpoint=%s, cost=%s, new_used=%s",
                user_id,
                endpoint,
                cost,
                new_used,
            )

            # Attach subscription info to request
            request.subscription_info = {
                "userId": user_id,
                "tier": subscription.plan_type,
                "used": new_used,
                "limit": subscription.optimizations_limit,
            }

            return await self.get_response(request)

        except IntegrityError as e:
            # Handling database constraint violations more specifically
            if "duplicate key value violates unique constraint" in str(e):
                if "subscriptions_pkey" in str(e):
                    logging.warning("Subscription creation race condition for user %s: %s", user_id, e)
                    return JsonResponse({"error": "Subscription initialization conflict. Please retry."}, status=429)
                if "subscriptions_user_id_key" in str(e):
                    logging.warning("User already has subscription for user %s: %s", user_id, e)
                    return JsonResponse({"error": "User subscription already exists. Please refresh."}, status=409)
            sentry_sdk.capture_exception(e)
            logging.exception("Database integrity error for user %s", user_id)
            return JsonResponse({"error": "Database constraint error. Please contact support."}, status=500)
        except Exception as e:
            sentry_sdk.capture_exception(e)
            logging.exception("Unexpected error tracking usage for user %s", user_id)
            return JsonResponse({"error": "Service temporarily unavailable. Please try again."}, status=500)
