"""MongoDB 连接 + collection 管理

3 个 collection:
  book            — 教辅
  book_document   — 教辅文档
  document_file   — 文档文件

设计思路：
- _id 用业务 ID（doc_id 16 字符 sha1 / book_id 自增字符串 / file_id 自增字符串）
- book_id 用 INT 自增（mongo 中用 find_one_and_update + $inc 实现）
- 不强制 schema（mongo 灵活），但用 helper 函数做字段约束
- 索引在 init_db() 时一次性建好
"""
import os
import yaml
from pymongo import MongoClient, ASCENDING, DESCENDING, ReturnDocument
from pymongo.database import Database

_CONFIG_PATH = os.path.join(os.path.dirname(os.path.dirname(__file__)), 'config.yaml')

_client: MongoClient | None = None


def load_config() -> dict:
    with open(_CONFIG_PATH) as f:
        return yaml.safe_load(f)


def get_client() -> MongoClient:
    global _client
    if _client is None:
        cfg = load_config()['database']
        _client = MongoClient(cfg['url'], serverSelectionTimeoutMS=5000)
    return _client


def get_db() -> Database:
    cfg = load_config()['database']
    return get_client()[cfg['db_name']]


def get_collection(name: str):
    return get_db()[name]


# ---- collection 名常量 ----
COL_BOOK = 'book'
COL_DOC = 'book_document'
COL_FILE = 'document_file'

# ---- counters collection（book_id / file_id 自增用）----
COL_COUNTER = 'counters'


# ---- 索引初始化 ----
def init_db():
    """建所有索引（幂等）"""
    db = get_db()

    # book: zip_sha256 唯一 + (name, version, is_deleted) 唯一
    db[COL_BOOK].create_index([('zip_sha256', ASCENDING)], unique=True, name='uk_book_zip_sha')
    db[COL_BOOK].create_index([('name', ASCENDING), ('version', ASCENDING), ('is_deleted', ASCENDING)],
                              unique=True, name='uk_book_name_version')
    db[COL_BOOK].create_index([('status', ASCENDING)], name='idx_book_status')
    db[COL_BOOK].create_index([('created_at', DESCENDING)], name='idx_book_created')

    # book_document: _id 已唯一（doc_id 即 _id）；额外索引仅用于查询加速
    db[COL_DOC].create_index([('book_id', ASCENDING), ('order_index', ASCENDING)], name='idx_book_order')
    db[COL_DOC].create_index([('status', ASCENDING)], name='idx_doc_status')

    # document_file: (doc_id, file_type) 唯一 + (bucket, key)
    db[COL_FILE].create_index([('doc_id', ASCENDING), ('file_type', ASCENDING)],
                              unique=True, name='uk_doc_type')
    db[COL_FILE].create_index([('object_bucket', ASCENDING), ('object_key', ASCENDING)],
                              name='idx_bucket_key')

    print(f"✓ MongoDB 已就绪（db: {db.name}, collections: {[COL_BOOK, COL_DOC, COL_FILE]}）")
    print(f"  ✓ 8 个索引创建完成")


def health() -> bool:
    try:
        get_client().admin.command('ping')
        return True
    except Exception:
        return False


def next_id(name: str) -> int:
    """原子自增 ID（mongo 没有 AUTO_INCREMENT，用 counters collection 模拟）"""
    db = get_db()
    r = db[COL_COUNTER].find_one_and_update(
        {'_id': name},
        {'$inc': {'seq': 1}},
        upsert=True,
        return_document=ReturnDocument.AFTER,
    )
    return int(r['seq'])