from __future__ import annotations

import datetime
from typing import TYPE_CHECKING, Any

from plain import models
from plain.auth import get_user_model
from plain.exceptions import ValidationError
from plain.models import transaction, types
from plain.models.db import IntegrityError
from plain.runtime import SettingsReference
from plain.utils import timezone

from .exceptions import OAuthUserAlreadyExistsError

if TYPE_CHECKING:
    from .providers import OAuthToken, OAuthUser


@models.register_model
class OAuthConnection(models.Model):
    created_at: datetime.datetime = types.DateTimeField(auto_now_add=True)
    updated_at: datetime.datetime = types.DateTimeField(auto_now=True)

    user = types.ForeignKeyField(
        SettingsReference("AUTH_USER_MODEL"),
        on_delete=models.CASCADE,
    )

    # The key used to refer to this provider type (in settings)
    provider_key: str = types.CharField(max_length=100)

    # The unique ID of the user on the provider's system
    provider_user_id: str = types.CharField(max_length=100)

    # Token data
    access_token: str = types.CharField(max_length=2000)
    refresh_token: str = types.CharField(max_length=2000, required=False)
    access_token_expires_at: datetime.datetime | None = types.DateTimeField(
        required=False, allow_null=True
    )
    refresh_token_expires_at: datetime.datetime | None = types.DateTimeField(
        required=False, allow_null=True
    )

    query: models.QuerySet[OAuthConnection] = models.QuerySet()

    model_options = models.Options(
        constraints=[
            models.UniqueConstraint(
                fields=["provider_key", "provider_user_id"],
                name="plainoauth_oauthconnection_unique_provider_key_user_id",
            )
        ],
        ordering=("provider_key",),
    )

    def __str__(self) -> str:
        return f"{self.provider_key}[{self.user}:{self.provider_user_id}]"

    def refresh_access_token(self) -> None:
        from .providers import OAuthToken, get_oauth_provider_instance

        provider_instance = get_oauth_provider_instance(provider_key=self.provider_key)
        oauth_token = OAuthToken(
            access_token=self.access_token,
            refresh_token=self.refresh_token,
            access_token_expires_at=self.access_token_expires_at,
            refresh_token_expires_at=self.refresh_token_expires_at,
        )
        refreshed_oauth_token = provider_instance.refresh_oauth_token(
            oauth_token=oauth_token
        )
        self.set_token_fields(refreshed_oauth_token)
        self.save()

    def set_token_fields(self, oauth_token: OAuthToken) -> None:
        self.access_token = oauth_token.access_token
        self.refresh_token = oauth_token.refresh_token
        self.access_token_expires_at = oauth_token.access_token_expires_at
        self.refresh_token_expires_at = oauth_token.refresh_token_expires_at

    def set_user_fields(self, oauth_user: OAuthUser) -> None:
        self.provider_user_id = oauth_user.provider_id

    def access_token_expired(self) -> bool:
        return (
            self.access_token_expires_at is not None
            and self.access_token_expires_at < timezone.now()
        )

    def refresh_token_expired(self) -> bool:
        return (
            self.refresh_token_expires_at is not None
            and self.refresh_token_expires_at < timezone.now()
        )

    @classmethod
    def get_or_create_user(
        cls, *, provider_key: str, oauth_token: OAuthToken, oauth_user: OAuthUser
    ) -> OAuthConnection:
        try:
            connection = cls.query.get(
                provider_key=provider_key,
                provider_user_id=oauth_user.provider_id,
            )
            connection.set_token_fields(oauth_token)
            connection.save()
            return connection
        except cls.DoesNotExist:
            with transaction.atomic():
                # If email needs to be unique, then we expect
                # that to be taken care of on the user model itself
                try:
                    user = get_user_model()(
                        **oauth_user.user_model_fields,
                    )
                    user.save()
                except (IntegrityError, ValidationError):
                    raise OAuthUserAlreadyExistsError()

                return cls.connect(
                    user=user,
                    provider_key=provider_key,
                    oauth_token=oauth_token,
                    oauth_user=oauth_user,
                )

    @classmethod
    def connect(
        cls,
        *,
        user: Any,
        provider_key: str,
        oauth_token: OAuthToken,
        oauth_user: OAuthUser,
    ) -> OAuthConnection:
        """
        Connect will either create a new connection or update an existing connection
        """
        try:
            connection = cls.query.get(
                user=user,
                provider_key=provider_key,
                provider_user_id=oauth_user.provider_id,
            )
        except cls.DoesNotExist:
            # Create our own instance (not using get_or_create)
            # so that any created signals contain the token fields too
            connection = cls(
                user=user,
                provider_key=provider_key,
                provider_user_id=oauth_user.provider_id,
            )

        connection.set_user_fields(oauth_user)
        connection.set_token_fields(oauth_token)
        connection.save()

        return connection
