import os
import time
import json
import sys

import gradio as gr

# এই ভার্সনটা ইচ্ছাকৃতভাবে CPU-only — এই Space-এ GPU (ZeroGPU) সাথে llama.cpp-এর
# persistent CUDA context ঠিকভাবে মিলছিল না (প্রতিটা রিকোয়েস্টে GPU attach/detach হওয়ায়
# মডেল আবার লোড করতে হতো), তাই এখানে GPU-চেষ্টা ও `spaces` ডিপেন্ডেন্সি বাদ দেওয়া হয়েছে।

# গুরুত্বপূর্ণ: os.cpu_count() কন্টেইনারের হোস্ট মেশিনের মোট কোর সংখ্যা রিপোর্ট করে,
# কিন্তু Spaces-এর মতো শেয়ার্ড হোস্টে আপনার কন্টেইনার আসলে তার চেয়ে অনেক কম CPU quota
# (cgroup affinity) পায়। os.cpu_count() ধরে থ্রেড বানালে বরাদ্দের চেয়ে বেশি থ্রেড
# একে অপরের সাথে প্রতিযোগিতা করে (context-switch overhead), যেটা গতি কমিয়ে দেয়।
# sched_getaffinity আসল বরাদ্দকৃত কোর সংখ্যা দেয় (শুধু Linux-এ কাজ করে, তাই fallback রাখা হলো)।
try:
    cpu_count = len(os.sched_getaffinity(0))
except AttributeError:
    cpu_count = os.cpu_count() or 4
print(f"🧵 ব্যবহারযোগ্য CPU কোর সনাক্ত হয়েছে: {cpu_count}")

# BLAS/OpenMP ব্যাকএন্ড যদি থাকে, সেটাও যেন একই সংখ্যক থ্রেড ব্যবহার করে —
# নাহলে llama.cpp-এর থ্রেড + BLAS-এর নিজস্ব থ্রেড মিলে ওভারসাবস্ক্রাইব হতে পারে।
# native library লোড হওয়ার আগেই সেট করা জরুরি, তাই llama_cpp ইম্পোর্টের আগে বসানো হলো।
os.environ.setdefault("OMP_NUM_THREADS", str(cpu_count))
os.environ.setdefault("OPENBLAS_NUM_THREADS", str(cpu_count))

# llama_cpp ইম্পোর্ট
#
# নোট: Bonsai-27B-Q1_0.gguf-এর Q1_0 (1-বিট) ফরম্যাট mainline llama.cpp-তেই সাপোর্টেড —
# CPU ব্যাকএন্ডে কাজ করার জন্য আলাদা কোনো কাস্টম কার্নেল/ফর্ক লাগে না।
try:
    from llama_cpp import Llama
    print("✅ llama-cpp-python ইম্পোর্ট সফল")
except Exception as e:
    print(f"❌ llama-cpp-python ইম্পোর্ট ব্যর্থ: {e}")
    sys.exit(1)

# --- ১. মডেল কনফিগারেশন ---
#MODEL_REPO = "prism-ml/Bonsai-27B-gguf"
#MODEL_FILE = "Bonsai-27B-Q1_0.gguf"

# ✅ বর্তমানে সক্রিয়: Bonsai 8B (~1.16 GB)
MODEL_REPO = "prism-ml/Bonsai-8B-gguf"
MODEL_FILE = "Bonsai-8B-Q1_0.gguf"

print("⏳ মডেল ডাউনলোড হচ্ছে...")

from huggingface_hub import hf_hub_download
model_path = hf_hub_download(
    repo_id=MODEL_REPO,
    filename=MODEL_FILE,
    local_dir="./models",
    token=None
)
print(f"✅ ডাউনলোড সম্পূর্ণ: {model_path}")

# --- ২. মডেল লোড ---
print("⏳ মডেল লোড হচ্ছে...")


def load_model():
    """CPU-তে সর্বোচ্চ থ্রুপুটের জন্য টিউন করা লোড"""
    print("🐢 CPU মোডে লোড হচ্ছে...")
    llm = Llama(
        model_path=model_path,
        n_ctx=4096,           # কনটেক্সট যত বড়, প্রম্পট প্রসেসিং তত ধীর — CPU-তে রক্ষণশীল রাখা ভালো
        n_gpu_layers=0,       # CPU-only
        n_threads=cpu_count,        # জেনারেশন থ্রেড — আসল বরাদ্দকৃত কোর সংখ্যা
        n_threads_batch=cpu_count,  # প্রম্পট-প্রসেসিং (ব্যাচ) থ্রেড
        n_batch=256,          # CPU-তে ছোট batch সাধারণত বড়টার চেয়ে কার্যকর হয়
        type_k=8,             # KV cache-এর K অংশ Q8_0-এ কোয়ান্টাইজ — মেমরি-ব্যান্ডউইথ বাঁচায়
                              # (V cache কোয়ান্টাইজেশনের জন্য flash_attn লাগে, তাই এটা বাদ)
        use_mmap=True,        # ডিস্ক থেকে সরাসরি ম্যাপ করে, RAM-এ পুরো কপি লাগে না
        use_mlock=False,      # RAM লক না করা — কম RAM এর হোস্টে সোয়াপ/ক্র্যাশ ঠেকায়
        verbose=True
    )
    print("✅ CPU মোডে সফলভাবে লোড হয়েছে!")
    return llm, "CPU"


try:
    llm, device_mode = load_model()
except Exception as e:
    print(f"❌ মডেল লোড করা সম্ভব হয়নি: {e}")
    sys.exit(1)

# --- ৩. ওয়ার্মআপ ---
try:
    print("🔥 ওয়ার্মআপ হচ্ছে...")
    list(llm("<|im_start|>user\nহাই<|im_end|>\n<|im_start|>assistant\n<think>\n\n</think>\n\n",
              max_tokens=1, stream=True))
    print(f"✅ ওয়ার্মআপ সম্পন্ন! ({device_mode})")
except Exception as e:
    print(f"⚠️ ওয়ার্মআপে সমস্যা: {e}")


# ===== ৪. চ্যাট জেনারেশন ফাংশন =====
def generate_reply(user_prompt):
    """মডেল থেকে স্ট্রিমিং টেক্সট জেনারেট করে; পূর্ণ টেক্সট + মেটা রিটার্ন করে (জেনারেটর)"""
    full_prompt = f"<|im_start|>user\n{user_prompt}<|im_end|>\n<|im_start|>assistant\n<think>\n\n</think>\n\n"

    start_time = time.time()
    start_str = time.strftime('%H:%M:%S', time.localtime(start_time))
    first_token_time = None
    full_text = ""

    try:
        stream = llm(
            full_prompt,
            max_tokens=256,
            temperature=0.1,
            top_p=0.2,
            top_k=40,
            repeat_penalty=1.1,
            stop=["<|im_end|>", "user:", "User:"],
            stream=True
        )

        for chunk in stream:
            token_text = chunk["choices"][0]["text"]
            if token_text:
                if first_token_time is None:
                    first_token_time = time.time() - start_time
                full_text += token_text
                yield full_text, None
    except Exception as e:
        print(f"❌ জেনারেশন ত্রুটি: {e}")
        full_text += f"\n⚠️ মডেল ত্রুটি: {str(e)}"
        yield full_text, None
    finally:
        end_time = time.time()
        end_str = time.strftime('%H:%M:%S', time.localtime(end_time))
        elapsed = round(end_time - start_time, 2)
        first_token = round(first_token_time, 2) if first_token_time else 0

        meta = {
            "start": start_str,
            "end": end_str,
            "elapsed": elapsed,
            "first_token_time": first_token
        }
        yield full_text, meta


def chat_fn(user_prompt, history):
    """Gradio Chatbot-এর জন্য wrapper: history আপডেট করে stream করে"""
    if not user_prompt or not user_prompt.strip():
        yield history, ""
        return

    history = history + [
        {"role": "user", "content": user_prompt},
        {"role": "assistant", "content": "⏳ উত্তর তৈরি হচ্ছে..."}
    ]
    yield history, ""

    last_meta = None
    for partial_text, meta in generate_reply(user_prompt):
        history[-1]["content"] = partial_text
        if meta is not None:
            last_meta = meta
        yield history, ""

    if last_meta is not None:
        time_info = (
            f"⏱️ **প্রথম টোকেন:** {last_meta['first_token_time']}s &nbsp;|&nbsp; "
            f"**মোট সময়:** {last_meta['elapsed']}s &nbsp;|&nbsp; "
            f"**শুরু:** {last_meta['start']} &nbsp;|&nbsp; "
            f"**শেষ:** {last_meta['end']}"
        )
        history.append({"role": "assistant", "content": time_info})

    yield history, ""


def clear_chat():
    return [], ""


# ===== ৫. কাস্টম CSS =====
CUSTOM_CSS = """
.gradio-container {
    max-width: 900px !important;
    margin: 20px auto !important;
    background: #f0f2f5 !important;
}
#header-box {
    background: linear-gradient(135deg, #667eea 0%, #764ba2 100%);
    color: white;
    padding: 15px 20px;
    border-radius: 16px 16px 0 0;
    margin-bottom: 0px;
}
#header-box h2 { margin: 0; font-weight: 400; }
#header-box small { opacity: 0.85; font-size: 14px; }
.message-wrap .message.user {
    background: #007bff !important;
    color: white !important;
    border-radius: 18px !important;
}
.message-wrap .message.bot {
    background: #e9ecef !important;
    color: #333 !important;
    border-radius: 18px !important;
}
footer { display: none !important; }
"""

DEVICE_BADGE = "🐢 CPU মোড"

with gr.Blocks(title="Bonsai 27B চ্যাট (CPU)") as demo:
    gr.HTML(f"""
    <div id="header-box">
        <h2>🌳 Bonsai 27B (১-বিট) চ্যাট</h2>
        <small>মডেল সাইজ: ~৩.৮ GB • {DEVICE_BADGE} — উত্তর আসতে কিছুটা সময় লাগতে পারে</small>
    </div>
    """)

    chatbot = gr.Chatbot(
        value=[{"role": "assistant", "content": "👋 হ্যালো! আমি Bonsai 27B। আপনি কী জানতে চান?"}],
        height=400,
        show_label=False
    )

    with gr.Row():
        user_input = gr.Textbox(
            placeholder="এখানে প্রশ্ন লিখুন...",
            show_label=False,
            scale=8
        )
        send_btn = gr.Button("পাঠান", scale=1, variant="primary")
        clear_btn = gr.Button("🗑️", scale=1)

    user_input.submit(chat_fn, [user_input, chatbot], [chatbot, user_input])
    send_btn.click(chat_fn, [user_input, chatbot], [chatbot, user_input])
    clear_btn.click(clear_chat, None, [chatbot, user_input])

    gr.Markdown(
        "<div style='text-align:center;color:#a0aec0;font-size:12px;'>"
        "Powered by llama-cpp-python • Bonsai 27B (Q1_0, 1-bit)</div>"
    )

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