from django.apps import apps
from django.conf import settings as s
from django.core.management.base import BaseCommand
from django.db.models import F, JSONField, TextField, Value
from django.db.models.functions import Cast, Replace

from epdb.models import EnviPathModel


class Command(BaseCommand):
    def add_arguments(self, parser):
        parser.add_argument(
            "--old",
            type=str,
            help="Old Host, most likely https://envipath.org/",
            required=True,
        )
        parser.add_argument(
            "--new",
            type=str,
            help="New Host, most likely http://localhost:8000/",
            required=True,
        )

    def handle(self, *args, **options):
        Package = s.GET_PACKAGE_MODEL()
        print("Localizing urls for Package")
        Package.objects.update(url=Replace(F("url"), Value(options["old"]), Value(options["new"])))

        MODELS = [
            "User",
            "Group",
            "Compound",
            "CompoundStructure",
            "Pathway",
            "Edge",
            "Node",
            "Reaction",
            "SimpleAmbitRule",
            "SimpleRDKitRule",
            "ParallelRule",
            "SequentialRule",
            "Scenario",
            "Setting",
            "MLRelativeReasoning",
            "RuleBasedRelativeReasoning",
            "EnviFormer",
            "ApplicabilityDomain",
            "EnzymeLink",
        ]
        for model in MODELS:
            obj_cls = apps.get_model("epdb", model)
            print(f"Localizing urls for {model}")
            obj_cls.objects.update(
                url=Replace(F("url"), Value(options["old"]), Value(options["new"]))
            )
            if issubclass(obj_cls, EnviPathModel):
                obj_cls.objects.update(
                    kv=Cast(
                        Replace(
                            Cast(F("kv"), output_field=TextField()),
                            Value(options["old"]),
                            Value(options["new"]),
                        ),
                        output_field=JSONField(),
                    )
                )