from stream import Stream
import threading

LLAMA3_MODEL = "meta/meta-llama-3-8b-instruct"

# Instances of this class lives very briefly, created when needed and immediately disposed.
# Subclasses must be careful in not setting up long term resources
class AgentControl:
    def __init__(self, app, agent_id):
        self.app = app
        self.agent_id = agent_id
        self.system_prompt = "You are a useful and precise AI assistant. Respond concisely."

    # Receives a single string prompt, and must respond with a string
    async def prompt(self, content):
        response = await replicate_llama_completion(self.app, [
            LlamaMessage("system", self.system_prompt),
            LlamaMessage("user", content)
        ])
        return response

    async def received_message(self, message):
        chat_context = self.app.get_chat_context(message.chat_id)
        def sender_name(message):
            sender = message['sender']
            is_digit = isinstance(sender, (int, str)) and str(sender).isdigit()
            return "assistant" if is_digit else sender
        chat_summary = "\n".join(f"**{sender_name(message)}**: {message['content']}" for message in chat_context.messages)

        prologue = "\nYou are given the entire context of the conversation, respond with your next message in the conversation."
        print(chat_summary)

        response = await replicate_llama_completion(self.app, [
            LlamaMessage("system", self.system_prompt + prologue),
            LlamaMessage("user", chat_summary)
        ])

        self.app.send_chat_message(
            chat_id=message.chat_id,
            sender=self.agent_id,
            content=response,
        )
    
    def create_from_data(app, agent_id, kind, data):
        return AgentControl(app, agent_id)

class LlamaMessage:
    def __init__(self, role, content):
        self.role = role
        self.content = content

async def replicate_llama_completion(app, messages):
    # Llama 3 instruct
    full_prompt = "<|begin_of_text|>"
    for message in messages:
        header_name = message.role if message.role != "agent" else "assistant"
        full_prompt += f"<|start_header_id|>{header_name}<|end_header_id|>\n\n{message.content}<|eot_id|>"

    full_prompt += f"<|start_header_id|>assistant<|end_header_id|>"
    
    content_stream = Stream(app)
    thread = threading.Thread(
        target=run_replicate_inference_loop,
        kwargs={
            'replicate_client': app.replicate_client,
            'model': LLAMA3_MODEL,
            'prompt': full_prompt,
            'content_stream': content_stream,
        }
    )

    thread.start()
    response = await content_stream.content_wait()
    
    return response.lstrip()


def run_replicate_inference_loop(replicate_client, model, prompt, content_stream):
    print("START INFERENCE")
    completion = replicate_client.run(
        model,
        input={
            "prompt": prompt,
            "temperature": 0.75,
            "max_new_tokens": 500,
            "system_prompt": "" # We handle system prompt in the formatted prompt
        }
    )
    for event in completion:
        if event:
            content_stream.feed(str(event))
    print("END INFERENCE")
    content_stream.close()