from csv import DictReader

from django.core.management.base import BaseCommand
from django.db import transaction

from epdb.models import Compound, CompoundStructure, Reaction, ExternalDatabase, ExternalIdentifier


class Command(BaseCommand):
    STR_TO_MODEL = {
        "Compound": Compound,
        "CompoundStructure": CompoundStructure,
        "Reaction": Reaction,
    }

    STR_TO_DATABASE = {
        "ChEBI": ExternalDatabase.objects.get(name="ChEBI"),
        "RHEA": ExternalDatabase.objects.get(name="RHEA"),
        "KEGG Reaction": ExternalDatabase.objects.get(name="KEGG Reaction"),
        "PubChem Compound": ExternalDatabase.objects.get(name="PubChem Compound"),
        "PubChem Substance": ExternalDatabase.objects.get(name="PubChem Substance"),
    }

    def add_arguments(self, parser):
        parser.add_argument(
            "--data",
            type=str,
            help="Path of the ID Mapping file.",
            required=True,
        )

        parser.add_argument(
            "--replace-host",
            type=str,
            help="Replace https://envipath.org/ with this host, e.g. http://localhost:8000/",
        )

    @transaction.atomic
    def handle(self, *args, **options):
        with open(options["data"]) as fh:
            reader = DictReader(fh)
            for row in reader:
                clz = self.STR_TO_MODEL[row["model"]]

                url = row["url"]
                if options["replace_host"]:
                    url = url.replace("https://envipath.org/", options["replace_host"])

                instance = clz.objects.get(url=url)
                db = self.STR_TO_DATABASE[row["identifier_type"]]

                ExternalIdentifier.objects.create(
                    content_object=instance,
                    database=db,
                    identifier_value=row["identifier_value"],
                    url=db.url_pattern.format(id=row["identifier_value"]),
                    is_primary=False,
                )