import gradio as gr
from llama_cpp import Llama

# Load the RPG fine-tuned model (MythoMax-L2-13B GGUF)
# llama-cpp-python automatically downloads and caches this file on space startup
llm = Llama.from_pretrained(
    repo_id="TheBloke/MythoMax-L2-13B-GGUF",
    filename="mythomax-l2-13b.Q4_K_M.gguf",
    n_ctx=2048,
    n_threads=4,
    verbose=False
)

def predict(system_prompt, user_message):
    # Ensure system prompt fallback if empty
    sys_content = system_prompt.strip() if system_prompt and system_prompt.strip() else (
        "You are an immersive RPG Game Master. Describe outcomes vividly, enforce campaign rules fairly, "
        "and maintain an engaging narrative tone."
    )
    
    # ChatML formatting tuned for MythoMax
    prompt = f"<|im_start|>system\n{sys_content}<|im_end|>\n<|im_start|>user\n{user_message}<|im_end|>\n<|im_start|>assistant\n"
    
    # Model inference configuration
    output = llm(
        prompt=prompt,
        max_tokens=350,
        temperature=0.75,
        top_p=0.90,
        repeat_penalty=1.15,
        stop=["<|im_end|>", "User:", "\nUser:"]
    )
    
    return output["choices"][0]["text"].strip()

# Set up Gradio interface exposed to your HappySeeds client
demo = gr.Interface(
    fn=predict,
    inputs=[
        gr.Textbox(lines=8, label="System Prompt"),
        gr.Textbox(lines=3, label="User Action")
    ],
    outputs=gr.Textbox(lines=8, label="GM Response"),
    title="MYRPG AI Engine Backend",
    description="Fine-tuned RPG Game Master API endpoint running MythoMax-L2-13B."
)

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