# Copyright (c) 2023 ETH Zurich.
#                    All rights reserved.
#
# Use of this source code is governed by a BSD-style license that can be
# found in the LICENSE file.
#
# main author: Nils Blach

import json
import logging
from typing import List, Optional, Dict, Any
from escargot.language_models import AbstractLanguageModel
from escargot.operations import GraphOfOperations, Thought
from escargot.prompter import ESCARGOTPrompter
from escargot.parser import ESCARGOTParser
from escargot.coder import Coder
import copy

class Controller:
    """
    Controller class to manage the execution flow of the Graph of Operations,
    generating the Graph Reasoning State.
    This involves language models, graph operations, prompting, and parsing.
    """

    def __init__(
        self,
        lm: AbstractLanguageModel,
        graph: GraphOfOperations,
        prompter: ESCARGOTPrompter,
        parser: ESCARGOTParser,
        logger: logging.Logger,
        problem_parameters: Dict[str, Any],
    ) -> None:
        """
        Initialize the Controller instance with the language model,
        operations graph, prompter, parser, and problem parameters.
        """
        self.logger = logger
        self.lm = lm
        self.graph = graph
        self.prompter = prompter
        self.parser = parser
        self.original_problem_parameters = problem_parameters
        self.problem_parameters = problem_parameters
        self.run_executed = False
        self.got_steps = {}
        self.final_thought = None
        self.execution_queue = []
        self.max_run_tries = 3
        self.max_operation_tries = 2
        self.coder = Coder()

    def initialize_execution_queue(self) -> None:
        """
        Initialize the execution queue with the root operations of the graph.
        """
        self.execution_queue = [
            operation for operation in self.graph.operations if operation.can_be_executed()
        ]

    def execute_step(self) -> Optional[Thought]:
        """
        Execute one step from the execution queue.

        :return: The thought generated by the executed operation, or None if no operation is left to execute.
        """
        if not self.execution_queue:
            self.initialize_execution_queue()

        if not self.execution_queue:
            self.logger.debug("No more operations to execute.")
            return None

        current_operation = self.execution_queue.pop(0)
        current_operation_backup = copy.copy(current_operation)
        
        tries = 0
        while tries < self.max_operation_tries:
            try:
                current_operation.execute(
                    self.lm, self.prompter, self.parser, self.got_steps, self.logger, self.coder, **self.problem_parameters
                )
                break
            except Exception as e:
                self.logger.error("Error executing operation %s: %s", current_operation.operation_type, e)
                current_operation = copy.copy(current_operation_backup)
                tries += 1

        del current_operation_backup

        if tries == self.max_operation_tries:
            self.logger.error("Max tries reached on executing operation %s", current_operation.operation_type)
            return None
        
        for operation in current_operation.successors:
            if operation.can_be_executed():
                self.execution_queue.append(operation)
        
        self.final_thought = current_operation.get_thoughts()[0]
        
        return self.final_thought

    def run(self) -> None:
        """
        Run the controller and execute the operations from the Graph of Operations based on their readiness.
        Ensures the program is in a valid state before execution.
        """
        assert self.graph.roots is not None, "The operations graph has no root"
        
        while not self.run_executed and self.max_run_tries > 0:
            self.max_run_tries -= 1
            self.execution_queue = []
            self.got_steps = {}
            self.problem_parameters = copy.copy(self.original_problem_parameters)
            self.initialize_execution_queue()
            
            while self.execution_queue:
                self.execute_step()

            if self.final_thought.state["phase"] == "output":
                self.logger.info("All operations executed")
                self.run_executed = True

        if self.max_run_tries == 0:
            self.logger.error("Max tries reached on executing controller")


    def get_final_thoughts(self) -> List[List[Thought]]:
        """
        Retrieve the final thoughts after all operations have been executed.

        :return: List of thoughts for each operation in the graph's leaves.
        """
        assert self.run_executed, "The run method has not been executed"
        return [operation.get_thoughts() for operation in self.graph.leaves]

    def serialize_operation(self, operation) -> Dict[str, Any]:
        """
        Serialize an operation to a dictionary.

        :param operation: The operation to serialize.
        :return: The serialized operation.
        """
        operation_serialized = {
            "operation": operation.operation_type.name,
            "thoughts": [thought.state for thought in operation.get_thoughts()],
        }
        return operation_serialized

    def output_graph(self, path: str) -> None:
        """
        Serialize the state and results of the operations graph to a JSON file.

        :param path: The path to the output file.
        """
        output = [self.serialize_operation(op) for op in self.graph.operations]

        output.append(
            {
                "prompt_tokens": self.lm.prompt_tokens,
                "completion_tokens": self.lm.completion_tokens,
                "cost": self.lm.cost,
            }
        )

        with open(path, "w") as file:
            json.dump(output, file, indent=2)