"""
多策略 RAG 文件問答系統 v2 — ChromaDB + PDF/DOCX 版本（優化版）

安裝依賴：
    pip install gradio groq pypdf python-docx sentence-transformers numpy chromadb scikit-learn

執行：
    python multistrategy_rag_chromadb_docx_v2.py
"""

import chromadb
import gradio as gr
import numpy as np
import os
import time
import re
from docx import Document
from docx.oxml.table import CT_Tbl
from docx.oxml.text.paragraph import CT_P
from docx.table import Table
from docx.text.paragraph import Paragraph
from groq import Groq
from pypdf import PdfReader
from sentence_transformers import SentenceTransformer
from sklearn.feature_extraction.text import TfidfVectorizer
from pathlib import Path
from typing import Any

# ══════════════════════════════════════════════════════════
#  RAG 核心邏輯（優化版）
# ══════════════════════════════════════════════════════════
class MultiStrategyRAG:

    STRATEGY_MAP = {
        "semantic":      "1  ChromaDB 語意搜尋",
        "tfidf":         "2  TF-IDF 關鍵詞",
        "hybrid":        "3  混合搜尋",
        "rerank":        "4  重新排序",
        "multi_query":   "5  多查詢擴展",
        "compress":      "6  上下文壓縮",
        "parent_child":  "7  父子文檔",
        "hyde":          "8  假設性答案 HyDE",
    }

    def __init__(
        self,
        chroma_path: str = "/tmp/chroma_db",
        collection_name: str = "audit_rag_chunks",
        child_collection_name: str = "audit_rag_child_chunks",
    ):
        # API client 改為 None，由使用者透過 UI 輸入後動態建立
        self.client: Groq | None = None

        self.embedding_model = SentenceTransformer(
            "sentence-transformers/paraphrase-multilingual-MiniLM-L12-v2"
        )

        self.chroma_client = chromadb.PersistentClient(path=chroma_path)
        self.collection = self.chroma_client.get_or_create_collection(
            name=collection_name,
            metadata={"hnsw:space": "cosine"},
        )
        self.child_collection = self.chroma_client.get_or_create_collection(
            name=child_collection_name,
            metadata={"hnsw:space": "cosine"},
        )

        self.session_id: str | None = None
        self.source_name: str = ""
        self.file_type: str = ""
        self.chunks: list[str] = []
        self.child_chunks: list[str] = []
        self.tfidf_vectorizer: TfidfVectorizer | None = None
        self.tfidf_matrix = None

    # ── API Key 管理 ─────────────────────────────────────
    def set_api_key(self, api_key: str) -> None:
        """動態設定 Groq API Key，建立或更新 client。"""
        key = (api_key or "").strip()
        self.client = Groq(api_key=key) if key else None

    # ── 文件載入 ─────────────────────────────────────────
    def load_document(self, file_path: str) -> str:
        try:
            path = Path(file_path)
            if not path.exists():
                return "✗ 載入失敗：找不到檔案"

            suffix = path.suffix.lower()
            if suffix not in (".pdf", ".docx"):
                return "✗ 目前僅支援 PDF 與 DOCX 檔案"

            self.source_name = path.name
            self.file_type = suffix.lstrip(".")
            self.session_id = (
                f"{int(time.time())}_{re.sub(r'[^a-zA-Z0-9]+', '_', path.stem)[:40]}"
            )

            if suffix == ".pdf":
                full_text, stats = self._extract_pdf(path)
            else:
                full_text, stats = self._extract_docx(path)

            if not full_text.strip():
                return "✗ 載入失敗：文件沒有可擷取文字，可能是掃描圖片檔，需先 OCR"

            self.chunks = self._split(full_text, chunk_size=800, overlap=150)
            if not self.chunks:
                return "✗ 載入失敗：切段後沒有有效內容"

            self._build_chroma_index()
            self._build_tfidf_index()
            self._build_child_index()

            return (
                f"✓ 成功載入 {self.source_name}\n"
                f"類型：{suffix.upper().lstrip('.')}\n"
                f"{len(self.chunks)} 個主片段 · ChromaDB Session：{self.session_id}"
            )
        except Exception as exc:
            return f"✗ 載入失敗：{type(exc).__name__}: {exc}"

    # ── 文字擷取 ─────────────────────────────────────────
    def _extract_pdf(self, path: Path) -> tuple[str, str]:
        reader = PdfReader(str(path))
        parts = []
        for idx, page in enumerate(reader.pages, 1):
            text = page.extract_text() or ""
            if text.strip():
                parts.append(f"\n[PDF 第 {idx} 頁]\n{text}")
        return "\n".join(parts), f"{len(reader.pages)} 頁"

    def _extract_docx(self, path: Path) -> tuple[str, str]:
        doc = Document(str(path))
        blocks: list[str] = []
        para_count = table_count = 0

        for child in doc.element.body.iterchildren():
            if isinstance(child, CT_P):
                text = Paragraph(child, doc).text.strip()
                if text:
                    para_count += 1
                    blocks.append(text)
            elif isinstance(child, CT_Tbl):
                table_count += 1
                tbl_text = self._table_to_text(Table(child, doc))
                if tbl_text.strip():
                    blocks.append(f"\n[DOCX 表格 {table_count}]\n{tbl_text}")

        return "\n\n".join(blocks), f"{para_count} 段落 / {table_count} 表格"

    def _table_to_text(self, table: Table) -> str:
        rows = []
        for row in table.rows:
            cells = [re.sub(r"\s+", " ", c.text).strip() for c in row.cells if c.text.strip()]
            if cells:
                rows.append(" | ".join(cells))
        return "\n".join(rows)

    def _split(self, text: str, chunk_size: int, overlap: int) -> list[str]:
        clean = re.sub(r"\s+", " ", text).strip()
        step = max(1, chunk_size - overlap)
        return [
            c for start in range(0, len(clean), step)
            if (c := clean[start: start + chunk_size].strip())
        ]

    # ── Index 建立 ───────────────────────────────────────
    def _encode(self, texts: list[str]) -> list[list[float]]:
        return (
            self.embedding_model
            .encode(texts, convert_to_numpy=True, normalize_embeddings=True, show_progress_bar=False)
            .astype("float32")
            .tolist()
        )

    def _build_chroma_index(self) -> None:
        sid = self.session_id
        ids = [f"{sid}_chunk_{i:05d}" for i in range(len(self.chunks))]
        metas = [
            {"session_id": sid, "source": self.source_name,
             "file_type": self.file_type, "chunk_index": i}
            for i in range(len(self.chunks))
        ]
        self.collection.add(ids=ids, documents=self.chunks,
                            metadatas=metas, embeddings=self._encode(self.chunks))

    def _build_tfidf_index(self) -> None:
        self.tfidf_vectorizer = TfidfVectorizer(analyzer="char", ngram_range=(2, 4), max_features=3000)
        self.tfidf_matrix = self.tfidf_vectorizer.fit_transform(self.chunks)

    def _build_child_index(self) -> None:
        sid = self.session_id
        child_docs, child_ids, child_metas = [], [], []
        for pidx, parent in enumerate(self.chunks):
            for cidx, child in enumerate(self._split(parent, chunk_size=300, overlap=50)):
                child_docs.append(child)
                child_ids.append(f"{sid}_parent_{pidx:05d}_child_{cidx:03d}")
                child_metas.append({"session_id": sid, "source": self.source_name,
                                    "file_type": self.file_type,
                                    "parent_index": pidx, "child_index": cidx})
        self.child_chunks = child_docs
        if child_docs:
            self.child_collection.add(ids=child_ids, documents=child_docs,
                                       metadatas=child_metas, embeddings=self._encode(child_docs))

    # ── 工具函式 ─────────────────────────────────────────
    def _where(self) -> dict[str, str]:
        return {"session_id": self.session_id or ""}

    def _chroma_search(self, query: str, k: int, child: bool = False) -> list[dict[str, Any]]:
        if not self.session_id:
            return []
        col = self.child_collection if child else self.collection
        results = col.query(
            query_embeddings=self._encode([query]),
            n_results=max(1, k),
            where=self._where(),
            include=["documents", "metadatas", "distances"],
        )
        docs = results.get("documents", [[]])[0] or []
        metas = results.get("metadatas", [[]])[0] or []
        dists = results.get("distances", [[]])[0] or []
        return [{"text": d, "metadata": m or {}, "distance": dist}
                for d, m, dist in zip(docs, metas, dists)]

    def _dedupe(self, chunks: list[str], k: int) -> list[str]:
        seen: set[str] = set()
        out: list[str] = []
        for c in chunks:
            key = c[:120]
            if key not in seen:
                seen.add(key)
                out.append(c)
            if len(out) >= k:
                break
        return out

    def _llm(self, prompt: str, max_tokens: int = 300, temperature: float = 0.3) -> str | None:
        if not self.client:
            return None
        try:
            r = self.client.chat.completions.create(
                model="llama-3.1-8b-instant",
                messages=[{"role": "user", "content": prompt}],
                max_tokens=max_tokens,
                temperature=temperature,
            )
            return r.choices[0].message.content
        except Exception:
            return None

    # ── 8 種策略 ──────────────────────────────────────────
    def s_semantic(self, query: str, k: int = 3) -> list[str]:
        return [r["text"] for r in self._chroma_search(query, k)]

    def s_tfidf(self, query: str, k: int = 3) -> list[str]:
        if self.tfidf_vectorizer is None or self.tfidf_matrix is None:
            return []
        qv = self.tfidf_vectorizer.transform([query])
        scores = (self.tfidf_matrix * qv.T).toarray().flatten()
        return [self.chunks[i] for i in scores.argsort()[-k:][::-1]]

    def s_hybrid(self, query: str, k: int = 3) -> list[str]:
        return self._dedupe(
            self.s_semantic(query, k * 2) + self.s_tfidf(query, k * 2), k
        )

    def s_rerank(self, query: str, k: int = 3) -> list[str]:
        candidates = self.s_semantic(query, k * 2)
        if not self.client:
            return candidates[:k]
        scored: list[tuple[str, float]] = []
        for chunk in candidates:
            prompt = (f"問題：{query}\n\n文本：{chunk[:500]}\n\n"
                      f"請只輸出 0 到 10 的相關度分數（僅數字）：")
            resp = self._llm(prompt, max_tokens=10, temperature=0)
            nums = re.findall(r"\d+(?:\.\d+)?", resp or "")
            scored.append((chunk, float(nums[0]) if nums else 0.0))
        scored.sort(key=lambda x: x[1], reverse=True)
        return [c for c, _ in scored[:k]]

    def s_multi_query(self, query: str, k: int = 3) -> list[str]:
        queries = [query]
        prompt = f"將以下問題改寫成 3 個角度不同的繁體中文問題，每行一題，不加編號：\n{query}"
        resp = self._llm(prompt, max_tokens=200, temperature=0.7)
        if resp:
            extras = [ln.strip("-• 1234567890.、 ") for ln in resp.splitlines() if ln.strip()]
            queries += extras[:3]
        chunks: list[str] = []
        for q in queries:
            chunks.extend(self.s_semantic(q, 2))
        return self._dedupe(chunks, k)

    def s_compress(self, query: str, k: int = 3) -> list[str]:
        chunks = self.s_semantic(query, k)
        if not self.client:
            return chunks
        compressed = []
        for chunk in chunks:
            prompt = (f"從以下文本中，提取與問題「{query}」最相關的 1-2 句，"
                      f"保留繁體中文，不要添加任何解釋：\n\n{chunk}")
            resp = self._llm(prompt, max_tokens=180, temperature=0)
            compressed.append((resp or "").strip() or chunk[:350])
        return compressed

    def s_parent_child(self, query: str, k: int = 3) -> list[str]:
        hits = self._chroma_search(query, k * 3, child=True)
        seen_parents: list[int] = []
        for h in hits:
            pidx = h.get("metadata", {}).get("parent_index")
            if isinstance(pidx, int) and pidx not in seen_parents:
                seen_parents.append(pidx)
            if len(seen_parents) >= k:
                break
        return [self.chunks[i] for i in seen_parents if 0 <= i < len(self.chunks)]

    def s_hyde(self, query: str, k: int = 3) -> list[str]:
        prompt = f"請對以下問題給出一段假設性簡短答案（繁體中文）：\n{query}"
        hypo = self._llm(prompt, max_tokens=250, temperature=0.7) or query
        return self.s_semantic(hypo, k)

    # ── 策略路由 ──────────────────────────────────────────
    _FN = {
        "semantic":     s_semantic,
        "tfidf":        s_tfidf,
        "hybrid":       s_hybrid,
        "rerank":       s_rerank,
        "multi_query":  s_multi_query,
        "compress":     s_compress,
        "parent_child": s_parent_child,
        "hyde":         s_hyde,
    }

    def generate_answer(self, query: str, strategy_key: str, top_k: int):
        if not self.chunks:
            return "請先上傳並載入 PDF 或 DOCX 文件。", ""
        if not query.strip():
            return "請輸入問題。", ""

        fn = self._FN.get(strategy_key, self.s_semantic)
        chunks = fn(self, query, int(top_k))
        context = "\n\n—\n\n".join(chunks)

        strategy_label = self.STRATEGY_MAP.get(strategy_key, strategy_key)
        source_preview = (
            f"文件：{self.source_name}\n"
            f"策略：{strategy_label} · 片段數：{len(chunks)}\n"
            f"ChromaDB Session：{self.session_id}\n\n"
            f"{'─' * 56}\n\n{context}"
        )

        if not self.client:
            return (
                "⚠ 尚未設定 Groq API Key。\n"
                "請在左欄「Step 00」輸入您的 Groq API Key 並點擊「套用」後再提問。\n\n"
                "（檢索已完成，可在下方「查看檢索到的文本片段」確認結果）",
                source_preview,
            )

        prompt = f"""請根據以下上下文回答問題。若上下文無相關資訊，請明確說明無法從文件回答，不要自行編造。

上下文：
{context}

問題：{query}

請用繁體中文詳細回答，並以條列方式整理重點："""

        try:
            r = self.client.chat.completions.create(
                model="llama-3.1-8b-instant",
                messages=[
                    {"role": "system", "content": "你是專業的文件分析與 RAG 問答助手。"},
                    {"role": "user", "content": prompt},
                ],
                max_tokens=1024,
                temperature=0.3,
            )
            return r.choices[0].message.content, source_preview
        except Exception as exc:
            return f"生成失敗：{type(exc).__name__}: {exc}", source_preview


# ══════════════════════════════════════════════════════════
#  Gradio UI
# ══════════════════════════════════════════════════════════
STRATEGY_INFO = [
    ("semantic",     "語意搜尋",   "ChromaDB 向量相似度，最通用",         "🔍"),
    ("tfidf",        "TF-IDF",     "字元 n-gram 關鍵詞統計",              "📊"),
    ("hybrid",       "混合搜尋",   "語意 + TF-IDF 結果合併去重",          "⚡"),
    ("rerank",       "重新排序",   "LLM 對候選片段二次評分",              "🎯"),
    ("multi_query",  "多查詢擴展", "自動生成多角度問題聯合搜尋",          "🔄"),
    ("compress",     "上下文壓縮", "LLM 提取最相關句子精簡上下文",        "✂️"),
    ("parent_child", "父子文檔",   "小片段定位 → 回傳對應大片段",        "📂"),
    ("hyde",         "HyDE",       "先生成假設答案再語意搜尋",            "💡"),
]

CSS = """
body, .gradio-container { background:#f5f4f1 !important; }

#hdr {
    background:#fff;
    border:1px solid #e5e0d8;
    border-radius:14px;
    padding:28px 36px;
    margin-bottom:20px;
    border-top: 4px solid #2d6a4f;
}
.hdr-eyebrow { font-size:11px; letter-spacing:2.5px; color:#2d6a4f; text-transform:uppercase; margin-bottom:6px; }
.hdr-title   { font-size:26px; font-weight:700; color:#1a1714; margin:0 0 6px; }
.hdr-sub     { font-size:14px; color:#6b5e56; }
.pill { display:inline-block; margin:10px 5px 0 0; padding:3px 10px; border-radius:16px;
        font-size:11px; background:#e8f4f0; color:#2d6a4f; border:1px solid rgba(45,106,79,.2); }
.pill-amber { background:#fdf4e3; color:#b87a1a; border-color:rgba(184,122,26,.25); }

/* API Key 區塊 */
#apikey-box {
    background: #fffbf2;
    border: 1.5px solid #f0c96a;
    border-radius: 10px;
    padding: 12px 14px;
    margin-bottom: 8px;
}

.strat-grid { display:grid; grid-template-columns:repeat(4,1fr); gap:10px; margin:10px 0 16px; }
.strat-card {
    background:#fff;
    border:1.5px solid #e5e0d8;
    border-radius:10px;
    padding:10px 12px;
    cursor:pointer;
    transition:border-color .15s, box-shadow .15s;
    text-align:left;
    width:100%;
}
.strat-card:hover { border-color:#2d6a4f; box-shadow:0 2px 8px rgba(45,106,79,.12); }
.strat-card.active { border-color:#2d6a4f; background:#f0f9f5; box-shadow:0 2px 10px rgba(45,106,79,.18); }
.strat-icon  { font-size:20px; margin-bottom:4px; }
.strat-name  { font-size:13px; font-weight:700; color:#1a1714; margin:0 0 2px; }
.strat-desc  { font-size:11px; color:#7a6e67; line-height:1.4; }

.sec-label { font-size:11px; letter-spacing:1.5px; text-transform:uppercase;
             color:#7a6e67; font-weight:700; margin:16px 0 8px; }
.card-box { background:#fff !important; border:1px solid #e5e0d8 !important;
            border-radius:12px !important; padding:16px !important; }
#ask-btn { background:#2d6a4f !important; color:#fff !important; border:0 !important; border-radius:8px !important; }
#apply-key-btn { background:#b87a1a !important; color:#fff !important; border:0 !important; border-radius:8px !important; }
"""

HEADER_HTML = """
<div id="hdr">
  <div class="hdr-eyebrow">Intelligent Document Analysis · v2</div>
  <div class="hdr-title">多策略 RAG 文件問答系統</div>
  <div class="hdr-sub">支援 PDF / DOCX 上傳，採用 ChromaDB 持久化向量資料庫與 8 種 RAG 檢索策略</div>
  <div>
    <span class="pill">▸ Groq API</span>
    <span class="pill">▸ llama-3.1-8b-instant</span>
    <span class="pill pill-amber">▸ ChromaDB</span>
    <span class="pill pill-amber">▸ PDF / DOCX</span>
    <span class="pill">▸ SentenceTransformers</span>
  </div>
</div>
"""

def build_strategy_menu(selected: str = "semantic") -> str:
    cards = []
    for key, name, desc, icon in STRATEGY_INFO:
        active_cls = "active" if key == selected else ""
        cards.append(
            f"""<button class="strat-card {active_cls}" onclick="selectStrategy('{key}', this)" type="button">
  <div class="strat-icon">{icon}</div>
  <div class="strat-name">{name}</div>
  <div class="strat-desc">{desc}</div>
</button>"""
        )
    return f'<div class="strat-grid">{" ".join(cards)}</div>'


STRATEGY_MENU_JS = """
<script>
function selectStrategy(key, el) {
    document.querySelectorAll('.strat-card').forEach(c => c.classList.remove('active'));
    el.classList.add('active');
    const inp = document.getElementById('strategy-hidden');
    if (inp) { inp.value = key; inp.dispatchEvent(new Event('input')); }
}
</script>
"""

EXAMPLE_QS = [
    ["這份文件的主要內容是什麼？"],
    ["文件中提到哪些重要概念或定義？"],
    ["有哪些關鍵數據、統計資料或案例？"],
    ["文件的結論或建議是什麼？"],
    ["文件提及哪些潛在風險或挑戰？"],
]


def create_interface():
    # 啟動時嘗試從環境變數讀取（可留空）
    env_key = os.getenv("GROQ_API_KEY", "").strip()
    rag = MultiStrategyRAG(chroma_path="/tmp/chroma_db")
    if env_key:
        rag.set_api_key(env_key)

    current_strategy = {"key": "semantic"}

    def apply_api_key(api_key: str):
        key = (api_key or "").strip()
        rag.set_api_key(key)
        if key:
            masked = key[:8] + "****" + key[-4:] if len(key) > 12 else "****"
            return f"✓ API Key 已套用（{masked}）"
        return "⚠ API Key 已清除，無法呼叫 LLM"

    def upload_document(file):
        if file is None:
            return "⚠ 請選擇 PDF 或 DOCX 檔案"
        return rag.load_document(file.name)

    def set_strategy(key: str):
        current_strategy["key"] = key
        return f"✓ 已選擇策略：{dict((k, n) for k, n, *_ in STRATEGY_INFO).get(key, key)}"

    def ask(query, top_k):
        return rag.generate_answer(query, current_strategy["key"], int(top_k))

    with gr.Blocks(
        title="多策略 RAG 文件問答 v2",
        css=CSS,
        theme=gr.themes.Base(
            primary_hue=gr.themes.colors.green,
            neutral_hue=gr.themes.colors.stone,
        ),
    ) as demo:
        gr.HTML(HEADER_HTML)

        with gr.Row(equal_height=False):
            # ── 左欄 ──────────────────────────────────
            with gr.Column(scale=1, min_width=320, elem_classes="card-box"):

                # ★ Step 00：API Key 輸入（新增）
                gr.HTML("<div class='sec-label'>Step 00 · Groq API Key</div>")
                with gr.Group(elem_id="apikey-box"):
                    api_key_input = gr.Textbox(
                        label="",
                        placeholder="gsk_xxxxxxxxxxxxxxxxxxxxxxxx",
                        value=env_key,       # 若環境變數已設定則預填
                        type="password",     # 輸入時遮蔽顯示
                        lines=1,
                        show_label=False,
                    )
                    apply_key_btn = gr.Button(
                        "套用 API Key", size="sm", elem_id="apply-key-btn"
                    )
                    api_key_status = gr.Textbox(
                        value="✓ API Key 已從環境變數載入" if env_key else "⚠ 尚未設定 API Key",
                        interactive=False,
                        lines=1,
                        label="",
                        show_label=False,
                    )

                # Step 01：上傳文件
                gr.HTML("<div class='sec-label'>Step 01 · 上傳文件</div>")
                file_input = gr.File(label="PDF / DOCX", file_types=[".pdf", ".docx"])
                load_btn = gr.Button("↑ 載入文件")
                status = gr.Textbox(label="狀態", interactive=False, lines=3)

                # Step 02：RAG 策略
                gr.HTML("<div class='sec-label'>Step 02 · 選擇 RAG 策略</div>")
                gr.HTML(build_strategy_menu("semantic"))
                strategy_input = gr.Textbox(
                    value="semantic",
                    elem_id="strategy-hidden",
                    label="",
                    visible=False,
                )
                strategy_status = gr.Textbox(
                    value="✓ 已選擇策略：語意搜尋",
                    interactive=False,
                    lines=1,
                    label="目前策略",
                )
                gr.HTML(STRATEGY_MENU_JS)

                # Step 03：參數
                gr.HTML("<div class='sec-label'>Step 03 · 搜尋參數</div>")
                topk = gr.Slider(minimum=1, maximum=10, value=3, step=1, label="Top-K 片段數量")

            # ── 右欄：問答 ────────────────────────────
            with gr.Column(scale=2, elem_classes="card-box"):
                gr.HTML("<div class='sec-label'>Step 04 · 輸入問題</div>")
                qin = gr.Textbox(
                    label="",
                    placeholder="例如：這份文件的核心論點是什麼？",
                    lines=4,
                )
                ask_btn = gr.Button("提問", variant="primary", size="lg", elem_id="ask-btn")

                gr.HTML("<div class='sec-label'>AI 回答</div>")
                ans = gr.Textbox(label="", lines=12, interactive=False)

                with gr.Accordion("▸ 查看檢索到的文本片段", open=False):
                    src = gr.Textbox(label="", lines=10, interactive=False)

                gr.Examples(examples=EXAMPLE_QS, inputs=qin, label="範例問題")

        # ── 事件綁定 ──────────────────────────────────
        apply_key_btn.click(fn=apply_api_key, inputs=[api_key_input], outputs=[api_key_status])
        api_key_input.submit(fn=apply_api_key, inputs=[api_key_input], outputs=[api_key_status])
        load_btn.click(fn=upload_document, inputs=[file_input], outputs=[status])
        strategy_input.change(fn=set_strategy, inputs=[strategy_input], outputs=[strategy_status])
        ask_btn.click(fn=ask, inputs=[qin, topk], outputs=[ans, src])
        qin.submit(fn=ask, inputs=[qin, topk], outputs=[ans, src])

    return demo


if __name__ == "__main__":
    demo = create_interface()
    demo.launch(share=False, server_name="0.0.0.0",show_error=True)