import ast
import os
import re
from time import sleep

from dotenv import load_dotenv
from inspect_ai.model import GenerateConfig, get_model
from rxn4chemistry import RXN4ChemistryWrapper

LLM_INSTRUCTIONS = \
("Here is a chemical synthesis described as a json.\nYour task is "
"to describe the synthesis, as if you were giving instructions for"
"a recipe. Use only the substances, quantities, temperatures and "
"in general any action mentioned in the json file. This is your "
"only source of information, do not make up anything else. Also, "
"add 15mL of DCM as a solvent in the first step. If you ever need "
'to refer to the json file, refer to it as "(by) the tool". '
"However avoid references to it. \nFor this task, give as many "
"details as possible.\n {json}")


def retry(times: int, exceptions, sleep_time: int = 5):
    """
    Retry Decorator.

    Retries the wrapped function/method `times` times if the exceptions
    listed in ``exceptions`` are thrown
    :param times: The number of times to repeat the wrapped function/method
    :type times: Int
    :param Exceptions: Lists of exceptions that trigger a retry attempt
    :type Exceptions: Tuple of Exceptions
    """
    def decorator(func):
        def newfn(*args, **kwargs):
            attempt = 0
            while attempt < times:
                try:
                    sleep(sleep_time)
                    return func(*args, **kwargs)
                except exceptions:
                    print(
                        "Exception thrown when attempting to run %s, "
                        "attempt %d of %d" % (func, attempt, times)
                    )
                    attempt += 1
            return func(*args, **kwargs)

        return newfn

    return decorator



def get_RXN4ChemistryWrapper() -> RXN4ChemistryWrapper:
    load_dotenv()

    rxn_api_key = os.getenv("RXN4CHEMISTRY_API_KEY")
    rxn_project_id = os.getenv("RXN4CHEMISTRY_PROJECT_ID")
    rxn = RXN4ChemistryWrapper(api_key=rxn_api_key)
    rxn.set_project(rxn_project_id)

    return rxn


class RXNRetroSynthesis:
    """Class for predicting retrosynthesis using RXN4Chemistry"""
    def __init__(self, summarize_model_name: str, sleep_time: int = 5):
        self.rxn4chem = get_RXN4ChemistryWrapper()
        self.sleep_time = sleep_time
        self.summarize_model = get_model(
            summarize_model_name,
            config=GenerateConfig(temperature=0.05, max_tokens=2000)
        )

    async def run(self, target_smiles: str):
        prediction_id = self.predict_retrosynthesis(target_smiles)
        paths = self.get_paths(prediction_id)
        procedure = await self.get_action_sequence(paths[0])

        return procedure

    @retry(10, KeyError)
    def predict_retrosynthesis(self, target_smiles: str):
        response = self.rxn4chem.predict_automatic_retrosynthesis(
            product=target_smiles,
            # These are default values in RXN4Chemistry as of 21st March 2025
            fap=0.6,
            max_steps=3,
            nbeams=10,
            pruning_steps=2,
            # This is the AI Model used by ChemCrow
            ai_model="12class-tokens-2021-05-14"
        )

        return response["prediction_id"]

    @retry(20, KeyError)
    def get_paths(self, prediction_id):
        results = self.rxn4chem.get_predict_automatic_retrosynthesis_results(prediction_id)
        if "retrosynthetic_paths" not in results.keys():
            raise KeyError
        paths = results["retrosynthetic_paths"]
        if results["status"] != "SUCCESS":
            sleep(self.sleep_time * 2)
            raise KeyError
        if paths is not None:
            if len(paths) > 0:
                return paths
        raise KeyError

    async def get_action_sequence(self, path):
        """Get sequence of actions."""
        response = self.synth_from_sequence(path["sequenceId"])
        if "synthesis_id" not in response.keys():
            return path

        synthesis_id = response["synthesis_id"]
        nodeids = self.get_node_ids(synthesis_id)
        if nodeids is None:
            return "Tool error"

        # Attempt to get actions for each node + product information
        real_nodes = []
        actions_and_products = []
        for node in nodeids:
            node_resp = self.get_reaction_settings(
                synthesis_id=synthesis_id, node_id=node
            )
            if "actions" in node_resp.keys():
                real_nodes.append(node)
                actions_and_products.append(node_resp)

        json_actions = self._preproc_actions(actions_and_products)
        llm_sum = await self._summary_llm(json_actions)
        return llm_sum

    @retry(20, KeyError)
    def synth_from_sequence(self, sequence_id: str) -> str:
        response = self.rxn4chem.create_synthesis_from_sequence(sequence_id=sequence_id)
        if "synthesis_id" in response.keys():
            return response
        raise KeyError

    @retry(20, KeyError)
    def get_node_ids(self, synthesis_id: str):
        response = self.rxn4chem.get_node_ids(synthesis_id=synthesis_id)
        if isinstance(response, list) and len(response) > 0:
            return response
        raise KeyError

    @retry(20, KeyError)
    def get_reaction_settings(self, synthesis_id: str, node_id: str):
        response = self.rxn4chem.get_reaction_settings(
            synthesis_id=synthesis_id, node_id=node_id
        )
        if "actions" in response.keys():
            return response
        elif "response" in response.keys():
            if "error" in response["response"].keys():
                if response["response"]["error"] == "Too Many Requests":
                    sleep(self.sleep_time * 2)
                    raise KeyError
            return response
        raise KeyError

    def _preproc_actions(self, actions_and_products):
        """Preprocess actions."""
        json_actions = {"number_of_steps": len(actions_and_products)}

        for i, actn in enumerate(actions_and_products):
            json_actions[f"Step_{i}"] = {}
            json_actions[f"Step_{i}"]["actions"] = actn["actions"]
            json_actions[f"Step_{i}"]["product"] = actn["product"]

        # Clean actions to use less tokens: Remove False, None, ''
        clean_act_str = re.sub(
            r"\'[A-Za-z]+\': (None|False|\'\'),? ?", "", str(json_actions)
        )
        json_actions = ast.literal_eval(clean_act_str)

        return json_actions

    async def _summary_llm(self, json: dict) -> str:
        """Describe synthesis in a more natural language-friendly way."""
        model_input = LLM_INSTRUCTIONS.format(json=json)
        model_output = await self.summarize_model.generate(input=model_input, cache=True)
        return model_output.completion

