import os
from fastapi import FastAPI
from pydantic import BaseModel
from sse_starlette.sse import EventSourceResponse
from huggingface_hub import hf_hub_download
from llama_cpp import Llama

app = FastAPI(title="Llama Chat & RAG Engine")

# 1. Define Model Settings
MODEL_REPO = "bartowski/Llama-3.2-3B-Instruct-GGUF"
MODEL_FILE = "Llama-3.2-3B-Instruct-Q4_K_M.gguf"

print("Downloading Llama-3.2-3B weights... (This takes a moment)")
model_path = hf_hub_download(repo_id=MODEL_REPO, filename=MODEL_FILE)

print("Loading model into memory...")
# Llama 3.2 3B uses very little RAM, leaving massive headroom for RAG context
llm = Llama(
    model_path=model_path,
    n_ctx=4096, 
    n_threads=2,
    n_batch=512,
    verbose=False
)

class GenerateRequest(BaseModel):
    prompt: str
    max_tokens: int = 1024
    temperature: float = 0.7 # Slightly higher temp for natural chat

@app.get("/")
def health_check():
    return {"status": "Llama Chat Engine is Online and Ready"}

@app.post("/generate")
async def generate(request: GenerateRequest):
    # Format the prompt using exact Llama 3 syntax
    formatted_prompt = (
        f"<|begin_of_text|><|start_header_id|>system<|end_header_id|>\n\n"
        f"You are a highly intelligent AI assistant.<|eot_id|>"
        f"<|start_header_id|>user<|end_header_id|>\n\n"
        f"{request.prompt}<|eot_id|><|start_header_id|>assistant<|end_header_id|>\n\n"
    )
    
    def token_generator():
        stream = llm(
            formatted_prompt,
            max_tokens=request.max_tokens,
            temperature=request.temperature,
            stream=True
        )
        for output in stream:
            token = output["choices"][0]["text"]
            if token:
                yield {"data": token}

    return EventSourceResponse(token_generator())