"""Interface to RDT (Reaction Decoder Tool)."""

from __future__ import annotations

import contextlib
import os
import subprocess
from subprocess import CompletedProcess
from typing import Optional

import rdkit.Chem as rdk
import rdkit.Chem.rdChemReactions as rdk_r
from chemrecon import ChemReaction
from chemrecon.chem.mol import Mol, MolInstance

from aam_annotator.utils.cache import cache

# RDT Jar location
rdt_jar_location = "./src/redistribute/rdt/"
rdt_jar_filename = "rdt-2.4.1-jar-with-dependencies.jar"
# TODO should be able to pass this as argument


# RDT interface
# ----------------------------------------------------------------------------------------------------------------------
def produce_aam(
    lhs: list[Mol], rhs: list[Mol], timeout: int = 30
) -> Optional[ChemReaction]:
    # Check
    if any(isinstance(m, MolInstance) for m in [*lhs, *rhs]):
        raise ValueError("Cannot use MolInstance in RDT!")

    # Establish reaction S_SMILES query
    smiles_lhs = ".".join(sorted([temp.to_smiles() for temp in lhs]))
    smiles_rhs = ".".join(sorted([temp.to_smiles() for temp in rhs]))

    smiles_query: str = smiles_lhs + ">>" + smiles_rhs

    # If trivial (LHS, RHS identical)
    if (lhs == rhs) and (smiles_lhs == smiles_rhs):
        print("  -> Trivial aam.")
        trivial_map = trivial_aam(lhs)
        return trivial_map

    aam_str: Optional[str] = run_rdt(
        smiles_query, timeout=timeout, cache_key=smiles_query
    )
    if aam_str is None:
        # No result found
        print("    -> No result.")
        return None

    # Convert to a chemreaction and return
    chemreaction = ChemReaction(rdk_r.ReactionFromSmarts(aam_str))
    return chemreaction


# noinspection PyUnusedLocal
# (decorator uses cache_key)
@cache("cache_rdt.shelve")
def run_rdt(smiles_query: str, timeout: int, cache_key: str) -> Optional[str]:
    """ """
    print(f"    -> Running RDT with query: {smiles_query} ...")
    with cd_temp(rdt_jar_location):
        try:
            result: CompletedProcess = subprocess.run(
                [
                    "java",
                    "-jar",
                    rdt_jar_filename,
                    # Arguments
                    "-Q",
                    "SMI",  # Input is smiles
                    "-q",
                    f'"{smiles_query}"',  # Input query
                    "-j",
                    "AAM",  # Perform atom-to-atom mapping
                    "-f",
                    "TEXT",  # Output as text
                ],
                stdout=subprocess.DEVNULL,  # Do not print
                timeout=timeout,
            )
            print(f"      -> Done ({result.returncode}).")
        except subprocess.TimeoutExpired:
            print("      -> Timeout.")
            return None

        # Look at the output and delete the file
        try:
            with open("ECBLAST_smiles_AAM.txt", "r") as out_file:
                outstr = out_file.readlines()
                for i, l in enumerate(outstr):
                    if l.startswith("SELECTED AAM MAPPING"):
                        # Next line is correct
                        aam_str = outstr[i + 1].strip()
                        break

                print(f"   -> Outstr: {aam_str}")

                # Remove the file again
                os.remove("ECBLAST_smiles_AAM.txt")
                os.remove("ECBLAST_smiles_AAM.rxn")
                return aam_str

        except FileNotFoundError as e:
            # Could not read generated ECBLAST
            print(f"   -> Could not read generated ECBLAST: {e}")
            return None


def trivial_aam(mols: list[Mol]) -> ChemReaction:
    """Produce a trivial atom-to-atom mapping with identical LHS and RHS."""

    def with_maps(
        mol: rdk.Mol, start_map: int = 1
    ) -> tuple[rdk.Mol, dict[int, int], int]:
        m = rdk.Mol(mol)
        aam = {}
        n = start_map
        for a in m.GetAtoms():
            a.SetAtomMapNum(n)
            aam[n] = a.GetIdx()
            n += 1
        return m, aam, n

    def copy_maps_from(src_mol, dst_mol, aam):
        for a in dst_mol.GetAtoms():
            a.SetAtomMapNum(aam[a.GetIdx()])

    # Run
    rdk_reaction = rdk_r.ChemicalReaction()

    next_map_num = 1
    for mol in mols:
        mapped_mol, aam, next_map_num = with_maps(mol.mol, next_map_num)
        rdk_reaction.AddReactantTemplate(mapped_mol)
        rdk_reaction.AddProductTemplate(mapped_mol)

    rdk_reaction.Initialize()
    return ChemReaction(rdk_reaction)


# Misc
# ----------------------------------------------------------------------------------------------------------------------
# Context to change working directory
@contextlib.contextmanager
def cd_temp(path: str):
    olddir = os.getcwd()
    os.chdir(path)
    try:
        yield
    finally:
        os.chdir(olddir)
