"""
HK UTM LLM Assistant — Hugging Face Spaces deployment
======================================================
Pure RAG pipeline: FAISS + sentence-transformers + Cerebras Inference API
No langchain dependency — avoids pydantic v1/v2 conflicts on Python 3.13.

Features:
  - Two-column layout: chat (left) + sources sidebar (right)
  - Streaming response (token-by-token output)
  - FAISS confidence scores + star ratings per source
  - Chunk preview in sidebar (first 200 chars of retrieved text)
  - Conversation memory (last 10 turns)
  - Eager pipeline load at startup (background thread)

Environment variables (set as HF Secrets):
  CEREBRAS_API_KEY : Your Cerebras Cloud API key (csk-...)
  HF_MODEL_ID : (optional) defaults to Qwen/Qwen2.5-72B-Instruct
"""

import os
import json
import threading
import numpy as np
import gradio as gr
from pathlib import Path

# ── Config ────────────────────────────────────────────────────────────────────
CEREBRAS_API_KEY = os.environ.get("CEREBRAS_API_KEY", "")
INDEX_DIR   = "data/processed/faiss_index"
DATA_DIR    = "data/raw"
EMBED_MODEL = "sentence-transformers/paraphrase-multilingual-MiniLM-L12-v2"
TOP_K        = 6
BM25_K       = 20
RRF_K        = 60
CE_POOL      = 6   # reduced from 12 → faster reranking, RRF pre-filters so quality impact minimal
# Upgraded from cross-encoder/ms-marco-MiniLM-L-6-v2 (English-only)
# mxbai-rerank-large-v2: Chinese benchmark 84.16, 100+ languages, 0.89s latency
# To revert: change back to "cross-encoder/ms-marco-MiniLM-L-6-v2"
# cross-encoder/ms-marco-MiniLM-L-6-v2: 22.7M params, fast, English-focused
# Stable and proven on HF Free CPU tier
CE_MODEL     = "cross-encoder/ms-marco-MiniLM-L-6-v2"

# ── Available LLM models (all free via HF Inference API) ─────────────────────
# Cerebras Inference API — OpenAI-compatible, ~2100 tok/s, 1M tokens/day free
LLM_MODELS = {
    "Llama-3.3-70B (Default ★)": {
        "id":   "llama-3.3-70b",
        "desc": "最佳質素 · 強大中英文推理 · Cerebras ~1s",
    },
    "Qwen3-32B": {
        "id":   "qwen-3-32b",
        "desc": "Qwen3 旗艦 · 思維模式 · 強中文 · ~1s",
    },
    "Llama-3.1-8B (Fast)": {
        "id":   "llama3.1-8b",
        "desc": "輕量快速 · 適合快速測試 · <1s",
    },
}
DEFAULT_MODEL_NAME = "Llama-3.3-70B (Default ★)"
MODEL_NAMES        = list(LLM_MODELS.keys())

UTM_SYSTEM_PROMPT = """You are an expert assistant in UAS Traffic Management (UTM) \
and U-space systems, specialising in Hong Kong and Greater Bay Area airspace operations.
You have deep knowledge of U-space services (U1-U3), Hong Kong CAD regulations, \
ICAO/FAA/SESAR UTM frameworks, strategic and tactical deconfliction, \
Demand-Capacity Balancing (DCB), and eVTOL operations.
LANGUAGE RULE (STRICT): \
If the user's question contains ANY Chinese characters, you MUST reply ENTIRELY in \
Traditional Chinese (繁體中文). This is MANDATORY. \
NEVER use Simplified Chinese (簡體字) under any circumstances — not even for proper nouns. \
If the user writes in English, respond in English.
Answer clearly and accurately. Always cite your source document.
If information is not in the provided context, say so explicitly."""

SIDEBAR_PLACEHOLDER = """*請輸入問題，以顯示已檢索的文件。*

---
### 三階段檢索流程

**第一階段 — 混合候選文件選取**
| 標記 | 方法 |
|---|---|
| 🔄 混合 | BM25 + FAISS 均命中 |
| 🧠 語義 | 僅 FAISS 向量搜尋 |
| 🔍 關鍵字 | 僅 BM25 精確配對 |

**第二階段 — 倒數排名融合**  
每種方法取前 20 → 合併成前 12 候選池

**第三階段 — 交叉編碼器重新排序**  
對 12 個候選項以「(查詢, 區塊)」配對評分  
→ 依 CE 分數選取最終前 6 名

*每個來源顯示 CE 分數，分數愈高代表愈相關。*

**知識庫：** 22 份文件 · 4,960 個區塊 · ICAO · FAA · SESAR · 香港民航處
"""


# ── Pipeline ──────────────────────────────────────────────────────────────────
_pipeline       = None
_pipeline_ready = False
_pipeline_error = None


def _build_index_from_pdfs(embed_model):
    import faiss
    from pypdf import PdfReader

    print("Building FAISS index from PDFs...")
    texts, metas = [], []
    chunk_size, overlap = 800, 100

    for fp in sorted(Path(DATA_DIR).rglob("*.pdf")):
        try:
            reader = PdfReader(str(fp))
            for page_num, page in enumerate(reader.pages):
                text = page.extract_text() or ""
                start = 0
                while start < len(text):
                    chunk = text[start:start + chunk_size].strip()
                    if len(chunk) >= 30 and any(c.isalpha() for c in chunk):
                        texts.append(chunk)
                        metas.append({
                            "content": chunk,
                            "source_file": fp.name,
                            "page": page_num,
                        })
                    start += chunk_size - overlap
        except Exception as e:
            print(f"Skip {fp.name}: {e}")

    print(f"Encoding {len(texts)} chunks...")
    embeddings = embed_model.encode(
        texts, batch_size=32, show_progress_bar=False,
        convert_to_numpy=True, normalize_embeddings=True
    )
    dim = embeddings.shape[1]
    index = faiss.IndexFlatIP(dim)
    index.add(embeddings.astype(np.float32))

    Path(INDEX_DIR).mkdir(parents=True, exist_ok=True)
    faiss.write_index(index, f"{INDEX_DIR}/index.faiss")
    with open(f"{INDEX_DIR}/metadata.json", "w", encoding="utf-8") as f:
        json.dump(metas, f, ensure_ascii=False)
    print(f"Index built: {len(texts)} chunks")
    return index, metas


def _load_pipeline():
    global _pipeline, _pipeline_ready, _pipeline_error
    try:
        import faiss
        from sentence_transformers import SentenceTransformer
        from openai import OpenAI  # Cerebras is OpenAI-compatible

        print("=== Loading pipeline at startup ===")
        print("Loading embedding model...")
        embed_model = SentenceTransformer(EMBED_MODEL)
        print("Embedding model loaded.")

        meta_path  = Path(f"{INDEX_DIR}/metadata.json")
        index_path = Path(f"{INDEX_DIR}/index.faiss")

        if index_path.exists() and meta_path.exists():
            print("Loading pre-built FAISS index...")
            index = faiss.read_index(str(index_path))
            with open(str(meta_path), "r", encoding="utf-8") as f:
                metadata = json.load(f)
            print(f"FAISS index loaded: {index.ntotal} vectors")
        else:
            index, metadata = _build_index_from_pdfs(embed_model)

        # ── BM25 keyword index ────────────────────────────────────────────────
        print("Building BM25 index...")
        from rank_bm25 import BM25Okapi
        corpus = [m["content"].lower().split() for m in metadata]
        bm25 = BM25Okapi(corpus)
        print(f"BM25 index built: {len(corpus)} docs")

        # ── Cross-Encoder reranker ────────────────────────────────────────
        print("Loading cross-encoder reranker...")
        from sentence_transformers import CrossEncoder
        cross_encoder = CrossEncoder(CE_MODEL, max_length=512)
        print("Cross-encoder (mxbai-rerank-large-v2) loaded.")

        print("Initialising LLM clients...")
        # Single Cerebras client — model is passed per-call
        cerebras_client = OpenAI(
            api_key=CEREBRAS_API_KEY,
            base_url="https://api.cerebras.ai/v1",
        )
        llm_clients = {name: cerebras_client for name in LLM_MODELS}
        print(f"Cerebras client ready: {list(llm_clients.keys())}")

        _pipeline = {
            "embed_model":   embed_model,
            "index":         index,
            "metadata":      metadata,
            "bm25":          bm25,
            "cross_encoder": cross_encoder,
            "llm_clients":   llm_clients,
            "history":       {},   # per-model: {model_name: [messages]}
        }
        _pipeline_ready = True
        print("=== Pipeline ready ===")

    except Exception as e:
        import traceback
        _pipeline_error = str(e)
        print(f"Pipeline load FAILED:\n{traceback.format_exc()}")


threading.Thread(target=_load_pipeline, daemon=True).start()


# ── Hybrid Retrieval: BM25 (keyword) + FAISS (semantic) via RRF ──────────────
def retrieve(query: str, k: int = TOP_K):
    p   = _pipeline
    N   = BM25_K
    K60 = RRF_K

    # 1. BM25 keyword ranking
    bm25_scores = p["bm25"].get_scores(query.lower().split())
    bm25_top    = sorted(enumerate(bm25_scores), key=lambda x: -x[1])[:N]
    bm25_ranks  = {idx: rank for rank, (idx, _) in enumerate(bm25_top)}

    # 2. FAISS semantic ranking
    q_emb = p["embed_model"].encode(
        [query], convert_to_numpy=True, normalize_embeddings=True
    ).astype(np.float32)
    _, faiss_idxs = p["index"].search(q_emb, N)
    faiss_ranks   = {int(i): rank for rank, i in enumerate(faiss_idxs[0]) if i >= 0}

    # 3. Reciprocal Rank Fusion (RRF) — get top CE_POOL candidates
    all_ids    = set(bm25_ranks) | set(faiss_ranks)
    rrf_scores = {
        doc_id: (
            1.0 / (K60 + bm25_ranks.get(doc_id,  N * 2)) +
            1.0 / (K60 + faiss_ranks.get(doc_id, N * 2))
        )
        for doc_id in all_ids
    }
    pool = sorted(rrf_scores.items(), key=lambda x: -x[1])[:CE_POOL]

    # 4. Cross-Encoder reranking on CE_POOL candidates
    pool_ids  = [doc_id for doc_id, _ in pool]
    pool_docs = [p["metadata"][doc_id] for doc_id in pool_ids]
    pairs     = [(query, d["content"]) for d in pool_docs]
    ce_scores = p["cross_encoder"].predict(pairs).tolist()

    # Sort by CE score, take final top-k
    ranked = sorted(zip(ce_scores, pool_ids), key=lambda x: -x[0])
    top_k  = ranked[:k]

    # Normalise CE scores to 70-100% display range
    ce_vals = [s for s, _ in top_k]
    ce_min, ce_max = min(ce_vals), max(ce_vals)

    # 5. Annotate with confidence % + retrieval method badges
    results = []
    for ce_s, doc_id in top_k:
        doc = dict(p["metadata"][doc_id])
        pct = int((ce_s - ce_min) / (ce_max - ce_min) * 30 + 70) if ce_max > ce_min else 85
        n   = max(1, min(5, round(pct / 20)))
        doc["relevance_pct"] = pct
        doc["ce_score"]       = round(ce_s, 2)
        doc["stars"]          = "★" * n + "☆" * (5 - n)
        in_b = doc_id in bm25_ranks
        in_f = doc_id in faiss_ranks
        doc["method"] = (
            "🔄 混合"   if in_b and in_f else
            "🧠 語義" if in_f else
            "🔍 關鍵字"
        )
        results.append(doc)
    return results


# ── Build sidebar markdown from retrieved docs ────────────────────────────────
def build_sidebar(docs: list) -> str:
    if not docs:
        return SIDEBAR_PLACEHOLDER

    lines = ["### 📚 已檢索來源\n"]
    seen = set()
    rank = 1
    for d in docs:
        key = f"{d['source_file']}::{d['page']}"
        if key in seen:
            continue
        seen.add(key)

        preview = d["content"].strip().replace("\n", " ")
        preview = preview[:220] + "…" if len(preview) > 220 else preview

        # Colour bar: green ≥85%, orange 70-84%, grey <70%
        pct = d["relevance_pct"]
        if pct >= 85:
            bar = "🟢"
        elif pct >= 70:
            bar = "🟡"
        else:
            bar = "⚪"

        method   = d.get("method", "")
        ce_score = d.get("ce_score", None)
        ce_str   = f" · CE `{ce_score}`" if ce_score is not None else ""
        lines.append(
            f"**{rank}. {d['stars']} {bar} `{pct}%`{ce_str} {method}**\n"
            f"📄 `{d['source_file']}`  \n"
            f"📖 Page {d['page']}\n\n"
            f"> {preview}\n"
        )
        lines.append("---")
        rank += 1

    return "\n".join(lines)


# ── Streaming chat handler ────────────────────────────────────────────────────
# Uses gr.State for sidebar text so Markdown is output-only (Gradio 5 requirement)
def chat(user_message: str, history: list, sidebar_state: str, model_name: str = DEFAULT_MODEL_NAME):
    if not user_message.strip():
        yield "", history, sidebar_state, sidebar_state
        return

    history = history or []

    if not _pipeline_ready:
        msg = ("❌ 管道錯誤：" + _pipeline_error) if _pipeline_error else (
            "⏳ 知識庫仍在載入中（約 1 分鐘），請稍後再試。"
        )
        history.append({"role": "user",      "content": user_message})
        history.append({"role": "assistant", "content": msg})
        yield "", history, sidebar_state, sidebar_state
        return

    # Retrieve + build sidebar immediately (before LLM call)
    docs = retrieve(user_message)
    new_sidebar = build_sidebar(docs)

    context_parts = [
        f"[Source: {d['source_file']}, Page: {d['page']}]\n{d['content']}"
        for d in docs
    ]
    context = "\n\n---\n\n".join(context_parts)

    # Detect if user message contains Chinese characters
    has_chinese = any('\u4e00' <= c <= '\u9fff' for c in user_message)
    lang_reminder = "\n\n【重要】請務必以繁體中文（Traditional Chinese）回答，嚴禁使用簡體中文。" if has_chinese else ""
    user_msg = (
        "Use the following UTM/U-space reference material to answer the question.\n\n"
        f"--- CONTEXT ---\n{context}\n--- END CONTEXT ---\n\n"
        f"Question: {user_message}"
        f"{lang_reminder}"
    )
    p          = _pipeline
    model_name = model_name or DEFAULT_MODEL_NAME
    llm_client = p["llm_clients"][model_name]
    model_hist = p["history"].setdefault(model_name, [])
    model_info = LLM_MODELS[model_name]

    messages = [{"role": "system", "content": UTM_SYSTEM_PROMPT}]
    for h in model_hist[-10:]:
        messages.append(h)
    messages.append({"role": "user", "content": user_msg})

    # Show thinking indicator immediately
    model_tag = f"\n\n---\n*模型：**{model_name}** — {model_info['desc']}*"
    history.append({"role": "user",      "content": user_message})
    history.append({"role": "assistant", "content": "⏳ 正在思考中…"})
    yield "", history, new_sidebar, new_sidebar

    # Non-streaming single call — more stable on mobile / weak connections
    try:
        response = llm_client.chat.completions.create(
            model=LLM_MODELS[model_name]["id"],
            messages=messages,
            max_tokens=1024,
            temperature=0.3,
        )
        answer = response.choices[0].message.content or ""
    except Exception as e:
        history[-1]["content"] = f"⚠️ 錯誤：{e}\n\n請重試。"
        yield "", history, new_sidebar, new_sidebar
        return

    if not answer.strip():
        history[-1]["content"] = "⚠️ 未能取得回答，請重試。"
        yield "", history, new_sidebar, new_sidebar
        return

    # Finalise — store in per-model history, append model tag
    model_hist.append({"role": "user",      "content": user_message})
    model_hist.append({"role": "assistant", "content": answer})
    history[-1]["content"] = answer + model_tag
    yield "", history, new_sidebar, new_sidebar


def reset_chat():
    if _pipeline:
        _pipeline["history"] = {}   # clear all per-model histories
    return [], SIDEBAR_PLACEHOLDER, SIDEBAR_PLACEHOLDER


def get_status():
    if _pipeline_error:
        return f"❌ 錯誤：{_pipeline_error}"
    if _pipeline_ready:
        return "✅ 就緒 — 4,960 個區塊 · 22 份文件 · 4 個模型可用"
    return "⏳ 正在載入知識庫…（首次啟動約需 1–2 分鐘）"


# ── Gradio UI — two-column layout ─────────────────────────────────────────────
css = """
#sidebar { border-left: 2px solid #e5e7eb; padding-left: 16px; }
#sidebar .prose blockquote {
    border-left: 3px solid #6b7280;
    padding: 6px 12px;
    margin: 4px 0;
    background: #f9fafb;
    font-size: 0.82em;
    color: #374151;
    border-radius: 4px;
}
#status-bar input { font-size: 0.85em; color: #6b7280; }

/* ── Mobile responsive ───────────────────────────────────────────── */
@media (max-width: 768px) {
    /* Stack chat + sidebar vertically on mobile */
    .equal_height > .gap {
        flex-direction: column !important;
    }
    /* Sidebar: remove left border, add top border instead */
    #sidebar {
        border-left: none !important;
        border-top: 2px solid #e5e7eb;
        padding-left: 0 !important;
        padding-top: 12px;
        margin-top: 8px;
    }
    /* Make model dropdown full width */
    .gr-dropdown { width: 100% !important; }
    /* Slightly smaller font in sidebar on mobile */
    #sidebar .prose { font-size: 0.82em; }
    /* Ensure chatbot takes full width */
    .chatbot { min-height: 300px !important; }
}
"""

with gr.Blocks(title="HK UTM LLM Assistant", theme=gr.themes.Soft(), css=css) as demo:

    # ── Header ────────────────────────────────────────────────────────────────
    gr.Markdown("""
    # ✈️ HK UTM LLM Assistant
    **檢索增強生成問答系統 · U-Space / 無人機交通管理 · 香港及國際框架**
    """)

    status_box = gr.Textbox(
        value=get_status,
        label="系統狀態",
        interactive=False,
        every=5,
        elem_id="status-bar",
    )

    # ── Model selector row ─────────────────────────────────────────────────────
    with gr.Row():
        with gr.Column(scale=5):
            model_dropdown = gr.Dropdown(
                choices=MODEL_NAMES,
                value=DEFAULT_MODEL_NAME,
                label="🤖 模型選擇（A/B 測試）",
                interactive=True,
            )
        with gr.Column(scale=7):
            model_desc_md = gr.Markdown(
                value=f"*{LLM_MODELS[DEFAULT_MODEL_NAME]['desc']}*",
                label="模型資訊",
            )

    # ── Main two-column area ──────────────────────────────────────────────────
    with gr.Row(equal_height=True):

        # Left column — chat (65% width)
        with gr.Column(scale=13):
            chatbot = gr.Chatbot(
                label="UTM 問答",
                height=500,
                bubble_full_width=False,
                type="messages",
                show_copy_button=True,
            )
            with gr.Row():
                msg_box = gr.Textbox(
                    placeholder="例如：香港民航處對無人機操作有何要求？",
                    label="您的問題",
                    scale=5,
                    autofocus=True,
                    lines=1,
                )
                send_btn = gr.Button("發送 ✈️", variant="primary", scale=1, min_width=100)

            reset_btn = gr.Button("🔄 新對話", variant="secondary", size="sm")

            gr.Examples(
                examples=[
                    "香港民航處對無人機系統操作有哪些主要要求？",
                    "U-space 的 U2 服務包含哪些內容？",
                    "戰略衝突解除與戰術衝突解除有何分別？",
                    "請解釋 UTM 中的需求容量平衡（DCB）。",
                    "USSP 在 U-space 生態系統中擔演什麼角色？",
                    "ICAO UTM 框架就互通性有何規定？",
                ],
                inputs=msg_box,
                label="範例問題",
            )

        # Right column — sources sidebar (35% width)
        # sidebar_state (gr.State) holds the text; sidebar_md (gr.Markdown) displays it
        with gr.Column(scale=7, elem_id="sidebar"):
            sidebar_md = gr.Markdown(
                value=SIDEBAR_PLACEHOLDER,
                label="已檢索來源",
                elem_id="sidebar",
            )

    # gr.State stores sidebar text between calls (Markdown is output-only in Gradio 5)
    sidebar_state = gr.State(value=SIDEBAR_PLACEHOLDER)

    # ── Footer ────────────────────────────────────────────────────────────────
    gr.Markdown("""
    ---
    *由 Gordon 建立 · 香港理工大學 AAE5302 · 民航資訊科技專業*  
    *由 Qwen2.5-72B · Llama-3.3-70B · Qwen3-8B · FAISS · sentence-transformers · Gradio 5 驅動*
    """)

    # ── Model description update ──────────────────────────────────────────────
    model_dropdown.change(
        fn=lambda m: f"*{LLM_MODELS[m]['desc']}*",
        inputs=[model_dropdown],
        outputs=[model_desc_md],
    )

    # ── Event wiring ──────────────────────────────────────────────────────────
    # outputs: msg_box, chatbot, sidebar_state (State), sidebar_md (Markdown display)
    send_btn.click(
        chat,
        inputs=[msg_box, chatbot, sidebar_state, model_dropdown],
        outputs=[msg_box, chatbot, sidebar_state, sidebar_md],
    )
    msg_box.submit(
        chat,
        inputs=[msg_box, chatbot, sidebar_state, model_dropdown],
        outputs=[msg_box, chatbot, sidebar_state, sidebar_md],
    )
    reset_btn.click(
        reset_chat,
        outputs=[chatbot, sidebar_state, sidebar_md],
    )


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