from django import forms
from django.contrib.auth.forms import UserCreationForm, AuthenticationForm
from .models import User, Profile

class CompanySignUpForm(UserCreationForm):
    class Meta(UserCreationForm.Meta):
        model = User
        fields = ('username', 'email')

    def save(self, commit=True):
        user = super().save(commit=False)
        user.role = User.Role.COMPANY
        if commit:
            user.save()
            Profile.objects.create(user=user)
        return user

class InfluencerSignUpForm(UserCreationForm):
    class Meta(UserCreationForm.Meta):
        model = User
        fields = ('username', 'email')

    def save(self, commit=True):
        user = super().save(commit=False)
        user.role = User.Role.INFLUENCER
        if commit:
            user.save()
            Profile.objects.create(user=user)
        return user

class LoginForm(AuthenticationForm):
    pass