# Generated by Django 5.1.9 on 2025-06-04 01:11

import uuid

from django.db import migrations, models


def fix_malformed_uuids(apps, schema_editor):
    KhojUser = apps.get_model("database", "KhojUser")

    # Track UUID changes for automation cleanup
    uuid_mappings = {}

    # Handle null or empty user UUIDs
    for user in KhojUser.objects.filter(uuid__isnull=True):
        old_uuid = str(user.uuid) if user.uuid else "None"
        user.uuid = uuid.uuid4()
        user.save()
        uuid_mappings[old_uuid] = str(user.uuid)

    # Handle malformed user UUIDs
    for user in KhojUser.objects.all():
        current_uuid_val = user.uuid
        try:
            if not isinstance(current_uuid_val, uuid.UUID):
                # Attempt to parse it as UUID. This will catch "None", "null" strings or other malformed hex.
                uuid.UUID(str(current_uuid_val))
        except (ValueError, TypeError, AttributeError):
            old_uuid_str = str(current_uuid_val)
            new_uuid_obj = uuid.uuid4()
            user.uuid = new_uuid_obj
            user.save(
                update_fields=["uuid"]
            )  # Important to use update_fields to avoid triggering full save logic if not needed
            uuid_mappings[old_uuid_str] = str(new_uuid_obj)
            print(f"Fixed malformed UUID for user (old: '{old_uuid_str}', new: {str(new_uuid_obj)})")

    # Clean up orphaned automations
    cleanup_orphaned_automations(uuid_mappings)


def cleanup_orphaned_automations(uuid_mappings):
    """Remove automations with malformed UUIDs in job_ids"""
    from apscheduler.jobstores.base import JobLookupError

    from khoj.utils import state

    if not state.scheduler:
        return

    all_jobs = state.scheduler.get_jobs()
    removed_orphaned_count = 0
    removed_malformed_count = 0

    for job in all_jobs:
        if job.id.startswith("automation_"):
            # Extract UUID from job_id: "automation_{uuid}_{query_id}"
            job_parts = job.id.split("_", 2)
            if len(job_parts) >= 2:
                job_uuid = job_parts[1]

                # Check if this UUID was malformed
                if job_uuid in uuid_mappings:
                    # Remove orphaned automation
                    try:
                        state.scheduler.remove_job(job.id)
                        removed_orphaned_count += 1
                        print(f"Removed orphaned automation: {job.id}")
                    except JobLookupError:
                        pass  # Job already removed

                # Also remove jobs with clearly malformed UUIDs
                elif job_uuid in ["None", "null"] or not is_valid_uuid(job_uuid):
                    try:
                        state.scheduler.remove_job(job.id)
                        removed_malformed_count += 1
                        print(f"Removed automation with malformed UUID: {job.id}")
                    except JobLookupError:
                        pass

    if removed_orphaned_count > 0 or removed_malformed_count > 0:
        print(f"Removed {removed_orphaned_count} orphaned and {removed_malformed_count} malformed automations.")


def is_valid_uuid(uuid_string):
    """Check if string is a valid UUID"""
    try:
        uuid.UUID(str(uuid_string))
        return True
    except (ValueError, TypeError, AttributeError):
        return False


class Migration(migrations.Migration):
    dependencies = [
        ("database", "0089_chatmodel_price_tier_and_more"),
    ]

    operations = [
        migrations.RunPython(fix_malformed_uuids, reverse_code=migrations.RunPython.noop),
        migrations.AlterField(
            model_name="khojuser",
            name="uuid",
            field=models.UUIDField(default=uuid.uuid4, editable=False, unique=True),
        ),
    ]
