import json
from typing import Optional

from injector import inject

from taskweaver.code_interpreter.code_executor import CodeExecutor
from taskweaver.code_interpreter.code_generator import CodeGeneratorPluginOnly
from taskweaver.config.module_config import ModuleConfig
from taskweaver.logging import TelemetryLogger
from taskweaver.memory import Memory, Post
from taskweaver.memory.attachment import AttachmentType
from taskweaver.role import Role


class CodeInterpreterConfig(ModuleConfig):
    def _configure(self):
        self._set_name("code_interpreter_plugin_only")
        self.use_local_uri = self._get_bool("use_local_uri", False)
        self.max_retry_count = self._get_int("max_retry_count", 3)


class CodeInterpreterPluginOnly(Role):
    @inject
    def __init__(
        self,
        generator: CodeGeneratorPluginOnly,
        executor: CodeExecutor,
        logger: TelemetryLogger,
        config: CodeInterpreterConfig,
    ):
        self.generator = generator
        self.executor = executor
        self.logger = logger
        self.config = config
        self.retry_count = 0
        self.return_index = 0

        self.logger.info("CodeInterpreter initialized successfully.")

    def reply(
        self,
        memory: Memory,
        event_handler: callable,
        prompt_log_path: Optional[str] = None,
        use_back_up_engine: Optional[bool] = False,
    ) -> Post:
        response: Post = self.generator.reply(
            memory,
            event_handler,
        )

        if response.message is not None:
            return response

        functions = json.loads(response.get_attachment(type=AttachmentType.function)[0])
        if len(functions) > 0:
            code = []
            for i, f in enumerate(functions):
                function_name = f["name"]
                function_args = json.loads(f["arguments"])
                function_call = (
                    f"r{self.return_index + i}={function_name}("
                    + ", ".join(
                        [
                            f'{key}="{value}"' if isinstance(value, str) else f"{key}={value}"
                            for key, value in function_args.items()
                        ],
                    )
                    + ")"
                )
                code.append(function_call)
            code.append(f'{", ".join([f"r{self.return_index + i}" for i in range(len(functions))])}')
            self.return_index += len(functions)

            event_handler("code", "\n".join(code))
            exec_result = self.executor.execute_code(
                exec_id=response.id,
                code="\n".join(code),
            )

            response.message = self.executor.format_code_output(
                exec_result,
                with_code=True,
                use_local_uri=self.config.use_local_uri,
            )
            event_handler("CodeInterpreter-> Planner", response.message)
        else:
            response.message = "No code is generated because no function is selected."
            event_handler("CodeInterpreter-> Planner", response.message)

        return response