from csv import DictReader

from django.core.management.base import BaseCommand

from epdb.models import *


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
                )