import os
from functools import lru_cache

import gradio as gr
from huggingface_hub import hf_hub_download
from llama_cpp import Llama

MODEL_REPO = "microsoft/Phi-3-mini-4k-instruct-gguf"
MODEL_FILE = "Phi-3-mini-4k-instruct-q4.gguf"

@lru_cache(maxsize=1)
def get_model():
    model_path = hf_hub_download(repo_id=MODEL_REPO, filename=MODEL_FILE)
    return Llama(
        model_path=model_path,
        n_ctx=4096,
        n_threads=min(4, os.cpu_count() or 2),
        n_batch=128,
        use_mmap=True,
        verbose=False,
    )

def build_prompt(message, history):
    prompt = ""

    for user, assistant in history:
        prompt += f"User: {user}\nAssistant: {assistant}\n"

    prompt += f"User: {message}\nAssistant:"
    return prompt

def chat(message, history):
    try:
        model = get_model()
        prompt = build_prompt(message, history)

        output = model(
            prompt,
            max_tokens=200,
            temperature=0.7,
            top_p=0.9,
            stop=["User:", "</s>"],
        )

        text = output["choices"][0]["text"].strip()

        if not text:
            return "..."

        return text

    except Exception as e:
        return f"Error: {str(e)}"

demo = gr.ChatInterface(
    fn=chat,
    title="Phi-3 Mini CPU Chat",
)

if __name__ == "__main__":
    demo.launch()