# Generated by Django 5.2 on 2025-05-11 18:50

import random
import string

from django.db import migrations


def generate_base62_uid():
    """Generate a 12-character base62 UID."""
    charset = string.ascii_uppercase + string.ascii_lowercase + string.digits
    return "".join(random.choice(charset) for _ in range(12))


def copy_params_to_features(apps, schema_editor):
    """Copy data from Param to Feature with generated UIDs."""
    Param = apps.get_model("lamindb", "Param")
    Feature = apps.get_model("lamindb", "Feature")

    # Create features from params with generated UIDs
    name_to_feature = {}

    # First pass: create all Features
    for param in Param.objects.all():
        feature = Feature.objects.create(
            uid=generate_base62_uid(),
            name=param.name,
            dtype=param.dtype,
            is_type=param.is_type,
            _expect_many=param._expect_many,
            created_at=param.created_at,
            updated_at=param.updated_at,
            created_by_id=param.created_by_id,
            run_id=param.run_id,
        )
        name_to_feature[param.name] = feature

    # Second pass: set up type relationships
    for param in Param.objects.all():
        if param.type_id is not None:
            try:
                # Find the parameter that is the type
                type_param = Param.objects.get(id=param.type_id)

                # Set the type relationship
                feature = name_to_feature[param.name]
                feature.type = name_to_feature[type_param.name]
                feature.save()
            except (Param.DoesNotExist, KeyError):
                # Handle case where the referenced type doesn't exist
                print(f"Warning: Could not set type for feature {param.name}")


def copy_param_values(apps, schema_editor):
    """Copy data from ParamValue to FeatureValue."""
    Param = apps.get_model("lamindb", "Param")
    Feature = apps.get_model("lamindb", "Feature")
    ParamValue = apps.get_model("lamindb", "ParamValue")
    FeatureValue = apps.get_model("lamindb", "FeatureValue")

    # Create a mapping from param ID to feature ID
    param_to_feature = {
        param.id: Feature.objects.get(name=param.name).id
        for param in Param.objects.all()
    }

    # Create FeatureValues from ParamValues
    pv_to_fv = {}

    for pv in ParamValue.objects.all():
        try:
            feature_id = param_to_feature[pv.param_id]

            # Create the FeatureValue
            fv = FeatureValue.objects.create(
                feature_id=feature_id,
                value=pv.value,
                hash=pv.hash,
                created_at=pv.created_at,
                created_by_id=pv.created_by_id,
            )

            # Remember the mapping
            pv_to_fv[pv.id] = fv.id

        except KeyError:
            print(f"Warning: Could not map ParamValue {pv.id} to a FeatureValue")

    return pv_to_fv


def copy_run_param_values(apps, schema_editor):
    """Copy data from RunParamValue to RunFeatureValue."""
    ParamValue = apps.get_model("lamindb", "ParamValue")
    FeatureValue = apps.get_model("lamindb", "FeatureValue")
    RunParamValue = apps.get_model("lamindb", "RunParamValue")
    RunFeatureValue = apps.get_model("lamindb", "RunFeatureValue")

    # Create a mapping from ParamValue ID to FeatureValue ID
    # This is complex because we need to match based on param->feature and value
    pv_to_fv = {}

    for pv in ParamValue.objects.all():
        try:
            # Find the corresponding Feature for this Param
            feature = apps.get_model("lamindb", "Feature").objects.get(
                name=apps.get_model("lamindb", "Param").objects.get(id=pv.param_id).name
            )

            # Find the corresponding FeatureValue
            if pv.hash is None:
                fv = FeatureValue.objects.get(
                    feature=feature, value=pv.value, hash__isnull=True
                )
            else:
                fv = FeatureValue.objects.get(feature=feature, hash=pv.hash)

            pv_to_fv[pv.id] = fv.id
        except Exception as e:
            print(f"Warning: Could not map ParamValue {pv.id} to a FeatureValue: {e}")

    # Copy RunParamValues to RunFeatureValues
    for rpv in RunParamValue.objects.all():
        try:
            featurevalue_id = pv_to_fv[rpv.paramvalue_id]

            RunFeatureValue.objects.create(
                run_id=rpv.run_id,
                featurevalue_id=featurevalue_id,
                created_at=rpv.created_at,
                created_by_id=rpv.created_by_id,
            )
        except KeyError:
            print(f"Warning: Could not map RunParamValue {rpv.id} to a RunFeatureValue")


def copy_artifact_param_values(apps, schema_editor):
    """Copy data from ArtifactParamValue to ArtifactFeatureValue."""
    ParamValue = apps.get_model("lamindb", "ParamValue")
    FeatureValue = apps.get_model("lamindb", "FeatureValue")
    ArtifactParamValue = apps.get_model("lamindb", "ArtifactParamValue")
    ArtifactFeatureValue = apps.get_model("lamindb", "ArtifactFeatureValue")

    # Create a mapping from ParamValue ID to FeatureValue ID
    # This is complex because we need to match based on param->feature and value
    pv_to_fv = {}

    for pv in ParamValue.objects.all():
        try:
            # Find the corresponding Feature for this Param
            feature = apps.get_model("lamindb", "Feature").objects.get(
                name=apps.get_model("lamindb", "Param").objects.get(id=pv.param_id).name
            )

            # Find the corresponding FeatureValue
            if pv.hash is None:
                fv = FeatureValue.objects.get(
                    feature=feature, value=pv.value, hash__isnull=True
                )
            else:
                fv = FeatureValue.objects.get(feature=feature, hash=pv.hash)

            pv_to_fv[pv.id] = fv.id
        except Exception as e:
            print(f"Warning: Could not map ParamValue {pv.id} to a FeatureValue: {e}")

    # Copy ArtifactParamValues to ArtifactFeatureValues
    for apv in ArtifactParamValue.objects.all():
        try:
            featurevalue_id = pv_to_fv[apv.paramvalue_id]

            ArtifactFeatureValue.objects.create(
                artifact_id=apv.artifact_id,
                featurevalue_id=featurevalue_id,
                created_at=apv.created_at,
                created_by_id=apv.created_by_id,
                run_id=apv.run_id,
            )
        except KeyError:
            print(
                f"Warning: Could not map ArtifactParamValue {apv.id} to an ArtifactFeatureValue"
            )


def get_features_count(apps, schema_editor):
    Feature = apps.get_model("lamindb", "Feature")
    FeatureValue = apps.get_model("lamindb", "FeatureValue")
    RunFeatureValue = apps.get_model("lamindb", "RunFeatureValue")
    ArtifactFeatureValue = apps.get_model("lamindb", "ArtifactFeatureValue")

    print("Features are now")
    print(f"Feature count: {Feature.objects.count()}")
    print(f"FeatureValue count: {FeatureValue.objects.count()}")
    print(f"RunFeatureValue count: {RunFeatureValue.objects.count()}")
    print(f"ArtifactFeatureValue count: {ArtifactFeatureValue.objects.count()}")


def get_params_count(apps, schema_editor):
    Param = apps.get_model("lamindb", "Param")
    ParamValue = apps.get_model("lamindb", "ParamValue")
    RunParamValue = apps.get_model("lamindb", "RunParamValue")
    ArtifactParamValue = apps.get_model("lamindb", "ArtifactParamValue")

    print("Moving params into features and deleting the params table")
    print(f"Param count: {Param.objects.count()}")
    print(f"ParamValue count: {ParamValue.objects.count()}")
    print(f"RunParamValue count: {RunParamValue.objects.count()}")
    print(f"ArtifactParamValue count: {ArtifactParamValue.objects.count()}")


class Migration(migrations.Migration):
    dependencies = [
        ("lamindb", "0095_remove_rundata_flextable"),
    ]

    operations = [
        # Step 0: Report the results
        migrations.RunPython(get_params_count),
        migrations.RunPython(get_features_count),
        # Step 1: Copy from Param to Feature
        migrations.RunPython(copy_params_to_features),
        # Step 2: Copy from ParamValue to FeatureValue
        migrations.RunPython(copy_param_values),
        # Step 3: Copy from RunParamValue to RunFeatureValue
        migrations.RunPython(copy_run_param_values),
        # Step 4: Copy from ArtifactParamValue to ArtifactFeatureValue
        migrations.RunPython(copy_artifact_param_values),
        # Step 5: Report the results
        migrations.RunPython(get_features_count),
        migrations.RemoveField(
            model_name="artifact",
            name="_param_values",
        ),
        migrations.RemoveField(
            model_name="param",
            name="created_by",
        ),
        migrations.RemoveField(
            model_name="param",
            name="run",
        ),
        migrations.RemoveField(
            model_name="param",
            name="schemas",
        ),
        migrations.RemoveField(
            model_name="param",
            name="space",
        ),
        migrations.RemoveField(
            model_name="param",
            name="type",
        ),
        migrations.RemoveField(
            model_name="run",
            name="_param_values",
        ),
        migrations.AlterUniqueTogether(
            name="schemaparam",
            unique_together=None,
        ),
        migrations.RemoveField(
            model_name="schemaparam",
            name="schema",
        ),
    ]
