import os
import gradio as gr
from transformers import pipeline

MODEL_ID = os.getenv("MODEL_ID", "GenueAI/Inelly-4.5-Blaze")
MODEL_NAME = os.getenv("MODEL_NAME", "Inelly 4.5 Blaze")
pipe = None

def load_model():
    global pipe
    pipe = pipeline(
        "text-generation", 
        model=MODEL_ID, 
        torch_dtype="auto",
        model_kwargs={
            "low_cpu_mem_usage": True, 
            "device_map": "sequential"
        }
    )

def build_prompt(message_text, history, system_prompt):
    messages = []
    if system_prompt.strip():
        messages.append({"role": "system", "content": system_prompt.strip()})

    for msg in history:
        messages.append({"role": msg["role"], "content": msg["content"]})
        
    messages.append({"role": "user", "content": message_text})
    return messages

def chat(message, history, system_prompt):
    message_text = message.get("text", "").strip() if isinstance(message, dict) else str(message).strip()
    
    if not message_text:
        return ""

    prompt = build_prompt(message_text, history, system_prompt)
    
    outputs = pipe(prompt) 
    
    try:
        return outputs[0]["generated_text"][-1]["content"]
    except (KeyError, IndexError, TypeError):
        try:
            return outputs["generated_text"][-1]["content"]
        except (KeyError, IndexError, TypeError):
            return str(outputs)

with gr.Blocks(title="Genue Chat") as demo:
    gr.Markdown("# Genue Chat")
    gr.Markdown(f"Chat with {MODEL_NAME} from Hugging Face.")

    with gr.Row():
        with gr.Column(scale=4):
            chatbot = gr.ChatInterface(
                fn=chat,
                additional_inputs=[
                    gr.Textbox(
                        label="System prompt",
                        value="You are a helpful assistant.",
                        lines=3,
                    ),
                ],
                textbox=gr.Textbox(
                    label="Prompt",
                    placeholder="Ask anything...",
                    container=False,
                    scale=7,
                ),
            )

if __name__ == "__main__":
    load_model()
    demo.queue()
    demo.start()