import logging

from django.core.management.base import BaseCommand
from django.db import transaction
from uuid import uuid4
from epdb.models import Package, ReactionExplanation
from utilities.chem import FormatConverter
from django.utils import timezone


logger = logging.getLogger(__name__)


class Command(BaseCommand):
    def add_arguments(self, parser):
        parser.add_argument(
            "--rule-package",
            action="append",
            default=["32de3cf4-e3e6-4168-956e-32fa5ddb0ce1"],
            type=str,
            help="UUID to process. Can be specified multiple times.",
        )

        parser.add_argument(
            "--reaction-package",
            action="append",
            default=[
                "32de3cf4-e3e6-4168-956e-32fa5ddb0ce1",  # BBD
                "f05e38d8-e9b4-4c3e-b0d8-9ab29966eccf",  # Sediment
                "521c547a-fd2a-491c-ad5b-7eaa1577fb65",  # Sludge
                "5882df9c-dae1-4d80-a40e-db4724271456",  # Soil
                "87a49584-d937-482c-9c33-25928dcb02a8",  # PFAS
            ],
            type=str,
            help="UUID to process. Can be specified multiple times.",
        )

        parser.add_argument(
            "--dry-run",
            default=False,
            action="store_true",
            help="Perform dry run",
        )

    @transaction.atomic
    def handle(self, *args, **options):
        RUN_UUID = uuid4()
        RUN_START = timezone.now()

        rule_packages = Package.objects.filter(uuid__in=options["rule_package"])
        reaction_packages = Package.objects.filter(uuid__in=options["reaction_package"])

        rules = []
        for rule_package in rule_packages:
            rules.extend(rule_package.get_applicable_rules())

        reactions = []
        for reaction_package in reaction_packages:
            reactions.extend(reaction_package.reactions)

        logger.debug(f"Collected {len(rules)} rules and {len(reactions)} reactions.")

        for i, reaction in enumerate(reactions):
            logger.debug(f"Reaction {i} / {len(reactions)}")
            for j, rule in enumerate(rules):
                reactants, products = reaction.smirks().split(">>")

                if len(reactants.split(".")) > 1:
                    logger.debug(f"Skipping reaction {reaction.uuid} as it has multiple reactants.")
                    break

                products = products.split(".")

                # Run reaction with rule
                rule_products = rule.apply(reactants)

                # Check if products match (in both directions if extras are not allowed)
                for product_set in rule_products:
                    covered, exact = FormatConverter.smiles_covered_by(
                        products,
                        product_set.product_set,
                        standardize=True,
                        canonicalize_tautomers=True,
                        return_exact_match=True,
                    )

                    if covered and not options["dry-run"]:
                        logger.debug(f"Reaction {reaction.uuid} explained by rule {rule.uuid}")
                        re = ReactionExplanation()
                        re.run_uuid = RUN_UUID
                        re.run_start = RUN_START
                        re.reaction = reaction
                        re.rule = rule
                        re.exact = exact
                        re.save()
                        # Its explained, if there are more sets skip them
                        break