from django.contrib.auth.models import AbstractUser
from django.db import models
from django.core.exceptions import ValidationError

class CustomUser(AbstractUser):
    default_friend_code = models.ForeignKey("FriendCode", on_delete=models.SET_NULL, null=True, blank=True)

    def __str__(self):
        return self.email

    def set_default_friend_code(self, friend_code):
        """Set a friend code as default if it belongs to the user."""
        if friend_code.user != self:
            raise ValidationError("Friend code does not belong to this user.")
        self.default_friend_code = friend_code
        self.save(update_fields=["default_friend_code"])

    def remove_default_friend_code(self, friend_code):
        """
        If the given friend code is the current default,
        assign another of the user's friend codes (if any) as default.
        """
        if self.default_friend_code == friend_code:
            other_codes = self.friend_codes.exclude(pk=friend_code.pk)
            self.default_friend_code = other_codes.first() if other_codes.exists() else None
            self.save(update_fields=["default_friend_code"])

class FriendCode(models.Model):
    friend_code = models.CharField(max_length=19)
    user = models.ForeignKey(CustomUser, on_delete=models.PROTECT, related_name='friend_codes')
    created_at = models.DateTimeField(auto_now_add=True)
    updated_at = models.DateTimeField(auto_now=True)

    def save(self, *args, **kwargs):
        """
        When a new friend code is saved,
        if the user has no default friend code yet,
        automatically set this as the default.
        """
        is_new = self.pk is None
        super().save(*args, **kwargs)
        if is_new and not self.user.default_friend_code:
            self.user.default_friend_code = self
            self.user.save(update_fields=["default_friend_code"])

    def __str__(self):
        return self.friend_code