import os
import time
import json
import sys

import gradio as gr
import spaces

# ===== GPU ডিটেক্ট =====
try:
    import torch
    HAS_GPU = torch.cuda.is_available()
    if HAS_GPU:
        print(f"✅ GPU পাওয়া গেছে: {torch.cuda.get_device_name(0)}")
    else:
        print("⚠️ GPU পাওয়া যায়নি, CPU ব্যবহার করা হবে")
except Exception as e:
    HAS_GPU = False
    print(f"⚠️ torch ইম্পোর্টে সমস্যা: {e}")

# llama_cpp ইম্পোর্ট
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-8B-gguf"
#MODEL_FILE = "Bonsai-8B-Q1_0.gguf"

MODEL_REPO = "prism-ml/Bonsai-27B-gguf"
MODEL_FILE = "Bonsai-27B-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("⏳ মডেল লোড হচ্ছে...")
cpu_count = os.cpu_count() or 4


def load_model():
    """GPU/CPU অটো ডিটেক্ট"""
    try:
        print("🚀 GPU মোডে লোড করার চেষ্টা করছি...")
        llm = Llama(
            model_path=model_path,
            n_ctx=2048,
            n_gpu_layers=-1,
            n_threads=cpu_count,
            n_threads_batch=cpu_count,
            n_batch=512,
            verbose=False
        )
        print("✅ GPU মোডে সফলভাবে লোড হয়েছে!")
        return llm, "GPU"
    except Exception as e:
        print(f"⚠️ GPU লোড করতে ব্যর্থ: {e}")
        print("🔄 CPU-তে ফিরে যাচ্ছি...")

    try:
        print("🐢 CPU মোডে লোড হচ্ছে...")
        llm = Llama(
            model_path=model_path,
            n_ctx=2048,
            n_gpu_layers=0,
            n_threads=cpu_count,
            n_threads_batch=cpu_count,
            n_batch=512,
            verbose=False
        )
        print("✅ CPU মোডে সফলভাবে লোড হয়েছে!")
        return llm, "CPU"
    except Exception as e:
        print(f"❌ CPU-তেও লোড করতে ব্যর্থ: {e}")
        return None, None


llm, device_mode = load_model()

if llm is None:
    print("❌ মডেল লোড করা সম্ভব হয়নি!")
    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}")


# ===== ৪. চ্যাট জেনারেশন ফাংশন (ZeroGPU ডেকোরেটর সহ) =====

@spaces.GPU(duration=120)
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 = "🚀 GPU (দ্রুত)" if device_mode == "GPU" else "🐢 CPU (ধীর)"

with gr.Blocks(title="Bonsai 8B চ্যাট (ZeroGPU)") as demo:
    gr.HTML(f"""
    <div id="header-box">
        <h2>🌳 Bonsai 8B (১-বিট) চ্যাট</h2>
        <small>ZeroGPU • মডেল সাইজ: ~১.১৫ GB • {DEVICE_BADGE}</small>
    </div>
    """)

    chatbot = gr.Chatbot(
        value=[{"role": "assistant", "content": "👋 হ্যালো! আমি Bonsai 8B। আপনি কী জানতে চান?"}],
        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 8B (Q1_0, 1-bit)</div>"
    )

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