import dsp
import dspy
from ..primitives.program import Module
from ..primitives.python_interpreter import CodePrompt, PythonInterpreter
import re

class ProgramOfThought(Module):
    def __init__(self, signature, max_iters=3):
        super().__init__()
        self.signature = signature = dspy.Predict(signature).signature
        self.max_iters = max_iters

        self.input_fields = signature.input_fields()
        self.output_fields = signature.output_fields()

        inputs_ = ', '.join([f"`{field_name}`" for field_name in self.input_fields.keys()])
        outputs_ = ', '.join([f"`{field_name}`" for field_name in self.output_fields.keys()])

        assert len(self.output_fields) == 1, "PoT only supports one output field."
        
        instr = []
        instr.append(f"You will be given {inputs_} and you will respond with {outputs_}.")
        instr.append(f"Generating executable Python code that programmatically computes the correct {outputs_}.")
        instr.append(f"After you're done with the computation, make sure the last line in your code evaluates to the correct value for {outputs_}.")
        instr = '\n'.join(instr)
        
        self.code_generate = dspy.ChainOfThought(dsp.Template(self._generate_instruction('generate'), **self._generate_signature('generate')))
        self.code_regenerate = dspy.ChainOfThought(dsp.Template(self._generate_instruction('regenerate'), **self._generate_signature('regenerate')))
        self.generate_answer = dspy.ChainOfThought(dsp.Template(self._generate_instruction('answer'), **self._generate_signature('answer')))

    def _generate_signature(self, mode):
        signature_dict = dict(self.input_fields)
        fields_for_mode = {
            'generate': {
                'generated_code': dspy.OutputField(prefix="Code:", desc="python code that answers the question", format=str)
            },
            'regenerate': {
                'previous_code': dspy.InputField(prefix="Previous Code:", desc="previously-generated python code that errored", format=str),
                'error': dspy.InputField(prefix="Error:", desc="error message from previously-generated python code"),
                'generated_code': dspy.OutputField(prefix="Code:", desc="python code that answers the question", format=str)
            },
            'answer': {
                'final_generated_code': dspy.InputField(prefix="Code:", desc="python code that answers the question", format=str),
                'code_output': dspy.InputField(prefix="Code Output:", desc="output of previously-generated python code"),
                'answer': self.signature.kwargs["answer"]
            }
        }
        signature_dict.update(fields_for_mode[mode])
        return signature_dict

    def _generate_instruction(self, mode):
        mode_inputs = ', '.join([f"`{field_name}`" for field_name in self._generate_signature(mode).keys() if isinstance(self._generate_signature(mode)[field_name], dspy.InputField)])
        mode_outputs = ', '.join([f"`{field_name}`" for field_name in self._generate_signature(mode).keys() if isinstance(self._generate_signature(mode)[field_name], dspy.OutputField)])
        if mode == 'generate':
            instr = [
                f"You will be given {mode_inputs} and you will respond with {mode_outputs}.",
                f"Generating executable Python code that programmatically computes the correct {mode_outputs}.",
                f"After you're done with the computation, make sure the last line in your code evaluates to the correct value for {mode_outputs}."
            ]
        elif mode == 'regenerate':
            instr = [
                f"You are given {mode_inputs} due to an error in previous code.",
                f"Your task is to correct the error and provide the new {mode_outputs}."
            ]
        else:  # mode == 'answer'
            instr = [
                f"Given the final code {mode_inputs}, provide the final {mode_outputs}."
            ]

        return '\n'.join(instr)

    def parse_code(self, code_data):
        code = code_data.get('generated_code', '').split('---', 1)[0].split('\n\n\n', 1)[0]
        code_match = re.search(r'```python[ \n](.*?)[ \n]```?', code, re.DOTALL)
        code_block = (code_match.group(1) if code_match else code).replace('\\n', '\n')
        if not code_block:
            return code, "Error: Empty code after parsing."
        if "\n" not in code_block and code_block.count('=') > 1:
            return code, "Error: Code format is not correct."
        lines = code_block.split('\n')
        last_line_match = re.match(r'^(\w+)\s*=', lines[-1].strip())
        if last_line_match and len(lines) > 1:
            code_block += '\n' + last_line_match.group(1)
        else:
            code_block = re.sub(r'([a-zA-Z_]\w* *=.*?)(?=[a-zA-Z_]\w* *=)', r'\1\n', code_block)
            code_block = re.sub(r'([a-zA-Z_]\w* *=.*?)([a-zA-Z_]\w*)$', r'\1\n\2', code_block)
        return code_block, None

    def execute_code(self, code):
        if not code:
            return code, None, 'Error: Empty code before execution.'
        code_prompt = CodePrompt(code, code_type="python")
        interpreter = PythonInterpreter(action_space={"print": print})
        try:
            output = str(code_prompt.execute(interpreter=interpreter)[0])
            return code, output, None
        except Exception as e:
            return code, None, str(e)
            
    def forward(self, **kwargs):
        code_data = self.code_generate(question=kwargs["question"])
        parsed_code, error = self.parse_code(code_data)
        code, output, error = self.execute_code(parsed_code)
        hop = 0
        while hop < self.max_iters and error:
            print('Error in code execution')
            code_data = self.code_regenerate(question=kwargs["question"], previous_code=code, error=error)
            parsed_code, error = self.parse_code(code_data)
            hop += 1
            if hop == self.max_iters:
                print('Max hops reached. Error persists.')
                return None
        answer_gen_result = self.generate_answer(question=kwargs["question"], final_generated_code=code, code_output=output)
        return answer_gen_result