import os
import gradio as gr
import torch

MODEL_ID = "SatyamSinghal/taskmind-1.1b-chat-lora"
HF_TOKEN = os.getenv("HF_TOKEN")

pipe = None

def load_model():
    global pipe
    if pipe is not None:
        return

    from peft import AutoPeftModelForCausalLM
    from transformers import AutoTokenizer, pipeline

    print("Loading tokenizer...")
    tokenizer = AutoTokenizer.from_pretrained(
        MODEL_ID,
        token=HF_TOKEN,
    )

    print("Loading model...")
    model = AutoPeftModelForCausalLM.from_pretrained(
        MODEL_ID,
        token=HF_TOKEN,
        torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32,
        low_cpu_mem_usage=True,
    )

    pipe = pipeline(
        "text-generation",
        model=model,
        tokenizer=tokenizer,
    )
    print("Model loaded successfully.")


def respond(message, history):
    try:
        load_model()
    except Exception as e:
        return f"❌ Model failed to load: {str(e)}"

    messages = []
    for user_msg, assistant_msg in history:
        if user_msg:
            messages.append({"role": "user", "content": user_msg})
        if assistant_msg:
            messages.append({"role": "assistant", "content": assistant_msg})
    messages.append({"role": "user", "content": message})

    result = pipe(
        messages,
        max_new_tokens=256,
        do_sample=True,
        temperature=0.7,
        top_p=0.9,
    )

    generated = result[0]["generated_text"]
    if isinstance(generated, list):
        return generated[-1]["content"]
    return str(generated)


demo = gr.ChatInterface(
    fn=respond,
    title="TaskMind Interface",
    description="Chat with the TaskMind LoRA model.",
    examples=[
        "Who are you?",
        "@Satyam fix the growstreams deck ASAP NO Delay",
        "done bhai, merged the PR",
        "login page 60% ho gaya",
        "getting 500 error on registration",
    ],
)

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