"""离线阶段：构建知识库。

RAG 全流程分为离线、在线两阶段，本脚本负责离线部分：
    1. 解析文档（Docling）   —— 把 PDF/DOCX 等转成纯文本
    2. 文本分片              —— 滑动窗口切分成可检索的小段
    3. 向量化（Embedding）   —— 用 Minimax embo-01 把文本转成向量
    4. 写入向量库（Chroma）  —— 持久化存储，便于后续在线检索

运行一次后，向量库会落盘到 chroma_db/，之后由 online_chat.py 加载使用。
"""
import argparse
import hashlib
import os
from pathlib import Path
from typing import Any, List, Optional, Tuple

import chromadb
import requests
from docling.document_converter import DocumentConverter
from dotenv import load_dotenv

# Minimax 模型配置：直接读取 .env 中的 MINIMAX_* 变量
DEFAULT_EMBEDDING_MODEL = "embo-01"  # Minimax 官方嵌入模型

def parse_page_range(s: str) -> Optional[Tuple[int, int]]:
    """解析 CLI 形式的页码范围字符串，返回 (start, end)。

    支持的格式：
        - "1-10" / "1,10"  -> (1, 10)
        - "" / None        -> None（表示不限制，处理全部页）

    与 Docling 的 page_range 语义一致：start≥1 且 end≥start。
    """
    if not s:
        return None
    for sep in ("-", ","):
        if sep in s:
            start, end = s.split(sep, 1)
            page_range = (int(start), int(end))
            if page_range[0] < 1 or page_range[1] < page_range[0]:
                raise ValueError(
                    f"无效的 page_range: {s}（start 必须 ≥ 1，end 必须 ≥ start）"
                )
            return page_range
    raise ValueError(f"无法解析 page_range: {s}（请使用 '1-10' 或 '1,10' 格式）")

def get_minimax_config() -> Tuple[str, str]:
    """从 .env 读取 Minimax 的 API Key 和 Base URL。"""
    api_key = os.getenv("MINIMAX_API_KEY")
    base_url = os.getenv("MINIMAX_BASE_URL")
    if not api_key:
        raise ValueError("缺少 MINIMAX_API_KEY，请在 .env 中配置。")
    return api_key, base_url

# ---------- 步骤 1：文档解析 ----------
def extract_text_with_docling(
    file_path: Path, page_range: Optional[Tuple[int, int]] = None
) -> str:
    """使用 Docling 解析 PDF/DOCX 等文档并导出为文本。

    Args:
        page_range: 指定解析的页码范围 (start, end)，1-based。
            None 表示不限制（处理全部页）。
            该参数直接透传给 Docling 的 convert(page_range=...)。
    """
    converter = DocumentConverter()
    # page_range 是截断处理（只解析第 start~end 页），
    # 不要用 max_num_pages，那是"超过即拒绝整个文档"的阈值，会导致大文档直接报错。
    convert_kwargs = {}
    if page_range is not None:
        convert_kwargs["page_range"] = page_range
    document = converter.convert(str(file_path), **convert_kwargs).document

    # 优先调用 Markdown 导出，保留更多结构信息
    for method_name in ("export_to_markdown", "export_to_text", "to_markdown", "to_text"):
        method = getattr(document, method_name, None)
        if callable(method):
            text = method()
            if isinstance(text, str) and text.strip():
                return text

    text = str(document)
    if text.strip():
        return text
    raise RuntimeError("Docling 解析成功，但未提取到有效文本。")

# ---------- 步骤 2：文本分片 ----------
def split_text(text: str, chunk_size: int = 500, chunk_overlap: int = 50) -> List[str]:
    """按固定窗口对文本做滑动分片，相邻 chunk 保留重叠以避免语义断裂。"""
    text = text.strip()
    if not text:
        return []
    if chunk_overlap >= chunk_size:
        raise ValueError("chunk_overlap 必须小于 chunk_size。")

    chunks = []
    start = 0
    step = chunk_size - chunk_overlap  # 每次前进的步长
    while start < len(text):
        end = min(start + chunk_size, len(text))
        chunk = text[start:end].strip()
        if chunk:
            chunks.append(chunk)
        if end == len(text):
            break
        start += step
    return chunks

# ---------- 步骤 3：向量化 ----------
def get_embeddings(
    api_key: str,
    base_url: str,
    texts: List[str],
    model: str,
    embed_type: str = "db",
) -> List[List[float]]:
    """调用 Minimax 嵌入 API，把文本批量转为向量。

    注意：Minimax 的嵌入接口与 OpenAI 格式不兼容，不能用 OpenAI SDK 调用：
      - 请求体用 texts（非 input），且需要 type 字段（'db' 建库 / 'query' 查询）
      - 响应体用 vectors（非 data[].embedding）
    因此这里直接发 HTTP 请求。
    """
    response = requests.post(
        f"{base_url}/embeddings",
        headers={
            "Authorization": f"Bearer {api_key}",
            "Content-Type": "application/json",
        },
        json={"model": model, "texts": texts, "type": embed_type},
    )
    response.raise_for_status()
    data = response.json()
    vectors = data.get("vectors")
    if not vectors:
        raise RuntimeError(f"Minimax 嵌入接口返回空向量: {data}")
    return vectors

# ---------- 步骤 4：构建向量库 ----------
def build_collection(
    doc_path: Path,
    collection_name: str,
    persist_dir: Path,
    api_key: str,
    base_url: str,
    embedding_model: str,
    chunk_size: int,
    chunk_overlap: int,
    page_range: Optional[Tuple[int, int]] = None,
) -> Any:
    """解析文档 -> 分片 -> 向量化 -> 写入 Chroma，构建可检索的向量库。"""
    persist_dir.mkdir(parents=True, exist_ok=True)
    chroma_client = chromadb.PersistentClient(path=str(persist_dir))

    # 演示场景下每次重建集合，避免历史数据和向量维度冲突
    try:
        chroma_client.delete_collection(name=collection_name)
    except Exception:
        pass

    collection = chroma_client.get_or_create_collection(
        name=collection_name,
        metadata={"description": "Minimax + Chroma 的最小 RAG 演示"},
    )

    # 步骤 1 + 2：解析 + 分片（page_range 非 None 时只解析指定页码范围）
    text = extract_text_with_docling(doc_path, page_range=page_range)
    chunks = split_text(text, chunk_size=chunk_size, chunk_overlap=chunk_overlap)
    if not chunks:
        raise RuntimeError("文档分片结果为空，无法构建向量库。")

    # 步骤 3：批量向量化（embed_type='db' 表示用于建库的文档向量）
    embeddings = get_embeddings(api_key, base_url, chunks, embedding_model, embed_type="db")

    # 为每个 chunk 生成稳定的 ID 与元数据，便于追溯来源
    doc_hash = hashlib.md5(str(doc_path).encode("utf-8")).hexdigest()[:8]
    ids = [f"{doc_hash}_chunk_{i}" for i in range(len(chunks))]
    metadatas = [
        {"source": str(doc_path), "chunk_index": i, "chunk_size": len(chunks[i])}
        for i in range(len(chunks))
    ]

    # 步骤 4：写入 Chroma
    collection.add(ids=ids, documents=chunks, embeddings=embeddings, metadatas=metadatas)
    print(f"[INFO] 知识库构建完成，共 {len(chunks)} 个 chunk，已写入 {persist_dir}。")
    return collection

def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(description="离线阶段：构建 RAG 知识库")
    parser.add_argument("--doc", required=True, help="文档路径（pdf/docx/txt 等）")
    parser.add_argument("--chunk-size", type=int, default=500, help="分片长度（字符数）")
    parser.add_argument("--chunk-overlap", type=int, default=100, help="分片重叠（字符数）")
    parser.add_argument(
        "--page-range",
        type=parse_page_range,
        default=None,
        metavar="START-END",
        help="只解析指定页码范围，如 '1-10' 或 '1,10'；不传则处理全部页",
    )
    parser.add_argument("--collection", default="chatpdf_demo", help="Chroma 集合名")
    parser.add_argument("--embedding-model", default=DEFAULT_EMBEDDING_MODEL)
    return parser.parse_args()

def main() -> None:
    load_dotenv()
    args = parse_args()

    doc_path = Path(args.doc).expanduser().resolve()
    if not doc_path.exists():
        raise FileNotFoundError(f"文档不存在: {doc_path}")

    api_key, base_url = get_minimax_config()

    # 向量库落盘目录，online_chat.py 会从同一目录加载
    persist_dir = Path(__file__).resolve().parent / "chroma_db"
    build_collection(
        doc_path=doc_path,
        collection_name=args.collection,
        persist_dir=persist_dir,
        api_key=api_key,
        base_url=base_url,
        embedding_model=args.embedding_model,
        chunk_size=args.chunk_size,
        chunk_overlap=args.chunk_overlap,
        page_range=args.page_range,
    )
    print(f"[INFO] 离线知识库已就绪，可运行 online_chat.py 开始问答。")

if __name__ == "__main__":
    main()

"""在线阶段：基于已有知识库进行问答。

RAG 全流程分为离线、在线两阶段，本脚本负责在线部分：
    5. 检索上下文            —— 把用户问题向量化，在向量库中找最相关的 chunk
    6. 生成回答（Chat）      —— 把检索到的上下文塞进 prompt，交给 LLM 回答

前置条件：先运行 offline_knowlegebase.py 构建知识库（chroma_db/）。
"""
import argparse
import os
from pathlib import Path
from typing import Any, List, Tuple

import chromadb
import requests
from dotenv import load_dotenv
from openai import OpenAI

# Minimax 模型配置：直接读取 .env 中的 MINIMAX_* 变量
DEFAULT_CHAT_MODEL = os.getenv("MINIMAX_MODEL", "MiniMax-M3")
DEFAULT_EMBEDDING_MODEL = "embo-01"  # Minimax 官方嵌入模型

def get_minimax_config() -> Tuple[str, str]:
    """从 .env 读取 Minimax 的 API Key 和 Base URL。"""
    api_key = os.getenv("MINIMAX_API_KEY")
    base_url = os.getenv("MINIMAX_BASE_URL")
    if not api_key:
        raise ValueError("缺少 MINIMAX_API_KEY，请在 .env 中配置。")
    return api_key, base_url

def build_chat_client(api_key: str, base_url: str) -> OpenAI:
    """构建 OpenAI 兼容客户端，用于调用 Minimax 的对话接口。

    Minimax 的对话接口（chat completions）兼容 OpenAI 格式，可直接用 SDK；
    但嵌入接口不兼容（请求用 texts+type，响应用 vectors），见 get_embeddings。
    """
    return OpenAI(api_key=api_key, base_url=base_url)

def get_embeddings(
    api_key: str,
    base_url: str,
    texts: List[str],
    model: str,
    embed_type: str = "query",
) -> List[List[float]]:
    """调用 Minimax 嵌入 API，把文本批量转为向量。

    注意：Minimax 的嵌入接口与 OpenAI 格式不兼容，不能用 OpenAI SDK 调用：
      - 请求体用 texts（非 input），且需要 type 字段（'db' 建库 / 'query' 查询）
      - 响应体用 vectors（非 data[].embedding）
    因此这里直接发 HTTP 请求。
    """
    response = requests.post(
        f"{base_url}/embeddings",
        headers={
            "Authorization": f"Bearer {api_key}",
            "Content-Type": "application/json",
        },
        json={"model": model, "texts": texts, "type": embed_type},
    )
    response.raise_for_status()
    data = response.json()
    vectors = data.get("vectors")
    if not vectors:
        raise RuntimeError(f"Minimax 嵌入接口返回空向量: {data}")
    return vectors

def load_collection(persist_dir: Path, collection_name: str) -> Any:
    """从磁盘加载离线阶段构建好的 Chroma 向量库。"""
    chroma_client = chromadb.PersistentClient(path=str(persist_dir))
    return chroma_client.get_collection(name=collection_name)

# ---------- 步骤 5：检索上下文 ----------
def retrieve_context(
    collection: Any,
    api_key: str,
    base_url: str,
    embedding_model: str,
    question: str,
    top_k: int,
) -> str:
    """把问题向量化，在向量库中检索最相关的 top_k 个 chunk，拼成上下文。"""
    # embed_type='query' 表示查询向量，与建库时的 'db' 向量在同一空间可检索
    question_embedding = get_embeddings(
        api_key, base_url, [question], embedding_model, embed_type="query"
    )[0]
    results = collection.query(
        query_embeddings=[question_embedding],
        n_results=top_k,
        include=["documents", "metadatas", "distances"],
    )

    docs = results.get("documents", [[]])[0]
    metadatas = results.get("metadatas", [[]])[0]
    if not docs:
        return ""

    # 标注每个片段的来源 chunk_index，方便定位
    parts = []
    for i, doc in enumerate(docs):
        chunk_index = metadatas[i].get("chunk_index", i) if i < len(metadatas) else i
        parts.append(f"[Chunk {chunk_index}]\n{doc}")
    return "\n\n".join(parts)

# ---------- 步骤 6：生成回答 ----------
def answer_question(client: OpenAI, chat_model: str, context: str, question: str) -> str:
    """把检索到的上下文塞进 prompt，交给 LLM 生成回答。"""
    if not context:
        return "没有检索到相关上下文，无法给出可靠回答。"

    system_prompt = (
        "你是一个文档问答助手。"
        "请仅依据给定的上下文回答。"
        "如果上下文不足以回答，请明确说不知道，不要编造。"
    )
    user_prompt = f"上下文:\n{context}\n\n问题: {question}"

    response = client.chat.completions.create(
        model=chat_model,
        messages=[
            {"role": "system", "content": system_prompt},
            {"role": "user", "content": user_prompt},
        ],
        temperature=0.2,  # 低温度让回答更稳定、更忠于上下文
    )
    return response.choices[0].message.content or ""

def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(description="在线阶段：基于知识库的 RAG 问答")
    parser.add_argument("--top-k", type=int, default=3, help="检索 chunk 数量")
    parser.add_argument("--collection", default="chatpdf_demo", help="Chroma 集合名")
    parser.add_argument("--embedding-model", default=DEFAULT_EMBEDDING_MODEL)
    parser.add_argument("--chat-model", default=DEFAULT_CHAT_MODEL)
    return parser.parse_args()

def main() -> None:
    load_dotenv()
    args = parse_args()

    api_key, base_url = get_minimax_config()
    chat_client = build_chat_client(api_key, base_url)

    # 加载离线阶段构建好的向量库
    persist_dir = Path(__file__).resolve().parent / "chroma_db"
    if not persist_dir.exists():
        raise FileNotFoundError(
            f"向量库不存在: {persist_dir}，请先运行 offline_knowlegebase.py。"
        )
    collection = load_collection(persist_dir, args.collection)

    print("[INFO] RAG 问答已就绪，输入问题开始问答，输入 exit 退出。")
    while True:
        question = input("\n你 > ").strip()
        if not question:
            continue
        if question.lower() in {"exit", "quit", "q"}:
            print("[INFO] 已退出。")
            break

        # 步骤 5：检索上下文
        context = retrieve_context(
            collection=collection,
            api_key=api_key,
            base_url=base_url,
            embedding_model=args.embedding_model,
            question=question,
            top_k=args.top_k,
        )
        # 步骤 6：生成回答
        answer = answer_question(
            client=chat_client,
            chat_model=args.chat_model,
            context=context,
            question=question,
        )
        print("\n助手 >", answer)

if __name__ == "__main__":
    main()