module/qdrant_manager.py

"""Qdrant 向量数据库管理器 — 嵌入式模式（无需独立服务进程）

使用 QdrantClient(path=...) 在本地文件系统中运行 Qdrant，
适合 HF Spaces 单容器部署，无需额外启动 Qdrant 服务。
"""
import os
import hashlib
import uuid

import numpy as np
import torch
from qdrant_client import QdrantClient
from qdrant_client.models import (
    Distance, VectorParams, PointStruct,
    OptimizersConfigDiff, HnswConfigDiff,
)

from module.config import QDRANT_PATH
from module.text_classifier import classify_scene

COLL_PREFIX = 'user_'

_client_instance = None


def _get_client():
    """获取 Qdrant 客户端单例（嵌入式本地文件模式）"""
    global _client_instance
    if _client_instance is None:
        os.makedirs(QDRANT_PATH, exist_ok=True)
        _client_instance = QdrantClient(path=QDRANT_PATH)
    return _client_instance


def pt_id(img_path):
    """根据图片绝对路径生成确定性 UUID"""
    return str(uuid.UUID(hex=hashlib.md5(os.path.abspath(img_path).encode()).hexdigest()))


def _collection_name(username, scene_type):
    return f'{COLL_PREFIX}{username}_{scene_type}'


def _bge_encode(text, bge_tokenizer, bge_model):
    """编码文本为 numpy float32[512]"""
    if not text.strip():
        return np.zeros(512, dtype=np.float32)
    encoded = bge_tokenizer(text, padding=True, truncation=True, return_tensors='pt', max_length=512)
    # CPU 模式：不调用 .cuda()
    with torch.no_grad():
        outputs = bge_model(**encoded)
    emb = outputs.last_hidden_state[:, 0]
    emb = torch.nn.functional.normalize(emb, p=2, dim=1)
    return emb.cpu().numpy()[0]


def _dinov2_encode(img_path, dinov2_extractor):
    """提取 DINOv2 特征向量（384维）"""
    if dinov2_extractor is None or not dinov2_extractor.is_available:
        return np.zeros(384, dtype=np.float32)
    feat = dinov2_extractor.extract_feature(img_path)
    if feat is None:
        return np.zeros(384, dtype=np.float32)
    return feat.cpu().numpy()


# ==================== 预存储: 构建用户缓存 ====================

def build_user_cache(username, user_dir, ocr_engine, dinov2_extractor,
                     bge_tokenizer, bge_model):
    """为用户目录下所有历史图片构建 Qdrant 缓存"""
    from module.file_manager import get_history_images
    client = _get_client()
    images = get_history_images(user_dir)
    if not images:
        return 0, 0

    text_pts, complex_pts = [], []

    for img_path in images:
        try:
            scene_type, doc_score, detail, ocr_result = classify_scene(img_path, ocr_engine)
            pid = pt_id(img_path)
            payload = {'path': img_path, 'filename': os.path.basename(img_path)}

            if scene_type == 'text':
                full_text = detail.get('full_text', '') if detail else ''
                payload['ocr_text'] = full_text
                vector = _bge_encode(full_text, bge_tokenizer, bge_model)
                text_pts.append(PointStruct(id=pid, vector=vector.tolist(), payload=payload))
            else:
                vector = _dinov2_encode(img_path, dinov2_extractor)
                complex_pts.append(PointStruct(id=pid, vector=vector.tolist(), payload=payload))

            print(f"  [qdrant] {os.path.basename(img_path):<30s} {scene_type}")
        except Exception as e:
            print(f"  [qdrant] {os.path.basename(img_path)} ERROR: {e}")

    for suffix, pts, dim in [('text', text_pts, 512), ('complex', complex_pts, 384)]:
        if not pts:
            continue
        cname = _collection_name(username, suffix)
        if client.collection_exists(cname):
            client.delete_collection(cname)
        client.create_collection(
            cname,
            vectors_config=VectorParams(size=dim, distance=Distance.COSINE),
            optimizers_config=OptimizersConfigDiff(indexing_threshold=0),
        )
        client.upsert(cname, points=pts, wait=True)

    for suffix in ['text', 'complex']:
        cname = _collection_name(username, suffix)
        if client.collection_exists(cname):
            client.update_collection(cname, hnsw_config=HnswConfigDiff(m=16, ef_construct=100))

    return len(text_pts), len(complex_pts)


# ==================== 查询: 从 Qdrant 拉取缓存 ====================

def qdrant_query_scene(username, scene_type):
    """从 Qdrant 拉取指定场景类型的用户缓存（scroll 一次，返回 dict）"""
    client = _get_client()
    cache = {}
    cname = _collection_name(username, scene_type)
    if not client.collection_exists(cname):
        return cache
    records, _ = client.scroll(cname, limit=1000, with_payload=True, with_vectors=True)
    for rec in records:
        entry = dict(rec.payload)
        entry['_vector'] = np.array(rec.vector, dtype=np.float32)
        if scene_type == 'text':
            entry['bge_vector'] = entry['_vector']
        cache[rec.id] = entry
    return cache


def qdrant_delete_user_collections(username):
    """删除指定用户的所有 Qdrant collection（text + complex）

    用于用户历史库为空时清理残留数据
    """
    client = _get_client()
    for scene_type in ('text', 'complex'):
        cname = _collection_name(username, scene_type)
        try:
            if client.collection_exists(cname):
                client.delete_collection(cname)
                print(f"[Qdrant] 已删除 collection: {cname}")
        except Exception as e:
            print(f"[Qdrant] 删除 collection {cname} 失败: {e}")


# ==================== 向量相似度搜索 ====================

def qdrant_search_similar(username, scene_type, query_vector, top_k=3, score_threshold=0.0):
    client = _get_client()
    cname = _collection_name(username, scene_type)
    if not client.collection_exists(cname):
        return []
    response = client.query_points(
        collection_name=cname,
        query=query_vector.tolist(),
        limit=top_k,
        score_threshold=score_threshold,
        with_payload=True,
        with_vectors=True,
    )
    results = []
    for hit in response.points:
        entry = dict(hit.payload)
        entry['_vector'] = np.array(hit.vector, dtype=np.float32)
        if scene_type == 'text':
            entry['bge_vector'] = entry['_vector']
        results.append((hit.score, hit.id, entry))
    return results


# ==================== 增量更新 ====================

def add_to_qdrant(username, img_path, ocr_engine, dinov2_extractor,
                  bge_tokenizer, bge_model, scene_type=None,
                  precomputed_text=None, precomputed_bge=None,
                  precomputed_dinov2=None):
    """将新图片添加到 Qdrant（增量插入）"""
    client = _get_client()
    if scene_type is None:
        scene_type, _, detail, _ = classify_scene(img_path, ocr_engine)
        # 复用 classify_scene 的 OCR 结果，避免对同一张图重复 OCR
        if precomputed_text is None and scene_type == 'text':
            precomputed_text = detail.get('full_text', '')
    pid = pt_id(img_path)
    payload = {'path': img_path, 'filename': os.path.basename(img_path)}

    if scene_type == 'text':
        if precomputed_text is not None:
            full_text = precomputed_text
        else:
            from module.text_matcher import extract_text
            full_text, _ = extract_text(img_path, ocr_engine)
        payload['ocr_text'] = full_text
        if precomputed_bge is not None:
            vector = precomputed_bge
        else:
            vector = _bge_encode(full_text, bge_tokenizer, bge_model)
        dim = 512
    else:
        if precomputed_dinov2 is not None:
            vector = precomputed_dinov2
        else:
            vector = _dinov2_encode(img_path, dinov2_extractor)
        dim = 384

    cname = _collection_name(username, scene_type)
    if not client.collection_exists(cname):
        client.create_collection(
            cname,
            vectors_config=VectorParams(size=dim, distance=Distance.COSINE),
            optimizers_config=OptimizersConfigDiff(indexing_threshold=0),
        )
        client.update_collection(cname, hnsw_config=HnswConfigDiff(m=16, ef_construct=100))

    client.upsert(cname, points=[PointStruct(id=pid, vector=vector.tolist(), payload=payload)], wait=True)
    return scene_type


def remove_from_qdrant(username, img_path, scene_type=None):
    """从 Qdrant 中移除指定图片"""
    client = _get_client()
    pid = pt_id(img_path)

    if scene_type:
        cname = _collection_name(username, scene_type)
        if client.collection_exists(cname):
            try:
                client.delete(cname, points=[pid])
            except Exception:
                pass
    else:
        for st in ['text', 'complex']:
            cname = _collection_name(username, st)
            if client.collection_exists(cname):
                try:
                    client.delete(cname, points=[pid])
                except Exception:
                    pass


def has_qdrant_cache(username):
    """检查用户是否有 Qdrant 缓存"""
    client = _get_client()
    for st in ['text', 'complex']:
        cname = _collection_name(username, st)
        if client.collection_exists(cname):
            count = client.count(cname).count
            if count > 0:
                return True
    return False