from __future__ import annotations

import re
from pathlib import Path

from app.config import Settings, get_settings
from app.embeddings import create_embedding_provider
from app.milvus_store import IndexedDocument, MilvusVectorStore
from app.sql_store import initialize_database


# 受控元数据键集合：每个键对应 Milvus 过滤字段，全部允许 `_a-zA-Z0-9-`。
DOC_METADATA_KEYS = ("文档类型", "服务名称", "适用系统")


def split_document(path: Path) -> list[IndexedDocument]:
    """读取 Markdown 排障文档，按 ## 二级标题切 chunk，并提取受控元数据。"""
    text = path.read_text(encoding="utf-8").strip()
    if not text:
        raise ValueError(f"Milvus 原始文档为空：{path}")

    # 只抽取标准元数据键，其它含全角冒号的行不进字典，避免污染 policy_type。
    metadata_lines: dict[str, str] = {}
    for line in text.splitlines():
        if "：" not in line:
            continue
        key, value = line.split("：", 1)
        if key.strip() in DOC_METADATA_KEYS:
            metadata_lines[key.strip()] = value.strip()

    doc_type = metadata_lines.get("文档类型", "")
    service_name = metadata_lines.get("服务名称", "")
    client_os = metadata_lines.get("适用系统", "")
    if not doc_type or not service_name or not client_os:
        raise ValueError(
            f"文档缺少受控元数据：{path.name}，需要 "
            f"{'/'.join(DOC_METADATA_KEYS)}"
        )

    # 课程文档以二级标题作为稳定 Chunk 边界，保留章节语义和来源定位。
    sections = [
        section.strip()
        for section in re.split(r"(?=^## )", text, flags=re.MULTILINE)
        if section.strip()
    ]
    title = text.splitlines()[0].lstrip("# ")
    return [
        IndexedDocument(
            content=section,
            source=path.name,
            doc_type=doc_type,
            chunk_index=index,
            # policy_type 在 IT 域里复用为故障类型语义，便于检索面兼容。
            policy_type="troubleshooting_guide",
            category=client_os,
            service_name=service_name,
            metadata={"title": title},
        )
        for index, section in enumerate(sections)
    ]


def prepare_sqlite(settings: Settings) -> int:
    return initialize_database(
        path=settings.sqlite_path,
        source_path=settings.orders_source_path,
        reset=True,
    )


def prepare_milvus(settings: Settings) -> tuple[int, int, int]:
    document_paths = sorted(settings.documents_path.glob("*.md"))
    if not document_paths:
        raise FileNotFoundError(f"Milvus 原始文档不存在：{settings.documents_path}")

    # 先完成所有文档校验和切分，再连接 Milvus，避免半批脏数据。
    documents = [
        document
        for path in document_paths
        for document in split_document(path)
    ]
    embeddings = create_embedding_provider(settings)
    store = MilvusVectorStore(
        uri=settings.milvus_uri,
        token=settings.milvus_token.get_secret_value(),
        collection_name=settings.milvus_collection,
        embeddings=embeddings,
    )
    # 演示数据采用全量重建，确保 Schema、索引和向量维度保持一致。
    store.create_collection(recreate=True)
    inserted = store.insert_documents(documents)
    return len(document_paths), inserted, embeddings.dimension


def main() -> None:
    settings = get_settings()
    print("数据准备开始")
    print(f"- SQLite 原始工单：{settings.orders_source_path}")
    print(f"- Milvus 原始文档：{settings.documents_path}")

    sqlite_rows = prepare_sqlite(settings)
    document_count, chunk_count, dimension = prepare_milvus(settings)

    print(
        f"SQLite 完成：database={settings.sqlite_path}，rows={sqlite_rows}"
    )
    print(
        "Milvus 完成："
        f"collection={settings.milvus_collection}，"
        f"documents={document_count}，chunks={chunk_count}，dimension={dimension}"
    )


if __name__ == "__main__":
    main()