import os

from django.db import migrations
from django.core.management import call_command


def load_data(apps, __):
    TreatmentStatus = apps.get_model("archaeological_finds", "treatmentstatus")
    if not TreatmentStatus.objects.count():
        migration = "0137_data_migration_treatmentstatus.json"
        json_path = os.sep.join(os.path.abspath(__file__).split(os.sep)[:-1] + [migration])
        call_command("loaddata", json_path)

    TreatmentInputStatus = apps.get_model("archaeological_finds", "treatmentinputstatus")
    if not TreatmentInputStatus.objects.count():
        migration = "0137_data_migration_treatmentinputstatus.json"
        json_path = os.sep.join(os.path.abspath(__file__).split(os.sep)[:-1] + [migration])
        call_command("loaddata", json_path)

    input_status_executed = TreatmentInputStatus.objects.filter(executed=True).all()[0]
    input_status_draft = TreatmentInputStatus.objects.filter(executed=False).all()[0]

    Treatment = apps.get_model("archaeological_finds", "treatment")
    try:
        TreatmentState = apps.get_model("archaeological_finds", "treatmentstate")
    except LookupError:
        # treatment state model deleted
        return
    for treatment_state in TreatmentState.objects.all():
        q = TreatmentStatus.objects.filter(txt_idx=treatment_state.txt_idx)
        if not q.exists():
            print(f"Equivalence for treatment state: '{treatment_state.txt_idx}' not found")
            continue
        equiv = q.all()[0]
        q = Treatment.objects.filter(treatment_state=treatment_state)
        q.update(treatment_status=equiv)
        if treatment_state.executed:
            q.update(input_status=input_status_executed)
        else:
            q.update(input_status=input_status_draft)


class Migration(migrations.Migration):

    dependencies = [
        ('archaeological_finds', '0136_treatment_status_refactoring'),
    ]

    operations = [
        migrations.RunPython(load_data)
    ]