import re
import json
import inspect
from string import Template

import inference
from inference import LlamaMessage
from agent import Agent, AgentField

CLASS_DESCRIPTION = """
Expert AI in any field, powered by Llama 3 70b.
Will get guidelines, documents, and other advanced features in the future.

Available models are: [llama/qwen]-[big/small] (replicate), ollama:MODEL_NAME
""".strip()

class ExpertAgent(Agent):
    identifier = "expert"
    description = CLASS_DESCRIPTION
    schema_fields = [
        AgentField('system_prompt', 'string', default="You are an expert AI assistant."),
        AgentField('model', 'string', default="qwen-big"),
    ]

    def __init__(self, app, data):
        Agent.__init__(self, app, data)

    async def prompt(self, instruction, input=None, **kwargs):
        # Instruction extends system if there's input or messages
        # Otherwise instruction is the user prompt
        # Input is always added as a user message at the end, if present

        # Use 'messages' as starting point when passed
        messages = []
        if 'messages' in kwargs:
            messages = [LlamaMessage(msg['role'], msg['content']) for msg in kwargs['messages']]
        
        # Construct the default messages
        system = self.config['system_prompt']
        prompt = input

        if input == None and not messages:
            prompt = instruction
        elif instruction:
            system += "\n" + instruction
        
        # Add default messages to list
        messages.insert(0, LlamaMessage("system", system))
        if prompt:
            messages.append(LlamaMessage("user", prompt))

        model = self.config['model']

        if model.startswith("ollama:"):
            response = await inference.ollama_chat_completion(self.app, messages, model=model[7:])
        elif model == "qwen-big":
            response = await inference.replicate_qwen_completion(self.app, messages, model=inference.Models.qwen3_235b)
        elif model == "qwen-small":
            response = await inference.replicate_qwen_completion(self.app, messages, model=inference.Models.qwen3_8b)
        elif model == "llama-big":
            response = await inference.replicate_llama_completion(self.app, messages, model=inference.Models.llama_70b)
        elif model == "llama-small":
            response = await inference.replicate_llama_completion(self.app, messages, model=inference.Models.llama_8b)
        else:
            # An error here is questionable, model is user configurable
            raise Exception("Unkonwn agent model: " + model)

        return response

    async def received_message(self, message):
        chat_view = await self.app.get_chat_view(message.chat_id)
        thread = chat_view.get_thread(message.thread_id)

        def msg_role(message):
            sender = message.sender
            is_digit = isinstance(sender, (int, str)) and str(sender).isdigit()

            if is_digit and str(sender) == str(self.id):
                return 'assistant'
            elif sender == 'user':
                return 'user'
            else:
                raise Exception("Unsupported sender: " + sender)
        
        messages = [{'role': msg_role(msg), 'content': msg.content} for msg in thread.messages]

        response = await self.prompt(
            instruction=None,
            input=None,
            messages=messages
        )

        await self.app.send_chat_message(
            chat_id=message.chat_id,
            sender=self.id,
            content=response,
            thread_id=message.thread_id,
        )