# -*- coding: utf-8 -*-
import json
import logging
import re
import typing

from archytas.react import ReActAgent, Undefined
from archytas.tool_utils import AgentRef, LoopControllerRef, tool
from archytas.tool_utils import AgentRef, LoopControllerRef, ReactContextRef, tool
from beaker_kernel.lib.utils import togglable_tool
from beaker_kernel.lib.agent import BaseAgent

if typing.TYPE_CHECKING:
    from beaker_kernel.lib.context import BaseContext
    from .bunsen_context import BunsenContext

logger = logging.getLogger(__name__)


class BunsenAgent(BaseAgent):

    context: "BunsenContext"

    def __init__(
        self,
        context: "BaseContext" = None,
        tools: list = None,
        **kwargs,
    ):
        self.context = context

        super().__init__(
            context=context,
            tools=tools,
            **kwargs
        )

    @tool()
    async def generate_code(self, code_request: str, agent: AgentRef, loop: LoopControllerRef):
        """
        Generated code to be run in an interactive Jupyter notebook.

        Input is a full grammatically correct question about or request for an action to be performed in the current environment.
        If you need more information on how to accomplish the request, you should use the other tools prior to using this one.

        Args:
            code_request (str): A fully grammatically correct question about the current model.
        """

        code_generation_prompt = f"""
DO NOT USE A TOOL to generate the code. The code should be generated by the LLM and not by a tool.

Please generate {self.context.subkernel.DISPLAY_NAME} (via jupyter kernel "{self.context.subkernel.KERNEL_NAME}") code to
satisfy the user's request below.

Request:
```
{code_request}
```

Please generate the code as if you were programming inside a Jupyter Notebook and the code is to be executed inside a cell.
You MUST wrap the code with a line containing three backticks (```) before and after the generated code.
No addtional text is needed in the response, just the code block.
""".strip()

        agent: BaseAgent
        response = await agent.inspect(code_generation_prompt)
        preamble, code, coda = re.split("```\w*", response)
        loop.set_state(loop.STOP_SUCCESS)
        agent

        result = {
            "action": "code_cell",
            "language": self.context.subkernel.SLUG,
            "content": code.strip(),
        }

        return json.dumps(result)

    @tool()
    async def get_source_code(self, asset_type: str, asset_name: str, agent: AgentRef, loop: LoopControllerRef) -> list[str]:
        """
        Retrieves source code for a module, class, or function code asset.

        Input is a distinct identifier that uniquely distinguishes which asset you are wanting the source code of.
        Note that as this searches over the vectorized embeddings of the code using a nearest-neighbor search it is possible
        that you will receive irrelevent results if there is not a good match found.

        Args:
            asset_type (str): The type of asset you are looking up. Should be: "module", "class", or "function".
            asset_name (str): A distinct identifier that uniquely distinguishes which asset you are wanting the source code of.

        Returns:
            list: A list of chunks (strings) of source code that should correspond with the indicated asset. Order by cosine distance.
        """

        matches = agent.context.corpus.store.query(
            f"{asset_type} {asset_name}",
            partition="code",
            limit=3,
        )["matches"]
        return matches
