from flask import Flask, request, render_template_string, Response, stream_with_context
from llama_cpp import Llama
import os
import time
import json
import re
import sqlite3
from datetime import datetime

app = Flask(__name__)

# --- ১. সঠিক মডেল লোকেশন ---
# Q1_0 (এন্ড-টু-এন্ড ১-বিট) ফরম্যাট এখন mainline llama.cpp-তে upstream সাপোর্টেড,
# তাই llama-cpp-python (রিসেন্ট ভার্সন) দিয়ে এটা সরাসরি চলে।
# (শুধু তাদের ternary/Q2_0 ফরম্যাটের জন্য বিশেষ ফর্ক লাগে, Q1_0-এর জন্য না।)

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

# 🔽 ছোট মডেল চাইলে উপরের দুই লাইন কমেন্ট করে নিচের যেকোনো একটা আনকমেন্ট করুন 🔽

# --- Bonsai 4B (~0.57 GB) — মাঝারি সাইজ, আরও দ্রুত ---
#MODEL_REPO = "prism-ml/Bonsai-4B-gguf"
#MODEL_FILE = "Bonsai-4B-Q1_0.gguf"

# --- Bonsai 1.7B (~0.25 GB) — সবচেয়ে ছোট (১ বিলিয়নের কাছাকাছি), সবচেয়ে দ্রুত ---
#MODEL_REPO = "prism-ml/Bonsai-1.7B-gguf"
#MODEL_FILE = "Bonsai-1.7B-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("✅ ডাউনলোড সম্পূর্ণ!")
print(f"📁 ফাইল লোকেশন: {model_path}")

print("⏳ মডেল লোড হচ্ছে (CPU)...")
try:
    # llama-cpp-python: GGUF মেটাডেটা থেকে নিজেই আর্কিটেকচার (qwen3) বুঝে নেয়,
    # তাই model_type জাতীয় কিছু দিতে হয় না।

    # CPU-এর সব কোর ব্যবহার করা হচ্ছে যাতে prompt processing (prefill) দ্রুত হয়।
    # n_threads     -> টোকেন জেনারেশনের সময় ব্যবহৃত থ্রেড সংখ্যা
    # n_threads_batch -> prompt/prefill প্রসেসিংয়ের সময় ব্যবহৃত থ্রেড সংখ্যা (এটাই প্রথম টোকেনের দেরির মূল কারণ)
    cpu_count = os.cpu_count() or 4

    llm = Llama(
        model_path=model_path,
        n_ctx=6144,
        n_threads=cpu_count,
        n_threads_batch=cpu_count,
        n_batch=512,       # prefill ব্যাচ সাইজ, বড় করলে prompt processing দ্রুত হয়
        verbose=False
    )

    # --- ওয়ার্মআপ ---
    # প্রথম ইনফারেন্স কলে llama.cpp কিছু অভ্যন্তরীণ বাফার/গ্রাফ তৈরি করে যা এক্সট্রা সময় নেয়।
    # সার্ভার চালু হওয়ার সময়ই একটা ডামি জেনারেশন চালিয়ে সেই ওয়ান-টাইম খরচ আগেই সেরে ফেলা হচ্ছে,
    # যাতে ইউজারের প্রথম আসল রিকোয়েস্টে এই পেনাল্টি না লাগে।
    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("✅ ওয়ার্মআপ সম্পন্ন!")

    print("✅ মডেল প্রস্তুত! সার্ভার চালু হচ্ছে...")
except Exception as e:
    print(f"⚠️ llama-cpp-python লোড করতে সমস্যা: {e}")
    llm = None

# --- ২. ডাটাবেজ (SQLite) ---
# ব্যবহারকারীর তথ্য (facts) সংরক্ষণের জন্য একটা সাধারণ SQLite ডাটাবেজ।
# মডেল নিজেই সিদ্ধান্ত নেয় কখন সংরক্ষণ/আপডেট/ডিলিট/সার্চ করতে হবে — এখানে কোনো
# কীওয়ার্ড বা প্যাটার্ন ম্যাচিং নেই, শুধু মডেলের টুল-কল অনুযায়ী কাজ করা হয়।
DB_PATH = os.path.join(os.path.dirname(os.path.abspath(__file__)), "memory.db")

def db_init():
    conn = sqlite3.connect(DB_PATH)
    conn.execute("""
        CREATE TABLE IF NOT EXISTS facts (
            id INTEGER PRIMARY KEY AUTOINCREMENT,
            topic TEXT NOT NULL,
            content TEXT NOT NULL,
            created_at TEXT NOT NULL,
            updated_at TEXT NOT NULL
        )
    """)
    conn.commit()
    conn.close()

db_init()

def _now():
    return datetime.now().strftime("%Y-%m-%d %H:%M:%S")

def db_find_best_match(topic):
    """topic-এর সাথে সবচেয়ে কাছাকাছি মিল থাকা রেকর্ড খুঁজে বের করে (সাবস্ট্রিং ম্যাচ)।"""
    conn = sqlite3.connect(DB_PATH)
    conn.row_factory = sqlite3.Row
    row = conn.execute(
        "SELECT * FROM facts WHERE topic LIKE ? ORDER BY updated_at DESC LIMIT 1",
        (f"%{topic}%",)
    ).fetchone()
    conn.close()
    return dict(row) if row else None

def tool_save_fact(topic, content):
    existing = db_find_best_match(topic)
    conn = sqlite3.connect(DB_PATH)
    if existing:
        conn.execute(
            "UPDATE facts SET topic=?, content=?, updated_at=? WHERE id=?",
            (topic, content, _now(), existing["id"])
        )
        conn.commit()
        conn.close()
        return {"status": "already_existed_so_updated", "topic": topic, "content": content}
    conn.execute(
        "INSERT INTO facts (topic, content, created_at, updated_at) VALUES (?, ?, ?, ?)",
        (topic, content, _now(), _now())
    )
    conn.commit()
    conn.close()
    return {"status": "created", "topic": topic, "content": content}

def tool_update_fact(topic, content):
    existing = db_find_best_match(topic)
    conn = sqlite3.connect(DB_PATH)
    if not existing:
        conn.execute(
            "INSERT INTO facts (topic, content, created_at, updated_at) VALUES (?, ?, ?, ?)",
            (topic, content, _now(), _now())
        )
        conn.commit()
        conn.close()
        return {"status": "not_found_so_created", "topic": topic, "content": content}
    conn.execute(
        "UPDATE facts SET topic=?, content=?, updated_at=? WHERE id=?",
        (topic, content, _now(), existing["id"])
    )
    conn.commit()
    conn.close()
    return {"status": "updated", "old_content": existing["content"], "topic": topic, "content": content}

def tool_delete_fact(topic):
    existing = db_find_best_match(topic)
    if not existing:
        return {"status": "not_found", "topic": topic}
    conn = sqlite3.connect(DB_PATH)
    conn.execute("DELETE FROM facts WHERE id=?", (existing["id"],))
    conn.commit()
    conn.close()
    return {"status": "deleted", "topic": existing["topic"], "content": existing["content"]}

def _as_list(value):
    """যেকোনো ইনপুটকে (single dict/string বা list) নিরাপদে একটা list-এ রূপান্তর করে।"""
    if isinstance(value, list):
        return value
    if isinstance(value, dict) or isinstance(value, str):
        return [value]
    return []

def tool_search_facts(query):
    conn = sqlite3.connect(DB_PATH)
    conn.row_factory = sqlite3.Row
    rows = conn.execute(
        "SELECT topic, content, updated_at FROM facts WHERE topic LIKE ? OR content LIKE ? ORDER BY updated_at DESC LIMIT 10",
        (f"%{query}%", f"%{query}%")
    ).fetchall()
    conn.close()
    if not rows:
        return {"status": "not_found", "query": query, "results": []}
    return {"status": "found", "query": query, "results": [dict(r) for r in rows]}

def tool_list_all_facts():
    """ব্যবহারকারী সম্পর্কে সংরক্ষিত সব তথ্য ফেরত দেয় (ব্যাপক/সামগ্রিক প্রশ্নের জন্য)।"""
    conn = sqlite3.connect(DB_PATH)
    conn.row_factory = sqlite3.Row
    rows = conn.execute(
        "SELECT topic, content, updated_at FROM facts ORDER BY updated_at DESC LIMIT 100"
    ).fetchall()
    conn.close()
    if not rows:
        return {"status": "empty", "results": []}
    return {"status": "ok", "count": len(rows), "results": [dict(r) for r in rows]}

TOOLS_SPEC = [
    {
        "type": "function",
        "function": {
            "name": "save_facts",
            "description": "Save one or more new personal facts about the user. If there are multiple facts, put them all in one call.",
            "parameters": {
                "type": "object",
                "properties": {
                    "items": {
                        "type": "array",
                        "description": "List of facts to save",
                        "items": {
                            "type": "object",
                            "properties": {
                                "topic": {"type": "string", "description": "Short label for the fact"},
                                "content": {"type": "string", "description": "The fact itself"}
                            },
                            "required": ["topic", "content"]
                        }
                    }
                },
                "required": ["items"]
            }
        }
    },
    {
        "type": "function",
        "function": {
            "name": "update_facts",
            "description": "Update one or more existing facts. Batch multiple updates into one call.",
            "parameters": {
                "type": "object",
                "properties": {
                    "items": {
                        "type": "array",
                        "description": "List of facts to update",
                        "items": {
                            "type": "object",
                            "properties": {
                                "topic": {"type": "string", "description": "Label of the fact to update"},
                                "content": {"type": "string", "description": "The new value"}
                            },
                            "required": ["topic", "content"]
                        }
                    }
                },
                "required": ["items"]
            }
        }
    },
    {
        "type": "function",
        "function": {
            "name": "delete_facts",
            "description": "Delete one or more stored facts. Batch multiple topics into one call.",
            "parameters": {
                "type": "object",
                "properties": {
                    "topics": {
                        "type": "array",
                        "description": "Labels of the facts to delete",
                        "items": {"type": "string"}
                    }
                },
                "required": ["topics"]
            }
        }
    },
    {
        "type": "function",
        "function": {
            "name": "search_facts",
            "description": "Search stored facts for one or more topics. Batch multiple queries into one call.",
            "parameters": {
                "type": "object",
                "properties": {
                    "queries": {
                        "type": "array",
                        "description": "List of topics/keywords to search for",
                        "items": {"type": "string"}
                    }
                },
                "required": ["queries"]
            }
        }
    },
    {
        "type": "function",
        "function": {
            "name": "list_all_facts",
            "description": "Return every fact stored about the user. Use for broad questions like 'what do you know about me', or when search_facts finds nothing. No arguments needed.",
            "parameters": {
                "type": "object",
                "properties": {},
                "required": []
            }
        }
    }
]
TOOLS_JSON = json.dumps(TOOLS_SPEC, ensure_ascii=False, indent=2)

def exec_save_facts(args):
    items = _as_list(args.get("items", []))
    results = []
    for it in items:
        if not isinstance(it, dict):
            continue
        topic = str(it.get("topic", "")).strip()
        content = str(it.get("content", "")).strip()
        if topic and content:
            results.append(tool_save_fact(topic, content))
    return {"status": "ok", "count": len(results), "results": results}

def exec_update_facts(args):
    items = _as_list(args.get("items", []))
    results = []
    for it in items:
        if not isinstance(it, dict):
            continue
        topic = str(it.get("topic", "")).strip()
        content = str(it.get("content", "")).strip()
        if topic and content:
            results.append(tool_update_fact(topic, content))
    return {"status": "ok", "count": len(results), "results": results}

def exec_delete_facts(args):
    topics = _as_list(args.get("topics", []))
    results = []
    for t in topics:
        topic = t.get("topic") if isinstance(t, dict) else t
        topic = str(topic or "").strip()
        if topic:
            results.append(tool_delete_fact(topic))
    return {"status": "ok", "count": len(results), "results": results}

def exec_search_facts(args):
    queries = _as_list(args.get("queries", []))
    results_by_query = {}
    for q in queries:
        q = str(q or "").strip()
        if q:
            results_by_query[q] = tool_search_facts(q)
    return {"status": "ok", "results_by_query": results_by_query}

TOOL_EXECUTORS = {
    "save_facts": exec_save_facts,
    "update_facts": exec_update_facts,
    "delete_facts": exec_delete_facts,
    "search_facts": exec_search_facts,
    "list_all_facts": lambda args: tool_list_all_facts(),
    # নিচেরগুলো ব্যাকওয়ার্ড-কম্প্যাটিবিলিটির জন্য — মডেল যদি কখনো একবচন নাম দিয়ে কল করে
    # (একটা মাত্র আইটেম নিয়ে) সেটাও যেন কাজ করে, ভুল হলে পুরো উত্তর যেন ব্যর্থ না হয়।
    "save_fact": lambda args: exec_save_facts({"items": [args]}),
    "update_fact": lambda args: exec_update_facts({"items": [args]}),
    "delete_fact": lambda args: exec_delete_facts({"topics": [args.get("topic", "")]}),
}

def execute_tool_call(name, arguments):
    executor = TOOL_EXECUTORS.get(name)
    if not executor:
        return {"error": f"unknown tool: {name}"}
    try:
        return executor(arguments or {})
    except Exception as e:
        return {"error": str(e)}

TOOL_CALL_RE = re.compile(r"<tool_call>\s*(\{.*?\})\s*</tool_call>", re.DOTALL)

SYSTEM_PROMPT = f"""You are a helpful assistant with a personal facts database about the user.

# Tools
<tools>
{TOOLS_JSON}
</tools>

Rules:
- New personal info -> save_facts (batch ALL items into one call, not one call per fact).
- Info to change -> update_facts (batch).
- Info to remove -> delete_facts (batch).
- Specific question -> search_facts (batch all queries into one call).
- Broad question ("what do you know about me") or search_facts finds nothing -> list_all_facts.
- No database need -> just answer directly, no tool call.
- Always merge multiple items into ONE tool_call using arrays. Only use separate blocks for genuinely different tools in the same turn.

To call a tool, output exactly this format (nothing else):
<tool_call>
{{"name": "...", "arguments": {{...}}}}
</tool_call>

Example (user gives several facts at once):
<tool_call>
{{"name": "save_facts", "arguments": {{"items": [
  {{"topic": "name", "content": "Mukul"}},
  {{"topic": "age", "content": "35"}}
]}}}}
</tool_call>

After a tool call you will receive results in <tool_response> tags. Then give a short, natural, human-sounding final answer, in the same language the user wrote in. Never show JSON, tags, function names, or the word "database" in your final answer."""

def build_base_prompt(history, user_prompt):
    """সিস্টেম প্রম্পট + কথোপকথনের হিস্টরি + বর্তমান প্রশ্ন দিয়ে সম্পূর্ণ প্রম্পট তৈরি করে।"""
    prompt = f"<|im_start|>system\n{SYSTEM_PROMPT}<|im_end|>\n"
    for turn in history:
        h_user = str(turn.get('user', '')).strip()
        h_assistant = str(turn.get('assistant', '')).strip()
        if not h_user:
            continue
        prompt += f"<|im_start|>user\n{h_user}<|im_end|>\n<|im_start|>assistant\n<think>\n\n</think>\n\n{h_assistant}<|im_end|>\n"
    prompt += f"<|im_start|>user\n{user_prompt}<|im_end|>\n<|im_start|>assistant\n<think>\n\n</think>\n\n"
    return prompt

# --- ৩. HTML টেমপ্লেট ---
HTML_TEMPLATE = """
<!DOCTYPE html>
<html>
<head>
    <title>Bonsai 8B চ্যাট (Hugging Face Space)</title>
    <meta charset="UTF-8">
    <meta name="viewport" content="width=device-width, initial-scale=1.0">
    <style>
        * { box-sizing: border-box; }
        body {
            font-family: 'Segoe UI', Tahoma, sans-serif;
            max-width: 900px;
            margin: 20px auto;
            padding: 15px;
            background: #f0f2f5;
        }
        .container {
            background: white;
            border-radius: 16px;
            box-shadow: 0 2px 10px rgba(0,0,0,0.1);
            overflow: hidden;
            padding: 20px;
        }
        .header {
            background: linear-gradient(135deg, #667eea 0%, #764ba2 100%);
            color: white;
            padding: 15px 20px;
            margin: -20px -20px 20px -20px;
            border-radius: 16px 16px 0 0;
        }
        .header h2 { margin: 0; font-weight: 400; }
        .header small { opacity: 0.8; font-size: 14px; }

        .chat-box {
            background: #f8f9fa;
            border-radius: 12px;
            padding: 15px;
            height: 400px;
            overflow-y: auto;
            margin-bottom: 15px;
            border: 1px solid #e0e0e0;
        }
        .message {
            margin: 8px 0;
            padding: 10px 15px;
            border-radius: 18px;
            max-width: 80%;
            word-wrap: break-word;
            animation: fadeIn 0.3s ease;
        }
        @keyframes fadeIn {
            from { opacity: 0; transform: translateY(10px); }
            to { opacity: 1; transform: translateY(0); }
        }
        .user-msg {
            background: #007bff;
            color: white;
            margin-left: auto;
            text-align: right;
            border-bottom-right-radius: 4px;
        }
        .bot-msg {
            background: #e9ecef;
            color: #333;
            margin-right: auto;
            border-bottom-left-radius: 4px;
            white-space: pre-wrap;
        }
        .bot-msg.thinking {
            background: #e9ecef;
            color: #666;
            font-style: italic;
        }
        .time-info {
            font-size: 12px;
            color: #4a5568;
            margin: 4px 4px 10px 4px;
            text-align: left;
            background: #f7fafc;
            padding: 6px 12px;
            border-radius: 8px;
            border-left: 3px solid #667eea;
            display: inline-block;
        }
        .time-info strong {
            color: #2d3748;
        }
        .input-area {
            display: flex;
            gap: 10px;
        }
        #userInput {
            flex: 1;
            padding: 12px 18px;
            border: 2px solid #ddd;
            border-radius: 25px;
            font-size: 15px;
            outline: none;
            transition: 0.2s;
        }
        #userInput:focus {
            border-color: #667eea;
        }
        button {
            padding: 12px 28px;
            background: #667eea;
            color: white;
            border: none;
            border-radius: 25px;
            font-size: 15px;
            cursor: pointer;
            transition: 0.2s;
            white-space: nowrap;
        }
        button:hover {
            background: #5a67d8;
            transform: scale(1.02);
        }
        button:disabled {
            background: #a0aec0;
            cursor: not-allowed;
            transform: none;
        }
        .status {
            color: #718096;
            font-size: 13px;
            margin-top: 10px;
            display: flex;
            align-items: center;
            gap: 8px;
        }
        .loader {
            display: inline-block;
            width: 14px;
            height: 14px;
            border: 2px solid #e2e8f0;
            border-top: 2px solid #667eea;
            border-radius: 50%;
            animation: spin 0.8s linear infinite;
        }
        @keyframes spin {
            0% { transform: rotate(0deg); }
            100% { transform: rotate(360deg); }
        }
        .header {
            position: relative;
        }
        .menu-wrapper {
            position: absolute;
            top: 15px;
            right: 20px;
        }
        .menu-btn {
            background: rgba(255,255,255,0.15);
            border: none;
            color: white;
            width: 34px;
            height: 34px;
            padding: 0;
            border-radius: 50%;
            font-size: 20px;
            line-height: 1;
            cursor: pointer;
        }
        .menu-btn:hover {
            background: rgba(255,255,255,0.28);
            transform: none;
        }
        .menu-dropdown {
            display: none;
            position: absolute;
            top: 42px;
            right: 0;
            background: white;
            border-radius: 10px;
            box-shadow: 0 4px 16px rgba(0,0,0,0.2);
            overflow: hidden;
            min-width: 160px;
            z-index: 10;
        }
        .menu-dropdown.open {
            display: block;
        }
        .menu-item {
            padding: 12px 16px;
            color: #333;
            font-size: 14px;
            cursor: pointer;
            white-space: nowrap;
        }
        .menu-item:hover {
            background: #f0f2f5;
        }
        .footer {
            margin-top: 15px;
            text-align: center;
            color: #a0aec0;
            font-size: 12px;
        }
        .first-token-badge {
            background: #48bb78;
            color: white;
            padding: 2px 10px;
            border-radius: 12px;
            font-size: 11px;
            margin-left: 8px;
        }
    </style>
</head>
<body>
    <div class="container">
        <div class="header">
            <h2>🌳 Bonsai 8B (১-বিট) চ্যাট</h2>
            <small>Hugging Face Space • CPU Inference • মডেল সাইজ: ~১.১৫ GB (Q1_0, 1-bit)</small>
            <div class="menu-wrapper">
                <button id="menuBtn" class="menu-btn">⋮</button>
                <div id="menuDropdown" class="menu-dropdown">
                    <div id="newChatItem" class="menu-item">➕ নতুন চ্যাট</div>
                </div>
            </div>
        </div>

        <div class="chat-box" id="chatBox">
            <div class="message bot-msg">👋 হ্যালো! আমি Bonsai 8B। আপনি কী জানতে চান?</div>
        </div>

        <div class="input-area">
            <input type="text" id="userInput" placeholder="এখানে প্রশ্ন লিখুন..." />
            <button id="sendBtn">পাঠান</button>
        </div>
        <div class="status" id="status">
            <span>✅</span> প্রস্তুত
        </div>
        <div class="footer">
            Powered by llama-cpp-python • Bonsai 8B (Q1_0, 1-bit)
        </div>
    </div>

    <script>
        const chatBox = document.getElementById('chatBox');
        const input = document.getElementById('userInput');
        const sendBtn = document.getElementById('sendBtn');
        const status = document.getElementById('status');
        const menuBtn = document.getElementById('menuBtn');
        const menuDropdown = document.getElementById('menuDropdown');
        const newChatItem = document.getElementById('newChatItem');

        // সর্বশেষ ৩টা প্রশ্ন-উত্তরের হিস্টরি রাখা হচ্ছে (মডেলকে কনটেক্সট দেওয়ার জন্য)
        let conversationHistory = [];

        function addMessage(text, isUser = false, isThinking = false) {
            const div = document.createElement('div');
            div.className = `message ${isUser ? 'user-msg' : 'bot-msg'}${isThinking ? ' thinking' : ''}`;
            div.textContent = text;
            chatBox.appendChild(div);
            chatBox.scrollTop = chatBox.scrollHeight;
            return div;
        }

        function addTimeInfo(firstTokenTime, totalTime, startTimeStr, endTimeStr) {
            const div = document.createElement('div');
            div.className = 'time-info';
            div.innerHTML = `
                <strong>⏱️ প্রথম টোকেন:</strong> ${firstTokenTime} সেকেন্ড &nbsp;|&nbsp;
                <strong>মোট সময়:</strong> ${totalTime} সেকেন্ড &nbsp;|&nbsp;
                <strong>শুরু:</strong> ${startTimeStr} &nbsp;|&nbsp;
                <strong>শেষ:</strong> ${endTimeStr}
            `;
            chatBox.appendChild(div);
            chatBox.scrollTop = chatBox.scrollHeight;
        }

        async function sendMessage() {
            const text = input.value.trim();
            if (!text) return;

            // সেন্ড বাটন ক্লিক করার মুহূর্ত থেকেই টাইমার শুরু
            const startTime = Date.now();

            addMessage(text, true);
            input.value = '';
            sendBtn.disabled = true;
            status.innerHTML = '<span class="loader"></span> প্রক্রিয়াকরণ শুরু হচ্ছে...';

            // থিংকিং ইন্ডিকেটর
            const thinkingDiv = addMessage('⏳ উত্তর তৈরি হচ্ছে...', false, true);

            try {
                const response = await fetch('/chat', {
                    method: 'POST',
                    headers: { 'Content-Type': 'application/json' },
                    body: JSON.stringify({ prompt: text, history: conversationHistory })
                });

                if (!response.ok || !response.body) {
                    throw new Error('স্ট্রিম শুরু করা যায়নি');
                }

                // থিংকিং ইন্ডিকেটর সরানো
                thinkingDiv.remove();
                const botDiv = addMessage('', false);

                const reader = response.body.getReader();
                const decoder = new TextDecoder();
                let fullText = '';
                let buffer = '';
                let metaFound = false;
                let firstToken = true;
                let firstTokenTime = 0;

                while (true) {
                    const { done, value } = await reader.read();
                    if (done) break;
                    buffer += decoder.decode(value, { stream: true });

                    // টুল-কল স্ট্যাটাস মার্কার হ্যান্ডলিং (ডাটাবেজ চেক/সংরক্ষণ চলাকালীন)
                    const toolStartIdx = buffer.indexOf('[[TOOL_START]]');
                    if (toolStartIdx !== -1) {
                        fullText += buffer.slice(0, toolStartIdx);
                        botDiv.textContent = fullText;
                        buffer = buffer.slice(toolStartIdx + '[[TOOL_START]]'.length);
                        status.innerHTML = '<span class="loader"></span> 🗄️ ডাটাবেজ চেক করা হচ্ছে...';
                    }
                    const toolDoneIdx = buffer.indexOf('[[TOOL_DONE]]');
                    if (toolDoneIdx !== -1) {
                        buffer = buffer.slice(toolDoneIdx + '[[TOOL_DONE]]'.length);
                        status.innerHTML = '<span class="loader"></span> ✍️ উত্তর লেখা হচ্ছে...';
                    }

                    const metaIndex = buffer.indexOf('[[META]]');
                    if (metaIndex !== -1) {
                        fullText += buffer.slice(0, metaIndex);
                        botDiv.textContent = fullText;
                        buffer = buffer.slice(metaIndex);
                        metaFound = true;
                    } else if (!metaFound) {
                        fullText += buffer;
                        botDiv.textContent = fullText;
                        buffer = '';
                        
                        // প্রথম টোকেন আসার সময় রেকর্ড
                        if (firstToken && fullText.length > 0) {
                            firstToken = false;
                            firstTokenTime = (Date.now() - startTime) / 1000;
                            status.innerHTML = `✍️ উত্তর লেখা হচ্ছে... <span class="first-token-badge">প্রথম টোকেন: ${firstTokenTime.toFixed(2)}s</span>`;
                        }
                    }
                    chatBox.scrollTop = chatBox.scrollHeight;
                }

                // মেটাডেটা পার্স করা
                const metaIndex = buffer.indexOf('[[META]]');
                if (metaIndex !== -1) {
                    try {
                        const meta = JSON.parse(buffer.slice(metaIndex + 8));
                        // firstTokenTime যদি সেট না হয়ে থাকে (খুব ছোট রেসপন্স)
                        if (firstTokenTime === 0) {
                            firstTokenTime = meta.first_token_time || 0;
                        }
                        addTimeInfo(
                            firstTokenTime.toFixed(2),
                            meta.elapsed.toFixed(2),
                            meta.start,
                            meta.end
                        );
                    } catch (e) {
                        // মেটাডেটা পার্স করতে ব্যর্থ হলে চুপচাপ বাদ দেওয়া হলো
                    }
                }

                chatBox.scrollTop = chatBox.scrollHeight;
                status.innerHTML = '✅ প্রস্তুত';

                // এই এক্সচেঞ্জটা হিস্টরিতে যোগ করা এবং শুধু সর্বশেষ ৩টা রাখা
                conversationHistory.push({ user: text, assistant: fullText.trim() });
                if (conversationHistory.length > 3) {
                    conversationHistory = conversationHistory.slice(-3);
                }
            } catch (error) {
                thinkingDiv.remove();
                addMessage('⚠️ সার্ভার ত্রুটি! আবার চেষ্টা করুন।', false);
                status.innerHTML = '❌ সংযোগ ত্রুটি';
                console.error('Error:', error);
            }
            sendBtn.disabled = false;
            input.focus();
        }

        function newChat() {
            chatBox.innerHTML = '';
            conversationHistory = [];
            addMessage('👋 হ্যালো! আমি Bonsai 8B। আপনি কী জানতে চান?', false);
            status.innerHTML = '✅ প্রস্তুত';
            menuDropdown.classList.remove('open');
            input.focus();
        }

        sendBtn.addEventListener('click', sendMessage);
        input.addEventListener('keypress', (e) => {
            if (e.key === 'Enter') sendMessage();
        });

        // থ্রি-ডট মেনু টগল
        menuBtn.addEventListener('click', (e) => {
            e.stopPropagation();
            menuDropdown.classList.toggle('open');
        });
        newChatItem.addEventListener('click', newChat);
        // মেনুর বাইরে ক্লিক করলে বন্ধ হয়ে যাবে
        document.addEventListener('click', () => {
            menuDropdown.classList.remove('open');
        });

        // Auto-focus on page load
        input.focus();
    </script>
</body>
</html>
"""

# --- ৪. Flask রাউট ---
@app.route('/')
def index():
    return render_template_string(HTML_TEMPLATE)

@app.route('/chat', methods=['POST'])
def chat():
    data = request.get_json()
    user_prompt = data.get('prompt', '')
    # ফ্রন্টএন্ড থেকে পাঠানো সর্বশেষ (সর্বোচ্চ ৩টা) প্রশ্ন-উত্তরের হিস্টরি
    # প্রতিটা আইটেম: {"user": "...", "assistant": "..."}
    history = data.get('history', [])

    if not user_prompt:
        return {"error": "কোনো প্রশ্ন নেই"}

    if llm is None:
        return {"error": "মডেল লোড হয়নি। সার্ভার লগ চেক করুন।"}

    # শুধু সর্বশেষ ৩টা এক্সচেঞ্জ ব্যবহার করা হচ্ছে (কনটেক্সট উইন্ডো সীমিত রাখতে)
    if isinstance(history, list):
        history = history[-3:]
    else:
        history = []

    # হিস্টরি + সিস্টেম প্রম্পট (টুল সংজ্ঞাসহ) দিয়ে প্রম্পট তৈরি
    prompt = build_base_prompt(history, user_prompt)

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

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

            # প্রথম কয়েক অক্ষর বাফার করে দেখা হচ্ছে মডেল টুল কল করছে নাকি সরাসরি উত্তর দিচ্ছে —
            # এটা কোনো কীওয়ার্ড/প্যাটার্ন ম্যাচিং না, মডেল নিজেই <tool_call> ট্যাগ দিয়ে
            # শুরু করে কিনা সেটা শুধু চেক করা হচ্ছে (এটাই মডেলের নিজস্ব ফাংশন-কলিং ফরম্যাট)।
            head_buffer = ""
            head_decided = False
            is_tool_mode = False
            tool_call_text = ""
            TOOL_TAG = "<tool_call>"

            for chunk in stream1:
                token_text = chunk["choices"][0]["text"]
                if not token_text:
                    continue
                if first_token_time is None:
                    first_token_time = time.time() - start_time

                if not head_decided:
                    head_buffer += token_text
                    stripped = head_buffer.lstrip()
                    if not stripped:
                        continue
                    if stripped.startswith(TOOL_TAG):
                        head_decided = True
                        is_tool_mode = True
                        tool_call_text = head_buffer
                        yield "[[TOOL_START]]"
                    elif len(stripped) >= len(TOOL_TAG) or not TOOL_TAG.startswith(stripped):
                        head_decided = True
                        is_tool_mode = False
                        yield head_buffer
                    continue

                if is_tool_mode:
                    tool_call_text += token_text
                else:
                    yield token_text

            if not head_decided:
                stripped = head_buffer.lstrip()
                if stripped.startswith(TOOL_TAG):
                    is_tool_mode = True
                    tool_call_text = head_buffer
                    yield "[[TOOL_START]]"
                elif head_buffer:
                    yield head_buffer

            if is_tool_mode:
                # মডেলের টুল কল(গুলো) পার্স ও এক্সিকিউট করা
                tool_calls = TOOL_CALL_RE.findall(tool_call_text)
                tool_response_blocks = ""
                for tc_raw in tool_calls:
                    try:
                        tc = json.loads(tc_raw)
                        name = tc.get("name", "")
                        args = tc.get("arguments", {}) or {}
                        result = execute_tool_call(name, args)
                    except Exception as e:
                        result = {"error": f"failed to parse tool call: {str(e)}"}
                    tool_response_blocks += (
                        "<|im_start|>tool\n<tool_response>\n"
                        f"{json.dumps(result, ensure_ascii=False)}\n"
                        "</tool_response><|im_end|>\n"
                    )

                yield "[[TOOL_DONE]]"

                follow_prompt = (
                    prompt + tool_call_text.strip() + "<|im_end|>\n"
                    + tool_response_blocks
                    + "<|im_start|>assistant\n<think>\n\n</think>\n\n"
                )

                stream2 = llm(
                    follow_prompt,
                    max_tokens=100,
                    temperature=0.1,
                    top_p=0.2,
                    top_k=20,
                    repeat_penalty=1.4,
                    stop=["<|im_end|>", "<tool_call>", "user:", "User:"],
                    stream=True
                )
                for chunk in stream2:
                    token_text = chunk["choices"][0]["text"]
                    if token_text:
                        yield token_text

        except Exception as e:
            print(f"❌ জেনারেশন ত্রুটি: {e}")
            yield f"\n⚠️ মডেল ত্রুটি: {str(e)}"
        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

            print(f"⏱️ প্রথম টোকেন: {first_token} সেকেন্ড | মোট সময়: {elapsed} সেকেন্ড")

            # স্ট্রিমের একদম শেষে টাইমিং তথ্য পাঠানো
            meta = {
                "start": start_str,
                "end": end_str,
                "elapsed": elapsed,
                "first_token_time": first_token
            }
            yield f"\n[[META]]{json.dumps(meta, ensure_ascii=False)}"

    return Response(stream_with_context(generate()), mimetype='text/plain')

# --- ৫. সার্ভার চালানো ---
if __name__ == '__main__':
    app.run(host='0.0.0.0', port=7860)