import inference
from agent import Agent, AgentField

import asyncio

CLASS_DESCRIPTION = """
Expert AI with customizable system prompt.
Will get guidelines, documents, and other advanced features in the future.
""".strip()

class ExpertAgent(Agent):
    identifier = "expert"
    description = CLASS_DESCRIPTION
    schema_fields = [
        AgentField('system_prompt', 'string', default="You are an expert AI assistant."),
        AgentField('model', 'model_id', default=None),
        AgentField('temperature', 'number', default=None),
        AgentField('repetition_penalty', 'number', default=None),
    ]

    def __init__(self, *args):
        Agent.__init__(self, *args)

    async def run(self, run):
        opts = {}

        # Is this roon using *Agent's* tools, as opposed to run tools
        using_agent_tools = False
        if 'tools' not in run.input:
            # Agent tools will fully replace run tools if none are present,
            # otherwise they won't be used
            if self.tools:
                run.input['tools'] = self.tools
                using_agent_tools = True

        if self.config.get("temperature") is not None:
            opts["temperature"] = self.config["temperature"]
        if self.config.get("repetition_penalty") is not None:
            opts["repetition_penalty"] = self.config["repetition_penalty"]

        # Build message list based on input
        messages = inference.build_messages(
            system_prompt=self.config.get("system_prompt", ""),
            instruction=run.input.get("instruction"),
            input=run.input.get("input"),
            messages=run.input.get("messages"),
        )

        # Collect inference output. These are resused across calls
        tool_calls = []
        output_stream = None
        inference_end = asyncio.Event()
        def on_output(stream):
            nonlocal output_stream
            output_stream = stream
        run.on('output', on_output)
        run.on('inference-end', lambda event: inference_end.set())
        run.on('tool-call', lambda tc: print(f"[expert] tool call: {tc}") or tool_calls.append(tc))

        while True:
            # Reset data
            inference_end.clear()
            output_stream = None

            # Call inference, passing the run to get events out
            await inference.chat_completion(
                self.app,
                model_id=self.config.get("model"),
                messages=messages,
                tools=run.input.get("tools", []),
                opts=opts,
                agent_run=run,
            )

            await inference_end.wait()

            messages.append(inference.LlamaMessage('assistant', output_stream and await output_stream.content_wait() or ""))
            print(f"[expert] collected tool calls: {len(tool_calls)}, using? {using_agent_tools}")

            if using_agent_tools and tool_calls:
                print("Will execute tool calls")
                for tc in tool_calls:
                    tool = next(t for t in self.tools if t.name == tc.tool_name)
                    if not tool:
                        print(f"Bad agent tool name: {tc.name}")
                        continue
                    tc.result = await tool.execute_async(**tc.arguments)
                    messages.append(inference.LlamaMessage('tool', tc))
                tool_calls = []
                # Buffer time to let openrouter rest?
                await asyncio.sleep(2)
            else:
                print("No tool calls to execute, finish")
                break
        
        run.close()