import threading
from stream import Stream

LLAMA3_MODEL = "meta/meta-llama-3-8b-instruct"
VALID_ROLES = ['user', 'agent', 'system']

class LlamaMessage:
    def __init__(self, role, content):
        if role not in VALID_ROLES:
            raise Exception("message role must be one of: " + ",".join(VALID_ROLES))
        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()