From 63b0e68fc2f7f943a7aa0b7600616da14e275f69 Mon Sep 17 00:00:00 2001
From: Seyed Alireza Hashemi <amoo.hashem.mehraban@gmail.com>
Date: Wed, 13 Nov 2024 16:08:21 +0330
Subject: [PATCH] feat: add origin for uuid-logins from shad

---
 apps/accounts/migrations/0044_user_origin.py | 18 +++++++++
 apps/accounts/models.py                      |  1 +
 apps/accounts/serializers/user_serializer.py |  2 +-
 apps/accounts/views/uuid_login.py            | 42 +++++++++++++++++++-
 4 files changed, 61 insertions(+), 2 deletions(-)
 create mode 100644 apps/accounts/migrations/0044_user_origin.py

diff --git a/apps/accounts/migrations/0044_user_origin.py b/apps/accounts/migrations/0044_user_origin.py
new file mode 100644
index 00000000..96fc10c9
--- /dev/null
+++ b/apps/accounts/migrations/0044_user_origin.py
@@ -0,0 +1,18 @@
+# Generated by Django 4.1.3 on 2024-11-13 12:33
+
+from django.db import migrations, models
+
+
+class Migration(migrations.Migration):
+
+    dependencies = [
+        ('accounts', '0043_rename_is_artificial_user_is_temporary'),
+    ]
+
+    operations = [
+        migrations.AddField(
+            model_name='user',
+            name='origin',
+            field=models.CharField(blank=True, max_length=50, null=True),
+        ),
+    ]
diff --git a/apps/accounts/models.py b/apps/accounts/models.py
index 7d1adf33..b250b71e 100644
--- a/apps/accounts/models.py
+++ b/apps/accounts/models.py
@@ -37,6 +37,7 @@ class Gender(models.TextChoices):
     city = models.CharField(max_length=50, null=True, blank=True)
     postal_code = models.CharField(max_length=10, null=True, blank=True)
     is_temporary = models.BooleanField(default=False)
+    origin = models.CharField(max_length=50, blank=True, null=True)
 
     def get_user_website(self, website):
         try:
diff --git a/apps/accounts/serializers/user_serializer.py b/apps/accounts/serializers/user_serializer.py
index e4c38858..0a8d3f5d 100644
--- a/apps/accounts/serializers/user_serializer.py
+++ b/apps/accounts/serializers/user_serializer.py
@@ -96,7 +96,7 @@ def update(self, instance, validated_data):
     class Meta:
         model = User
         fields = ['id', 'phone_number', 'first_name', 'last_name',
-                  'password', 'username', 'email', 'is_temporary']
+                  'password', 'username', 'email', 'is_temporary', 'origin']
         read_only_fields = ['id']
 
 
diff --git a/apps/accounts/views/uuid_login.py b/apps/accounts/views/uuid_login.py
index a09ba5b8..a61e35fd 100644
--- a/apps/accounts/views/uuid_login.py
+++ b/apps/accounts/views/uuid_login.py
@@ -1,12 +1,19 @@
 import uuid
+from django.db import transaction
 from drf_yasg.utils import swagger_auto_schema
 from rest_framework import status
 from rest_framework.response import Response
 
 from apps.accounts.serializers.user_serializer import UserSerializer
+from apps.accounts.utils import create_or_get_user, find_user_in_website
 from apps.accounts.views.base_login import BaseLoginView
 
 
+import uuid
+from rest_framework import status
+from rest_framework.response import Response
+
+
 class UserIDLoginView(BaseLoginView):
     user_data_field = 'id'
 
@@ -20,6 +27,7 @@ class UserIDLoginView(BaseLoginView):
     )
     def post(self, request):
         user_id = request.data.get("user_id")
+        origin = request.data.get("origin", "")
 
         # Check if user_id is provided
         if not user_id:
@@ -29,7 +37,8 @@ def post(self, request):
         if not self.is_valid_uuid(user_id):
             return Response({"error": "Invalid UserID format. Must be a valid UUID."}, status=status.HTTP_400_BAD_REQUEST)
 
-        return self.handle_post(request)
+        # Pass the origin to handle_post
+        return self.handle_post(request, origin=origin)
 
     def get_user_identifier(self, request):
         return request.data.get("user_id")
@@ -40,3 +49,34 @@ def is_valid_uuid(self, value):
             return True
         except ValueError:
             return False
+
+    @transaction.atomic
+    def handle_post(self, request, origin=""):
+        user_identifier = self.get_user_identifier(request)
+        website = request.headers.get("Website")
+
+        try:
+            user = find_user_in_website(
+                user_data={self.user_data_field: user_identifier},
+                website=website,
+                raise_exception=True
+            )
+            created = False
+            response_status = status.HTTP_200_OK
+        except:
+            user_data = self.get_user_data(user_identifier)
+            user, created = create_or_get_user(user_data, website=website)
+
+            if created:
+                user.set_unusable_password()
+                user.origin = origin  # Save origin for new users
+                user.save()
+            response_status = status.HTTP_201_CREATED
+
+        # Update origin for existing users
+        if not created and origin:
+            user.origin = origin
+            user.save(update_fields=["origin"])
+
+        self.create_login_event(user, website)
+        return self.generate_response(user, created, response_status)