"""This file and its contents are licensed under the Apache License 2.0. Please see the included NOTICE for copyright information and LICENSE for a copy of the license.
"""
import logging
from typing import List

from django.conf import settings
from django.db import models
from django.utils.translation import gettext_lazy as _
from tasks.models import PredictionMeta

logger = logging.getLogger(__name__)


class ModelProviders(models.TextChoices):
    OPENAI = 'OpenAI', _('OpenAI')
    AZURE_OPENAI = 'AzureOpenAI', _('AzureOpenAI')
    AZURE_AI_FOUNDRY = 'AzureAIFoundry', _('AzureAIFoundry')
    VERTEX_AI = 'VertexAI', _('VertexAI')
    GEMINI = 'Gemini', _('Gemini')
    ANTHROPIC = 'Anthropic', _('Anthropic')
    CUSTOM = 'Custom', _('Custom')


class ModelProviderConnectionScopes(models.TextChoices):
    ORG = 'Organization', _('Organization')
    USER = 'User', _('User')
    MODEL = 'Model', _('Model')


class ModelProviderConnection(models.Model):

    provider = models.CharField(max_length=255, choices=ModelProviders.choices, default=ModelProviders.OPENAI)

    api_key = models.TextField(_('api_key'), null=True, blank=True, help_text='Model provider API key')

    auth_token = models.TextField(_('auth_token'), null=True, blank=True, help_text='Model provider Auth token')

    deployment_name = models.CharField(max_length=512, null=True, blank=True, help_text='Azure OpenAI deployment name')

    endpoint = models.CharField(max_length=512, null=True, blank=True, help_text='Azure OpenAI endpoint')

    google_application_credentials = models.TextField(
        _('google application credentials'),
        null=True,
        blank=True,
        help_text='The content of GOOGLE_APPLICATION_CREDENTIALS json file',
    )

    google_project_id = models.CharField(
        _('google project id'), max_length=255, null=True, blank=True, help_text='Google project ID'
    )

    google_location = models.CharField(
        _('google location'), max_length=255, null=True, blank=True, help_text='Google project location'
    )

    cached_available_models = models.CharField(
        max_length=4096, null=True, blank=True, help_text='List of available models from the provider'
    )

    scope = models.CharField(
        max_length=255, choices=ModelProviderConnectionScopes.choices, default=ModelProviderConnectionScopes.ORG
    )

    organization = models.ForeignKey(
        'organizations.Organization', on_delete=models.CASCADE, related_name='model_provider_connections', null=True
    )

    created_by = models.ForeignKey(
        settings.AUTH_USER_MODEL,
        related_name='created_model_provider_connections',
        on_delete=models.SET_NULL,
        null=True,
    )

    # Future work - add foreign key for modelinterface / modelinstance

    created_at = models.DateTimeField(_('created at'), auto_now_add=True)

    updated_at = models.DateTimeField(_('updated at'), auto_now=True)

    is_internal = models.BooleanField(
        _('is_internal'),
        default=False,
        help_text='Whether the model provider connection is internal, not visible to the user',
        null=True,
        blank=True,
    )

    budget_limit = models.FloatField(
        _('budget_limit'),
        null=True,
        blank=True,
        default=None,
        help_text='Budget limit for the model provider connection (null if unlimited)',
    )

    budget_last_reset_date = models.DateTimeField(
        _('budget_last_reset_date'),
        null=True,
        blank=True,
        default=None,
        help_text='Date and time the budget was last reset',
    )

    budget_reset_period = models.CharField(
        _('budget_reset_period'),
        max_length=20,
        choices=[
            ('Monthly', 'Monthly'),
            ('Yearly', 'Yearly'),
        ],
        null=True,
        blank=True,
        default=None,
        help_text='Budget reset period for the model provider connection (null if not reset)',
    )

    budget_total_spent = models.FloatField(
        _('budget_total_spent'),
        null=True,
        blank=True,
        default=None,
        help_text='Tracked total budget spent for the given provider connection within the current budget period',
    )

    budget_alert_threshold = models.FloatField(
        _('budget_alert_threshold'),
        null=True,
        blank=True,
        default=None,
        help_text='Budget alert threshold for the given provider connection',
    )

    # Check if user is Admin or Owner
    # This will need to be updated if we ever use this model in LSO as `is_owner` and
    # `is_administrator` only exist in LSE
    def has_permission(self, user):
        return (
            user.is_administrator or user.is_owner or user.is_manager
        ) and user.active_organization_id == self.organization_id

    def update_budget_total_spent_from_predictions_meta(self, predictions_meta: List[PredictionMeta]):
        total_cost = sum(meta.total_cost or 0 for meta in predictions_meta)
        # opting for the goofy "self.budget_total_spent or 0" to avoid a db migration
        self.budget_total_spent = (self.budget_total_spent or 0) + total_cost
        self.save(update_fields=['budget_total_spent'])

    def has_reached_budget_limit(self):
        if (
            self.is_internal
            and self.budget_total_spent
            and self.budget_limit
            and self.budget_total_spent > self.budget_limit
        ):
            logger.info(
                f'Model connection {self.id} has reached the budget limit: '
                f'{self.budget_total_spent} > {self.budget_limit}'
            )
            return True
        return False