import os
import sys
import torch
from threading import Thread
import gradio as gr
from transformers import AutoModelForCausalLM, AutoTokenizer, TextIteratorStreamer

MODEL_ID = "Bur3hani/Machi-Know-DeepSeek-8B"
HF_TOKEN = os.getenv("HF_TOKEN")

SYSTEM_PROMPT = (
    "Wewe ni Machi-Know, msaidizi mwenye maarifa tele na mtaalamu wa kila jambo (know-it-all). "
    "Unazungumza Kiswahili sanifu, chenye uchangamfu na busara. Jibu maswali yote kwa usahihi na kwa kina, "
    "kisha MARA ZOTE malizia jibu lako kwa kuuliza swali la kufuatilia (follow-up question) ili kuendeleza mazungumzo."
)

print(f"Loading model and tokenizer: {MODEL_ID}...")

tokenizer = None
model = None
load_error = None

try:
    tokenizer = AutoTokenizer.from_pretrained(
        MODEL_ID, 
        token=HF_TOKEN,
        trust_remote_code=True
    )

    device = "cuda" if torch.cuda.is_available() else "cpu"
    dtype = torch.float16 if torch.cuda.is_available() else torch.bfloat16

    print(f"Loading weights on device: {device} with dtype: {dtype}...")
    model = AutoModelForCausalLM.from_pretrained(
        MODEL_ID,
        token=HF_TOKEN,
        torch_dtype=dtype,
        low_cpu_mem_usage=True,
        device_map="auto" if torch.cuda.is_available() else None,
        trust_remote_code=True,
    )
    if not torch.cuda.is_available():
        model = model.to(device)
    print("Model loaded successfully!")

except Exception as e:
    load_error = str(e)
    print(f"Error loading model: {e}", file=sys.stderr)

def respond(message, history):
    if load_error or model is None or tokenizer is None:
        yield (
            f"⚠️ **Hitilafu ya Upakiaji wa Mfano (Model Load Error):**\n\n"
            f"```{load_error or 'AI Model failed to initialize. RAM memory limit exceeded or HF_TOKEN missing.'}```\n\n"
            f"📌 *Tafadhali hakikisha umeongeza `HF_TOKEN` katika Space Secrets na umechagua GPU Hardware ikiwa mfano ni mkubwa.*"
        )
        return

    messages = [{"role": "system", "content": SYSTEM_PROMPT}]
    
    for item in history:
        if isinstance(item, (list, tuple)) and len(item) == 2:
            u, a = item
            if u:
                messages.append({"role": "user", "content": u})
            if a:
                messages.append({"role": "assistant", "content": a})
        elif isinstance(item, dict) and "role" in item and "content" in item:
            messages.append({"role": item["role"], "content": item["content"]})
            
    messages.append({"role": "user", "content": message})

    try:
        prompt = tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
    except Exception:
        formatted = f"<|im_start|>system\n{SYSTEM_PROMPT}<|im_end|>\n"
        for item in history:
            if isinstance(item, (list, tuple)) and len(item) == 2:
                u, a = item
                if u:
                    formatted += f"<|im_start|>user\n{u}<|im_end|>\n"
                if a:
                    formatted += f"<|im_start|>assistant\n{a}<|im_end|>\n"
            elif isinstance(item, dict):
                formatted += f"<|im_start|>{item.get('role', 'user')}\n{item.get('content', '')}<|im_end|>\n"
        formatted += f"<|im_start|>user\n{message}<|im_end|>\n<|im_start|>assistant\n"
        prompt = formatted

    inputs = tokenizer([prompt], return_tensors="pt").to(model.device)
    streamer = TextIteratorStreamer(tokenizer, skip_prompt=True, skip_special_tokens=True)

    generation_kwargs = dict(
        inputs,
        streamer=streamer,
        max_new_tokens=512,
        temperature=0.7,
        top_p=0.9,
        do_sample=True,
    )

    thread = Thread(target=model.generate, kwargs=generation_kwargs)
    thread.start()

    partial_text = ""
    for new_text in streamer:
        partial_text += new_text
        yield partial_text

demo = gr.ChatInterface(
    fn=respond,
    title="Machi-Know Chatbot 🧠💬",
    description="Karibu! Mimi ni Machi-Know, msaidizi wako wa Kiswahili anayejua kila kitu. Niulize swali lolote!",
    examples=["Mambo vipi? Nieleze kuhusu akili mbandia (AI).", "Jinsi gani naweza kujifunza kuprogramu?", "Kwanini anga ni ya bluu?"],
)

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