import os

# ২ vCPU-তে বাড়তি থ্রেড (BLAS/OpenMP) খোলা বন্ধ রাখতে Llama() লোডের আগেই সেট করতে হবে
os.environ["OMP_NUM_THREADS"] = "2"
os.environ["OPENBLAS_NUM_THREADS"] = "2"

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

# ========================
# ১. মডেল ডাউনলোড
# ========================
#MODEL_REPO = "prism-ml/Bonsai-8B-gguf"
#MODEL_FILE = "Bonsai-8B-Q1_0.gguf"

#MODEL_REPO = "prism-ml/Bonsai-27B-gguf"
#MODEL_FILE = "Bonsai-27B-Q1_0.gguf"

MODEL_REPO = "prism-ml/Bonsai-1.7B-gguf"
MODEL_FILE = "Bonsai-1.7B-Q1_0.gguf"

print("📥 মডেল ডাউনলোড হচ্ছে...")
model_path = hf_hub_download(
    repo_id=MODEL_REPO,
    filename=MODEL_FILE,
    local_dir="./models"
)
print("✅ মডেল ডাউনলোড সম্পূর্ণ!")

# ========================
# ২. LLM লোড করুন
# HF Spaces free tier: 2 vCPU, 16GB RAM, no GPU
# ========================
llm = Llama(
    model_path=model_path,
    n_ctx=1024,             # কমানো হলো — MAX_TURNS দিয়ে history ছোট রাখা হচ্ছে, তাই বড় ctx দরকার নেই
    n_threads=2,
    n_threads_batch=2,
    n_batch=128,
    n_gpu_layers=0,
    use_mmap=True,
    use_mlock=False,
    # NOTE: quantized KV cache (type_k/type_v = Q8_0) বাদ দেওয়া হলো —
    # এটার জন্য flash_attn=True লাগে, আর CPU build-এ flash attention প্রায়ই
    # অসমর্থিত/অস্থিতিশীল, ফলে "Failed to create llama_context" এরর হয়েছিল।
    verbose=False,
)

print("🔥 মডেল ওয়ার্ম-আপ হচ্ছে...")
try:
    list(llm("Hi", max_tokens=1, stream=False))
except Exception as e:
    print(f"ওয়ার্ম-আপ স্কিপ হলো: {e}")
print("✅ ওয়ার্ম-আপ সম্পূর্ণ, সার্ভার রেডি!")

# ========================
# ৩. চ্যাট ফাংশন (স্ট্রিমিং + কনটেক্সট-সেফ হিস্ট্রি)
# ========================
MAX_TURNS = 4  # সীমিত RAM/ctx-তে পুরো history না রেখে শেষ কয়েক টার্ন রাখাই ভালো

def build_prompt(message, history):
    trimmed = history[-MAX_TURNS:] if history else []
    prompt = ""
    for item in trimmed:
        if isinstance(item, (list, tuple)) and len(item) >= 2:
            user_msg, bot_msg = item[0], item[1]
            # আগের রেসপন্স থেকে স্পিড-মেট্রিক অংশ বাদ দিন যাতে prompt নোংরা না হয়
            if bot_msg:
                bot_msg = bot_msg.split("\n\n⚡")[0]
            prompt += f"User: {user_msg}\nAssistant: {bot_msg}\n"
    prompt += f"User: {message}\nAssistant:"
    return prompt

def generate_response(message, history):
    prompt = build_prompt(message, history)
    start_time = time.time()

    try:
        stream = llm(
            prompt,
            max_tokens=200,              # সামান্য কমানো — লেটেন্সি cap
            stop=["User:", "\nUser"],
            echo=False,
            stream=True,
            temperature=0.2,
            top_p=0.9,
            repeat_penalty=1.15,
        )
    except Exception as e:
        yield f"⚠️ ত্রুটি হয়েছে: {e}"
        return

    response_text = ""
    token_count = 0
    UI_UPDATE_EVERY = 3   # প্রতি টোকেনে UI re-render না করে ৩ টোকেন পরপর আপডেট — actual generation-এ বেশি CPU সময় যায়

    try:
        for chunk in stream:
            choices = chunk.get("choices", [])
            if choices:
                delta = choices[0].get("text", "")
                if delta:
                    response_text += delta
                    token_count += 1
                    if token_count % UI_UPDATE_EVERY == 0:
                        yield response_text + "▌"
    except Exception as e:
        yield response_text + f"\n\n⚠️ স্ট্রিমিং-এ সমস্যা: {e}"
        return

    elapsed = time.time() - start_time
    speed = token_count / elapsed if elapsed > 0 else 0
    yield f"{response_text}\n\n⚡ {speed:.1f} tok/s | {elapsed:.1f}s"

# ========================
# ৪. Gradio UI
# ========================
with gr.Blocks(title="Bonsai-8B (CPU)", theme=gr.themes.Soft()) as demo:
    gr.Markdown("""
    # 🌳 Bonsai-8B Chat
    ### ⚡ ২-কোর CPU | ১৬ GB RAM | ১-বিট কোয়ান্টাইজড | স্ট্রিমিং
    """)

    gr.ChatInterface(
        fn=generate_response,
        title="Bonsai-8B",
        description="আপনার প্রশ্ন করুন... (টাইপিং ইফেক্ট দেখতে পাবেন)",
        concurrency_limit=1,   # ২ vCPU-তে একসাথে একাধিক জেনারেশন চালালে দুটোই স্লো হয়ে যায়; একবারে একটাই প্রসেস হবে
    )

if __name__ == "__main__":
    demo.queue().launch(server_name="0.0.0.0", server_port=7860)