from dateutil.relativedelta import relativedelta
from flask import session, current_app, flash
from sqlalchemy.exc import SQLAlchemyError
from sqlalchemy.sql.expression import text

from common.extensions import db, cache_manager
from common.models.entitlements import PartnerServiceLicenseTier, License, LicenseUsage, LicensePeriod, PeriodStatus
from common.models.user import Partner, PartnerTenant
from common.services.entitlements import LicensePeriodServices
from common.utils.database import Database
from common.utils.eveai_exceptions import EveAINoManagementPartnerService, EveAINoActiveLicense, \
    EveAIStorageQuotaExceeded, EveAIEmbeddingQuotaExceeded, EveAIInteractionQuotaExceeded, EveAILicensePeriodsExceeded, \
    EveAIException
from common.utils.model_logging_utils import set_logging_information, update_logging_information
from datetime import datetime as dt, timezone as tz

from common.utils.security_utils import current_user_has_role


class LicenseUsageServices:
    @staticmethod
    def check_storage_and_embedding_quota(tenant_id: int, file_size_mb: float) -> None:
        """
        Check if a tenant can add a new document without exceeding storage and embedding quotas

        Args:
            tenant_id: ID of the tenant
            file_size_mb: Size of the file in MB

        Raises:
            EveAIStorageQuotaExceeded: If storage quota would be exceeded
            EveAIEmbeddingQuotaExceeded: If embedding quota would be exceeded
            EveAINoActiveLicense: If no active license is found
            EveAIException: For other errors
        """
        # Get active license period
        license_period = LicensePeriodServices.find_current_license_period_for_usage(tenant_id)
        # Early return if both overruns are allowed - no need to check usage at all
        if license_period.additional_storage_allowed and license_period.additional_embedding_allowed:
            return

        # Check storage quota only if overruns are not allowed
        if not license_period.additional_storage_allowed:
            LicenseUsageServices._validate_storage_quota(license_period, file_size_mb)

        # Check embedding quota only if overruns are not allowed
        if not license_period.additional_embedding_allowed:
            LicenseUsageServices._validate_embedding_quota(license_period, file_size_mb)

    @staticmethod
    def check_embedding_quota(tenant_id: int, file_size_mb: float) -> None:
        """
        Check if a tenant can re-embed a document without exceeding embedding quota

        Args:
            tenant_id: ID of the tenant
            file_size_mb: Size of the file in MB

        Raises:
            EveAIEmbeddingQuotaExceeded: If embedding quota would be exceeded
            EveAINoActiveLicense: If no active license is found
            EveAIException: For other errors
        """
        # Get active license period
        license_period = LicensePeriodServices.find_current_license_period_for_usage(tenant_id)
        # Early return if both overruns are allowed - no need to check usage at all
        if license_period.additional_embedding_allowed:
            return

        # Check embedding quota
        LicenseUsageServices._validate_embedding_quota(license_period, file_size_mb)

    @staticmethod
    def check_interaction_quota(tenant_id: int) -> None:
        """
        Check if a tenant can execute a specialist without exceeding interaction quota. As it is impossible to estimate
        the number of interaction tokens, we only check if the interaction quota are exceeded. So we might have a
        limited overrun.

        Args:
            tenant_id: ID of the tenant

        Raises:
            EveAIInteractionQuotaExceeded: If interaction quota would be exceeded
            EveAINoActiveLicense: If no active license is found
            EveAIException: For other errors
        """
        # Get active license period
        license_period = LicensePeriodServices.find_current_license_period_for_usage(tenant_id)
        # Early return if both overruns are allowed - no need to check usage at all
        if license_period.additional_interaction_allowed:
            return

        # Convert tokens to M tokens and check interaction quota
        LicenseUsageServices._validate_interaction_quota(license_period)

    @staticmethod
    def _validate_storage_quota(license_period: LicensePeriod, additional_mb: float) -> None:
        """Check storage quota and raise exception if exceeded"""
        current_storage = license_period.license_usage.storage_mb_used or 0
        projected_storage = current_storage + additional_mb
        max_storage = license_period.max_storage_mb

        # Hard limit check (we only get here if overruns are NOT allowed)
        if projected_storage > max_storage:
            raise EveAIStorageQuotaExceeded(
                current_usage=current_storage,
                limit=max_storage,
                additional=additional_mb
            )

    @staticmethod
    def _validate_embedding_quota(license_period: LicensePeriod, additional_mb: float) -> None:
        """Check embedding quota and raise exception if exceeded"""
        current_embedding = license_period.license_usage.embedding_mb_used or 0
        projected_embedding = current_embedding + additional_mb
        max_embedding = license_period.included_embedding_mb

        # Hard limit check (we only get here if overruns are NOT allowed)
        if projected_embedding > max_embedding:
            raise EveAIEmbeddingQuotaExceeded(
                current_usage=current_embedding,
                limit=max_embedding,
                additional=additional_mb
            )

    @staticmethod
    def _validate_interaction_quota(license_period) -> None:
        """Check interaction quota and raise exception if exceeded (tokens in millions). We might have an overrun!"""
        current_tokens = license_period.license_usage.interaction_total_tokens_used / 1_000_000 or 0
        max_tokens = license_period.included_interaction_tokens

        # Hard limit check (we only get here if overruns are NOT allowed)
        if current_tokens > max_tokens:
            raise EveAIInteractionQuotaExceeded(
                current_usage=current_tokens,
                limit=max_tokens
            )