app.py

"""
多策略 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
"""

from __future__ import annotations

import os
import re
import time
from pathlib import Path
from typing import Any

import chromadb
import gradio as gr
import numpy as np
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


# ══════════════════════════════════════════════════════════
#  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 = "./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('.')} · {stats}\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 = """
:root {
    --paper: #f4eee3;
    --paper-2: #ece4d4;
    --ink: #1b2434;
    --gold: #b88a42;
    --gold-soft: #d5b070;
    --line: rgba(126, 97, 52, .28);
    --shadow: 0 18px 44px rgba(18, 21, 27, .14);
}

body, .gradio-container {
    background:
        radial-gradient(circle at top left, rgba(184,138,66,.12), transparent 28%),
        radial-gradient(circle at bottom right, rgba(27,36,52,.10), transparent 26%),
        linear-gradient(180deg, #f8f3e9 0%, #efe6d7 100%) !important;
    color: var(--ink) !important;
    font-family: 'Noto Serif TC', 'Microsoft JhengHei', serif !important;
}

.gradio-container { max-width: 1340px !important; }
.contain, .app { background: transparent !important; }

#hdr {
    position: relative;
    overflow: hidden;
    background:
        linear-gradient(135deg, rgba(255,248,236,.98), rgba(247,238,222,.96));
    border: 1px solid rgba(145, 116, 68, .34);
    border-radius: 28px;
    padding: 34px 36px 30px;
    margin-bottom: 20px;
    box-shadow: var(--shadow);
}
#hdr::before {
    content: '';
    position: absolute;
    inset: 0;
    pointer-events: none;
    background:
        radial-gradient(circle at 14% 18%, rgba(184,138,66,.13), transparent 16%),
        radial-gradient(circle at 88% 12%, rgba(27,36,52,.10), transparent 18%),
        linear-gradient(90deg, rgba(184,138,66,.18), transparent 22%, transparent 78%, rgba(184,138,66,.18));
    opacity: .78;
}
.hdr-seal {
    position: absolute;
    top: 18px;
    right: 22px;
    color: rgba(171, 52, 42, .75);
    border: 1px solid rgba(171, 52, 42, .3);
    border-radius: 12px;
    padding: 6px 10px;
    font-size: 12px;
    letter-spacing: .24em;
    background: rgba(255,245,238,.65);
    z-index: 1;
}
.hdr-eyebrow { position: relative; z-index: 1; font-size: 11px; letter-spacing: .34em; color: #8f6d32; text-transform: uppercase; margin-bottom: 10px; }
.hdr-title { position: relative; z-index: 1; font-size: 36px; line-height: 1.18; font-weight: 800; color: #1c2433; margin: 0 0 12px; }
.hdr-sub { position: relative; z-index: 1; font-size: 15px; line-height: 1.9; color: #5e5649; max-width: 920px; }
.hdr-quote { position: relative; z-index: 1; margin-top: 12px; color: #6f624d; font-size: 13px; letter-spacing: .08em; }
.hdr-pills { position: relative; z-index: 1; margin-top: 16px; }
.pill { display: inline-block; margin: 8px 8px 0 0; padding: 5px 12px; border-radius: 999px; font-size: 11px; color: #7f5a1e; background: rgba(245,233,205,.95); border: 1px solid rgba(184,138,66,.32); }
.pill-dark { color: #edf0f5; background: rgba(27,36,52,.92); border-color: rgba(27,36,52,.92); }

.card-box {
    background: linear-gradient(180deg, rgba(255,250,242,.94), rgba(250,244,233,.92)) !important;
    border: 1px solid rgba(145, 116, 68, .26) !important;
    border-radius: 24px !important;
    padding: 20px !important;
    box-shadow: var(--shadow);
}
.sec-label {
    display: flex; align-items: center; gap: 10px;
    font-size: 15px; letter-spacing: .22em; text-transform: uppercase;
    color: #26324a; font-weight: 800; margin: 10px 0 14px;
}
.sec-label::after { content: ''; flex: 1; height: 1px; background: linear-gradient(90deg, rgba(184,138,66,.6), rgba(184,138,66,0)); }
.sub-note { margin: -4px 0 14px; color: #7a6a57; font-size: 12px; line-height: 1.75; }
#apikey-box {
    background: linear-gradient(180deg, rgba(255,244,219,.84), rgba(255,249,238,.92));
    border: 1.5px solid rgba(190,135,33,.85);
    border-radius: 18px; padding: 16px 16px 10px; margin-bottom: 10px;
    box-shadow: inset 0 0 0 1px rgba(255,255,255,.5);
}

.strat-grid { display:grid; grid-template-columns:repeat(4,1fr); gap:12px; margin:12px 0 16px; }
.strat-card {
    background: linear-gradient(180deg, rgba(255,249,239,.98), rgba(246,238,223,.94));
    border:1px solid rgba(148,120,73,.28);
    border-radius:16px; padding:12px; cursor:pointer; text-align:left; width:100%;
    transition:transform .18s ease, border-color .18s ease, box-shadow .18s ease; box-shadow: 0 6px 18px rgba(0,0,0,.04);
}
.strat-card:hover { border-color:#b88a42; box-shadow:0 10px 24px rgba(184,138,66,.12); transform: translateY(-2px); }
.strat-card.active { border-color:#b88a42; background: linear-gradient(180deg, rgba(255,246,224,.98), rgba(249,238,212,.96)); box-shadow:0 12px 26px rgba(184,138,66,.16); }
.strat-icon { font-size:20px; margin-bottom:6px; }
.strat-name { font-size:14px; font-weight:800; color:#1c2433; margin:0 0 3px; }
.strat-desc { font-size:11.5px; color:#6d6253; line-height:1.55; }

#ask-btn { background: linear-gradient(180deg, #22314b, #172131) !important; color:#fff !important; border:1px solid #172131 !important; border-radius:12px !important; }
#apply-key-btn { background: linear-gradient(180deg, #c18a22, #a97012) !important; color:#fff !important; border:1px solid #8e5d11 !important; border-radius:12px !important; }
#ask-btn:hover, #apply-key-btn:hover { filter: brightness(1.08); }
button.secondary, .secondary-button { border-radius: 12px !important; }
.gr-textbox, .gr-file, .gr-slider, .gr-accordion, .gr-examples { --block-background-fill: transparent !important; }
textarea, input, .gradio-container textarea, .gradio-container input { font-family: 'Noto Serif TC', 'Microsoft JhengHei', serif !important; }
.gr-textbox label, .gr-file label, .gr-slider label { color: #3a2e1b !important; font-size: 12px !important; letter-spacing: .08em !important; }
.gr-textbox textarea, .gr-textbox input, textarea.scroll-hide, input[type='text'], input[type='password'] {
    background: linear-gradient(180deg, rgba(29,38,55,.96), rgba(36,49,70,.94)) !important;
    color: #f6efe3 !important; border: 1px solid rgba(164, 137, 94, .55) !important; border-radius: 16px !important;
}
.gr-textbox textarea::placeholder, .gr-textbox input::placeholder, textarea::placeholder, input::placeholder { color: #b9b3aa !important; }
.gr-textbox textarea[disabled], .gr-textbox input[disabled], textarea[disabled], input[disabled] {
    background: linear-gradient(180deg, rgba(58,53,47,.96), rgba(70,64,56,.94)) !important; color: #f3ebdd !important; opacity: 1 !important;
}
.gr-file { background: rgba(255,251,243,.76) !important; border: 1px dashed rgba(147,117,69,.45) !important; border-radius: 16px !important; padding: 6px !important; }
.gr-accordion { border: 1px solid rgba(147,117,69,.30) !important; border-radius: 16px !important; overflow: hidden !important; background: rgba(255,250,242,.78) !important; }
.gr-accordion summary { background: linear-gradient(180deg, rgba(248,240,225,.95), rgba(241,232,213,.94)) !important; color: #26324a !important; font-weight: 700 !important; }
.gr-accordion .label-wrap span { color: #26324a !important; }
.gr-examples { background: rgba(255,251,243,.72) !important; border: 1px solid rgba(147,117,69,.25) !important; border-radius: 16px !important; padding: 10px 12px !important; }
.gr-examples .label-wrap span, .gr-examples .label-wrap label { color: #7b6238 !important; font-weight: 700 !important; }
.gr-examples button { background: rgba(255,247,230,.96) !important; border: 1px solid rgba(184,138,66,.30) !important; color: #4f4334 !important; border-radius: 12px !important; }
.gr-examples button:hover { border-color: #b88a42 !important; color: #26324a !important; }
input[type='range'] { accent-color: #b88a42 !important; }
@media (max-width: 980px) { .strat-grid { grid-template-columns: repeat(2, 1fr); } .hdr-title { font-size: 29px; } #hdr { padding: 26px 22px; } }
@media (max-width: 640px) { .strat-grid { grid-template-columns: 1fr; } .sec-label { font-size: 13px; letter-spacing: .16em; } .hdr-title { font-size: 24px; } .hdr-seal { display: none; } }

/* ===== 頁首可讀性修正：改成深色墨境底，避免白字吃掉 ===== */
#hdr {
    background:
        radial-gradient(circle at 72% 18%, rgba(214, 176, 103, .16), transparent 24%),
        radial-gradient(circle at 18% 80%, rgba(54, 77, 114, .38), transparent 30%),
        linear-gradient(135deg, #101928 0%, #19283d 48%, #253855 100%) !important;
    border: 1px solid rgba(214, 176, 103, .38) !important;
    box-shadow: 0 18px 44px rgba(16, 25, 40, .30) !important;
}
#hdr::before {
    background:
        linear-gradient(90deg, rgba(214,176,103,.18), transparent 24%, transparent 75%, rgba(214,176,103,.13)),
        radial-gradient(circle at 86% 20%, rgba(255,255,255,.08), transparent 18%) !important;
    opacity: 1 !important;
}
#hdr .hdr-eyebrow {
    color: #d7b679 !important;
    opacity: 1 !important;
    text-shadow: 0 1px 2px rgba(0,0,0,.45) !important;
}
#hdr .hdr-title {
    color: #fff5df !important;
    opacity: 1 !important;
    text-shadow: 0 2px 8px rgba(0,0,0,.55) !important;
}
#hdr .hdr-sub,
#hdr .hdr-quote {
    color: #f2e6cc !important;
    opacity: 1 !important;
    text-shadow: 0 1px 4px rgba(0,0,0,.45) !important;
}
#hdr .hdr-seal {
    color: #d7b679 !important;
    border-color: rgba(215,182,121,.42) !important;
    background: rgba(255,245,220,.08) !important;
    text-shadow: 0 1px 3px rgba(0,0,0,.45) !important;
}
#hdr .pill {
    color: #fff1d0 !important;
    background: rgba(255, 241, 208, .10) !important;
    border-color: rgba(215, 182, 121, .44) !important;
    opacity: 1 !important;
}
#hdr .pill-dark {
    color: #101928 !important;
    background: #d7b679 !important;
    border-color: #d7b679 !important;
    font-weight: 800 !important;
}


/* ===== 全域可讀性修正：避免淺色紙面上的文字太白看不見 ===== */
.gradio-container {
    --body-text-color: #211a12 !important;
    --block-title-text-color: #211a12 !important;
    --block-label-text-color: #2d2418 !important;
    --input-placeholder-color: #b8ad9d !important;
}

.card-box {
    color: #241c13 !important;
}

.card-box .sub-note,
.sub-note {
    color: #4b3a25 !important;
    opacity: 1 !important;
    font-weight: 650 !important;
    text-shadow: none !important;
}

.card-box .sec-label,
.sec-label {
    color: #1c2433 !important;
    opacity: 1 !important;
    font-weight: 900 !important;
    text-shadow: none !important;
}

.card-box .sec-label *,
.sec-label * {
    color: #1c2433 !important;
    opacity: 1 !important;
}

/* Gradio 產生的標籤文字：強制深色 */
.gradio-container label,
.gradio-container .label-wrap,
.gradio-container .label-wrap span,
.gradio-container .block-title,
.gradio-container .block-label,
.gradio-container .form label,
.gradio-container .input-container label {
    color: #2b2116 !important;
    opacity: 1 !important;
    font-weight: 700 !important;
}

/* 深色輸入框與輸出框內的 label / 內容維持亮色 */
.gr-textbox label,
#apikey-box label {
    color: #f5e8c8 !important;
    opacity: 1 !important;
    font-weight: 800 !important;
}

.gr-textbox textarea,
.gr-textbox input,
textarea.scroll-hide,
input[type='text'],
input[type='password'] {
    color: #fff3d6 !important;
    opacity: 1 !important;
}

/* API 狀態與卷宗狀態等唯讀框，文字更亮更清楚 */
.gr-textbox textarea[disabled],
.gr-textbox input[disabled],
textarea[disabled],
input[disabled] {
    color: #fff1cf !important;
    opacity: 1 !important;
    font-weight: 700 !important;
}

/* 檔案上傳區是深色底，所以裡面的文字維持亮色 */
.gr-file,
.gr-file *,
.gr-upload,
.gr-upload * {
    color: #fff3d6 !important;
    opacity: 1 !important;
}

/* 範例問題與 Accordion 標題：避免變成淡白或藍色 */
.gr-examples,
.gr-examples *,
.gr-accordion,
.gr-accordion summary,
.gr-accordion summary * {
    opacity: 1 !important;
}

.gr-examples .label-wrap span,
.gr-examples .label-wrap label {
    color: #5f431d !important;
}

.gr-examples button {
    color: #2b2116 !important;
    font-weight: 650 !important;
}

.gr-accordion summary,
.gr-accordion summary * {
    color: #fff3d6 !important;
    font-weight: 800 !important;
}

/* 策略卡片：Gradio button 樣式會覆蓋背景，這裡重新壓回可讀配色 */
.strat-card {
    background: linear-gradient(180deg, rgba(54, 49, 42, .96), rgba(39, 35, 31, .96)) !important;
    border: 1px solid rgba(215, 182, 121, .38) !important;
    color: #fff3d6 !important;
}

.strat-card .strat-name,
.strat-card .strat-desc,
.strat-card .strat-icon {
    color: #fff3d6 !important;
    opacity: 1 !important;
}

.strat-card.active {
    background: linear-gradient(180deg, rgba(184, 138, 42, .98), rgba(125, 83, 21, .98)) !important;
    border-color: #ffd98a !important;
}

/* 避免整個介面被滑鼠拖曳選成藍色；輸入框仍可複製 */
.gradio-container :not(textarea):not(input) {
    user-select: none !important;
}
.gradio-container textarea,
.gradio-container input {
    user-select: text !important;
}
::selection {
    background: rgba(184, 138, 66, 0.45);
    color: #fff5df;
}


/* ===== 深色元件最終潤色：修正黑字 / 暗字卡在深色框內看不見 ===== */
/* Textbox 的標題、label、說明文字都改成米金色 */
.gradio-container .gr-textbox label,
.gradio-container .gr-textbox .label-wrap,
.gradio-container .gr-textbox .label-wrap *,
.gradio-container [data-testid="textbox"] label,
.gradio-container [data-testid="textbox"] .label-wrap,
.gradio-container [data-testid="textbox"] .label-wrap *,
.gradio-container .input-container label,
.gradio-container .wrap label {
    color: #f6e6c4 !important;
    opacity: 1 !important;
    font-weight: 800 !important;
    text-shadow: 0 1px 2px rgba(0,0,0,.45) !important;
}

/* Textbox 本體：深墨底、亮字、亮 placeholder */
.gradio-container .gr-textbox,
.gradio-container [data-testid="textbox"] {
    color: #f6e6c4 !important;
}

.gradio-container textarea,
.gradio-container input[type="text"],
.gradio-container input[type="password"] {
    background: linear-gradient(180deg, #292522 0%, #3a332b 100%) !important;
    color: #fff4d8 !important;
    -webkit-text-fill-color: #fff4d8 !important;
    border: 1px solid rgba(215, 182, 121, .45) !important;
    opacity: 1 !important;
}

.gradio-container textarea::placeholder,
.gradio-container input::placeholder {
    color: #d0c0a8 !important;
    opacity: 1 !important;
    -webkit-text-fill-color: #d0c0a8 !important;
}

/* 唯讀回答框 / 狀態框也保持亮字，不要變黑 */
.gradio-container textarea[disabled],
.gradio-container input[disabled],
.gradio-container textarea[readonly],
.gradio-container input[readonly] {
    background: linear-gradient(180deg, #312b24 0%, #463e34 100%) !important;
    color: #fff1cf !important;
    -webkit-text-fill-color: #fff1cf !important;
    opacity: 1 !important;
    font-weight: 650 !important;
}

/* 檔案上傳區：修正「上傳卷宗、拖放檔案、點擊上傳」過暗問題 */
.gradio-container .gr-file,
.gradio-container .gr-file *,
.gradio-container .gr-upload,
.gradio-container .gr-upload *,
.gradio-container [data-testid="file"],
.gradio-container [data-testid="file"] *,
.gradio-container [data-testid="file-upload"],
.gradio-container [data-testid="file-upload"] * {
    color: #fff1cf !important;
    fill: #fff1cf !important;
    stroke: #fff1cf !important;
    opacity: 1 !important;
    text-shadow: 0 1px 2px rgba(0,0,0,.45) !important;
}

.gradio-container .gr-file,
.gradio-container .gr-upload,
.gradio-container [data-testid="file"],
.gradio-container [data-testid="file-upload"] {
    background: #27221f !important;
    border: 1px dashed rgba(215, 182, 121, .38) !important;
    border-radius: 14px !important;
}

/* 檢索片段 Accordion：深色底時，標題與箭頭都改亮 */
.gradio-container .gr-accordion,
.gradio-container .gr-accordion *,
.gradio-container details,
.gradio-container details *,
.gradio-container summary,
.gradio-container summary * {
    color: #fff1cf !important;
    fill: #fff1cf !important;
    stroke: #fff1cf !important;
    opacity: 1 !important;
}

.gradio-container .gr-accordion summary,
.gradio-container details summary {
    background: #241f1d !important;
    border: 1px solid rgba(215, 182, 121, .30) !important;
    border-radius: 8px !important;
    font-weight: 800 !important;
}

/* 問題 / AI 回答這種元件內部黑底區域，所有小標都用亮色 */
.gradio-container .block,
.gradio-container .form,
.gradio-container .wrap,
.gradio-container .container {
    --block-label-text-color: #f6e6c4 !important;
    --block-title-text-color: #f6e6c4 !important;
}

/* 但外層章節標題仍維持深色紙面可讀 */
.gradio-container .sec-label,
.gradio-container .sec-label * {
    color: #1c2433 !important;
    text-shadow: none !important;
}

/* 範例問題是淺色底，不套深色字不清楚 */
.gradio-container .gr-examples,
.gradio-container .gr-examples .label-wrap,
.gradio-container .gr-examples .label-wrap *,
.gradio-container .gr-examples button,
.gradio-container .gr-examples button * {
    color: #2b2116 !important;
    text-shadow: none !important;
    opacity: 1 !important;
}

.gradio-container .gr-examples button {
    background: rgba(255, 248, 234, .98) !important;
    border: 1px solid rgba(94, 74, 40, .45) !important;
}

/* ===== 修正右下方 Accordion / Details 標題太暗 ===== */
.gradio-container details,
.gradio-container .gr-accordion,
.gradio-container [data-testid="accordion"],
.gradio-container [class*="accordion"] {
    background: #241f1d !important;
    border: 1px solid rgba(184, 138, 66, 0.55) !important;
    border-radius: 8px !important;
    color: #fff2c7 !important;
}

.gradio-container details summary,
.gradio-container .gr-accordion summary,
.gradio-container [data-testid="accordion"] summary,
.gradio-container [class*="accordion"] summary,
.gradio-container details summary *,
.gradio-container .gr-accordion summary *,
.gradio-container [data-testid="accordion"] summary *,
.gradio-container [class*="accordion"] summary * {
    background: #241f1d !important;
    color: #fff2c7 !important;
    opacity: 1 !important;
    font-weight: 800 !important;
    text-shadow: none !important;
}

.gradio-container details summary::marker {
    color: #fff2c7 !important;
}

.gradio-container details svg,
.gradio-container .gr-accordion svg,
.gradio-container [data-testid="accordion"] svg {
    color: #fff2c7 !important;
    fill: #fff2c7 !important;
    stroke: #fff2c7 !important;
}

/* 範例問法標題不要太白或太淡 */
.gradio-container .gr-examples,
.gradio-container .gr-examples *,
.gradio-container [data-testid="examples"],
.gradio-container [data-testid="examples"] *,
.gradio-container .examples,
.gradio-container .examples * {
    color: #1c2433 !important;
    opacity: 1 !important;
}

.gradio-container .gr-examples > label,
.gradio-container .gr-examples .label-wrap,
.gradio-container .gr-examples .label-wrap *,
.gradio-container [data-testid="examples"] label,
.gradio-container [data-testid="examples"] .label-wrap,
.gradio-container [data-testid="examples"] .label-wrap * {
    color: #9a6b28 !important;
    font-weight: 800 !important;
}

/* 深色框內的輸出/文字標籤統一亮金色 */
.gradio-container .gr-textbox label,
.gradio-container .gr-textbox label *,
.gradio-container .gr-file label,
.gradio-container .gr-file label *,
.gradio-container .file-preview,
.gradio-container .file-preview *,
.gradio-container [data-testid="file"] label,
.gradio-container [data-testid="file"] label * {
    color: #fff2c7 !important;
    opacity: 1 !important;
    font-weight: 800 !important;
}



/* ===== 範例問法標題修正：改用獨立標題，避免 Gradio label 變淡 ===== */
.examples-title {
    display: flex;
    align-items: center;
    gap: 10px;
    margin: 14px 0 10px;
    color: #1c2433 !important;
    opacity: 1 !important;
    font-size: 14px;
    font-weight: 900;
    letter-spacing: .12em;
    text-shadow: none !important;
}
.examples-title::after {
    content: '';
    flex: 1;
    height: 1px;
    background: linear-gradient(90deg, rgba(184,138,66,.45), rgba(184,138,66,0));
}

/* 隱藏 Gradio Examples 內建淡色 label，保留上方自訂標題 */
.gradio-container .gr-examples > label,
.gradio-container .gr-examples .label-wrap,
.gradio-container [data-testid="examples"] > label,
.gradio-container [data-testid="examples"] .label-wrap,
.gradio-container .examples > label,
.gradio-container .examples .label-wrap {
    display: none !important;
}

/* 範例按鈕：維持紙面深色文字，避免被全域亮字覆蓋 */
.gradio-container .gr-examples button,
.gradio-container .gr-examples button *,
.gradio-container [data-testid="examples"] button,
.gradio-container [data-testid="examples"] button *,
.gradio-container .examples button,
.gradio-container .examples button * {
    color: #21170f !important;
    -webkit-text-fill-color: #21170f !important;
    opacity: 1 !important;
    font-weight: 700 !important;
    text-shadow: none !important;
}

.gradio-container .gr-examples button,
.gradio-container [data-testid="examples"] button,
.gradio-container .examples button {
    background: rgba(255, 248, 234, .98) !important;
    border: 1px solid rgba(66, 49, 26, .55) !important;
    border-radius: 8px !important;
}
"""

HEADER_HTML = """
<div id="hdr">
  <div class="hdr-seal">封印卷宗</div>
  <div class="hdr-eyebrow">Arcane Archive · Multi-Strategy Retrieval</div>
  <div class="hdr-title">墨境卷宗｜多策略 RAG 文件問答系統</div>
  <div class="hdr-sub">以「古卷、夜藍、鎏金、紙墨」為視覺核心，重塑你的文件問答介面。支援 PDF / DOCX 上傳，採用 ChromaDB 持久化向量資料庫與 8 種 RAG 檢索策略。</div>
  <div class="hdr-quote">啟封文件、建立索引、選擇檢索術式 —— 從卷宗片段中召回最接近答案的線索。</div>
  <div class="hdr-pills">
    <span class="pill pill-dark">▸ 多策略檢索</span>
    <span class="pill">▸ Groq API</span>
    <span class="pill">▸ llama-3.1-8b-instant</span>
    <span class="pill">▸ ChromaDB</span>
    <span class="pill">▸ 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="./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"✦ 金鑰已啟封｜卷宗核心已連線（{masked}）"
        return "⚠ 金鑰已封存｜暫時無法召喚 AI 回答"

    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"):

                gr.HTML("<div class='sub-note'>左側控制台化作卷軸儀式區：先啟封金鑰，再載入卷宗，最後選擇檢索術式。</div>")

                # ★ Step 00：API Key 輸入（新增）
                gr.HTML("<div class='sec-label'>Step 00 · 啟封金鑰</div>")
                with gr.Group(elem_id="apikey-box"):
                    api_key_input = gr.Textbox(
                        label="Groq API Key",
                        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="✦ 金鑰已由環境變數啟封｜卷宗核心已連線" if env_key else "⚠ 尚未啟封金鑰｜請先輸入 Groq 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 · 選擇檢索術式</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='sub-note'>右側為發問主場域，採用海報式紙墨配色與深色輸入框，營造強烈奇幻敘事感。</div>")
                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="AI 回答內容", lines=12, interactive=False)

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

                gr.HTML("<div class='examples-title'>三　範例問法</div>")
                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")