app.py

"""FastAPI 主服务 — HF Spaces CPU 部署版

改动点：
1. 端口 7860（HF Spaces）
2. MASt3R 已禁用，SceneMatcher 使用 DINOv2-only
3. Qdrant 嵌入式本地文件模式
4. 历史图片管理 API
5. 图片文件服务 API
6. 可视化在删除文件之前生成，支持复杂+文本场景
7. 可视化按需展示（前端点击按钮触发）
"""
import os
import asyncio
from concurrent.futures import ThreadPoolExecutor
from contextlib import asynccontextmanager
import time
import base64

from fastapi import FastAPI, UploadFile, File, Form, HTTPException
from fastapi.middleware.cors import CORSMiddleware
from fastapi.responses import FileResponse, JSONResponse

from module.model_loader import load_model, load_dinov2, load_ocr, load_bge
from module.scene_matcher import SceneMatcher
from module.file_manager import (
    validate_username, get_user_dir, get_history_images,
    save_uploaded_file, delete_file
)
from module.evaluator import (
    evaluate_scene_consistency, evaluate_scene_consistency_cached, update_metadata
)
from module.system_monitor import _snapshot, ResourceMonitor, write_session_log
from module.qdrant_manager import (
    build_user_cache, qdrant_query_scene, has_qdrant_cache,
    add_to_qdrant, remove_from_qdrant,
)
from module.qdrant_manager import qdrant_search_similar, pt_id
from module.text_classifier import classify_scene
from module.demo_images import ensure_demo_images

BASE_DIR = os.path.dirname(os.path.abspath(__file__))
FRONTEND_PATH = os.path.join(BASE_DIR, 'index.html')
HISTORY_BASE = os.path.join(BASE_DIR, 'History_imgs')

matcher = None
dinov2_extractor = None
ocr_engine = None
bge_tokenizer = None
bge_model = None
executor = ThreadPoolExecutor(max_workers=2)

session_data = {
    'startup_stats': None,
    'model_loaded_stats': None,
    'detections': [],
}

resource_monitor = ResourceMonitor(sample_interval=0.5)

_qd_caches = {}


@asynccontextmanager
async def lifespan(app):
    global matcher, dinov2_extractor, ocr_engine, bge_tokenizer, bge_model, session_data

    startup_t0 = time.perf_counter()
    session_data['startup_stats'] = _snapshot()

    model = load_model()
    dinov2_extractor = load_dinov2()
    matcher = SceneMatcher(model, dinov2_extractor=dinov2_extractor)

    ocr_engine = load_ocr()
    bge_tokenizer, bge_model = load_bge()

    session_data['model_loaded_stats'] = _snapshot()
    model_load_time = time.perf_counter() - startup_t0
    print(f'[启动] 模型加载完成: {model_load_time:.1f}s')

    qdrant_t0 = time.perf_counter()
    if os.path.isdir(HISTORY_BASE):
        for uname in os.listdir(HISTORY_BASE):
            udir = os.path.join(HISTORY_BASE, uname)
            if not os.path.isdir(udir):
                continue
            imgs = get_history_images(udir)
            if not imgs:
                continue
            try:
                if not has_qdrant_cache(uname):
                    print(f"[qdrant] 启动预构建: 用户 [{uname}]")
                    build_user_cache(uname, udir, ocr_engine, dinov2_extractor,
                                     bge_tokenizer, bge_model)
                _qd_caches[uname] = {
                    'text': qdrant_query_scene(uname, 'text'),
                    'complex': qdrant_query_scene(uname, 'complex'),
                }
                _t_n = len(_qd_caches[uname]['text'])
                _c_n = len(_qd_caches[uname]['complex'])
                print(f"  [{uname}] 缓存已就绪 (text={_t_n}, complex={_c_n})")
            except Exception as e:
                print(f"  [{uname}] 缓存加载失败: {e}")
    qdrant_load_time = time.perf_counter() - qdrant_t0
    print(f"[qdrant] 缓存检查完成: {len(_qd_caches)} 个用户 ({qdrant_load_time:.1f}s)")
    total_startup = time.perf_counter() - startup_t0
    print(f'[启动] 总启动时间: {total_startup:.1f}s')

    # 生成示例场景图片（用于 /api/demo 端点）
    ensure_demo_images(BASE_DIR)

    yield

    session_data['shutdown_stats'] = _snapshot()
    try:
        log_path = write_session_log(session_data)
        print(f"日志已写入: {log_path}")
    except Exception as e:
        print(f"日志写入失败: {e}")


app = FastAPI(lifespan=lifespan)

app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_methods=["*"],
    allow_headers=["*"],
)


@app.get("/")
async def serve_frontend():
    return FileResponse(FRONTEND_PATH)


@app.get("/cover.png")
async def serve_cover():
    """提供背景图片 cover.png"""
    cover_path = os.path.join(BASE_DIR, 'cover.png')
    if os.path.exists(cover_path):
        return FileResponse(cover_path, media_type='image/png')
    raise HTTPException(status_code=404, detail="cover not found")


@app.post("/api/login")
async def login(username: str = Form(...)):
    try:
        uname = validate_username(username)
    except ValueError as e:
        raise HTTPException(status_code=400, detail=str(e))

    user_dir = get_user_dir(uname)
    history = get_history_images(user_dir)

    cached = uname in _qd_caches
    if not cached and history and has_qdrant_cache(uname):
        _qd_caches[uname] = {
            'text': qdrant_query_scene(uname, 'text'),
            'complex': qdrant_query_scene(uname, 'complex'),
        }
        cached = True

    return {
        "username": uname,
        "history_count": len(history),
        "cached": cached,
    }


# ==================== 历史图片管理 API ====================

@app.get("/api/history")
async def list_history(username: str):
    try:
        uname = validate_username(username)
    except ValueError as e:
        raise HTTPException(status_code=400, detail=str(e))

    user_dir = get_user_dir(uname)
    images = get_history_images(user_dir)

    result = []
    for img_path in images:
        fname = os.path.basename(img_path)
        fsize = os.path.getsize(img_path) if os.path.exists(img_path) else 0
        result.append({
            'filename': fname,
            'size_kb': round(fsize / 1024, 1),
            'url': f'/api/image/{uname}/{fname}',
        })
    return {"username": uname, "images": result, "count": len(result)}


@app.get("/api/all_users")
async def list_all_users():
    users = []
    if os.path.isdir(HISTORY_BASE):
        for uname in sorted(os.listdir(HISTORY_BASE)):
            udir = os.path.join(HISTORY_BASE, uname)
            if not os.path.isdir(udir):
                continue
            imgs = get_history_images(udir)
            if imgs:
                users.append({
                    'username': uname,
                    'image_count': len(imgs),
                })
    return {"users": users}


@app.get("/api/image/{username}/{filename}")
async def serve_image(username: str, filename: str):
    uname = validate_username(username)
    safe_name = os.path.basename(filename)
    user_dir = get_user_dir(uname)
    img_path = os.path.join(user_dir, safe_name)

    if not os.path.exists(img_path):
        raise HTTPException(status_code=404, detail="图片不存在")

    if not os.path.abspath(img_path).startswith(os.path.abspath(user_dir)):
        raise HTTPException(status_code=403, detail="无权访问")

    return FileResponse(img_path)


@app.post("/api/delete_image")
async def delete_image(username: str = Form(...), filename: str = Form(...)):
    global _qd_caches
    try:
        uname = validate_username(username)
    except ValueError as e:
        raise HTTPException(status_code=400, detail=str(e))

    safe_name = os.path.basename(filename)
    user_dir = get_user_dir(uname)
    img_path = os.path.join(user_dir, safe_name)

    if not os.path.exists(img_path):
        raise HTTPException(status_code=404, detail="图片不存在")

    if not os.path.abspath(img_path).startswith(os.path.abspath(user_dir)):
        raise HTTPException(status_code=403, detail="无权操作")

    try:
        remove_from_qdrant(uname, img_path, scene_type=None)
        if uname in _qd_caches:
            _qd_caches[uname] = {
                'text': qdrant_query_scene(uname, 'text'),
                'complex': qdrant_query_scene(uname, 'complex'),
            }
    except Exception as e:
        print(f"  [删除] Qdrant 清理失败: {e}")

    delete_file(img_path)

    from module.file_manager import load_metadata, save_metadata
    metadata = load_metadata(user_dir)
    metadata['images'] = [m for m in metadata['images']
                          if m.get('filename') != safe_name]
    save_metadata(user_dir, metadata)

    remaining = get_history_images(user_dir)
    return {
        "success": True,
        "deleted": safe_name,
        "remaining_count": len(remaining),
    }


# ==================== 检测 API ====================

def _run_detect_pipeline(uploaded_paths, uname, ocr_engine, bge_tokenizer,
                         bge_model, dinov2_extractor, matcher, history_before,
                         scene_mode='auto'):
    scene_type_q = 'complex'
    ocr_result_q = None
    precomputed_text = None
    precomputed_bge = None
    precomputed_dinov2 = None

    _t0 = time.perf_counter()
    if scene_mode in ('text', 'complex'):
        # 用户手动指定场景，跳过 classify_scene
        scene_type_q = scene_mode
        print(f'  [SKIP] classify_scene (用户指定场景: {scene_mode})')
        if scene_mode == 'text' and ocr_engine is not None and uploaded_paths:
            # 文本场景仍需 OCR 提取文本
            _, _, detail, ocr_result_q = classify_scene(uploaded_paths[0], ocr_engine)
            precomputed_text = detail.get('full_text', '')
            if precomputed_text and precomputed_text.strip():
                from module.qdrant_manager import _bge_encode
                precomputed_bge = _bge_encode(precomputed_text, bge_tokenizer, bge_model)
                _t1 = time.perf_counter()
                print(f'  [TIME] ocr+bge (skip classify): {_t1-_t0:.3f}s')
    elif ocr_engine is not None and uploaded_paths:
        scene_type_q, _, _, ocr_result_q = classify_scene(uploaded_paths[0], ocr_engine)
        _t1 = time.perf_counter()
        print(f'  [TIME] classify_scene: {_t1-_t0:.3f}s')
        if ocr_result_q is not None and ocr_result_q and ocr_result_q[0]:
            precomputed_text = ' '.join(line[1][0] for line in ocr_result_q[0])
        if precomputed_text and precomputed_text.strip():
            from module.qdrant_manager import _bge_encode
            precomputed_bge = _bge_encode(precomputed_text, bge_tokenizer, bge_model)
            _t2 = time.perf_counter()
            print(f'  [TIME] bge_encode: {_t2-_t1:.3f}s')

    user_qd = _qd_caches.get(uname, {})
    has_cache = isinstance(user_qd, dict) and bool(user_qd)

    cache_used = False
    final_result = True
    evaluation_logs = []
    best_match = None

    if has_cache:
        cache_used = True
        final_result, evaluation_logs, best_match = evaluate_scene_consistency_cached(
            matcher, uploaded_paths, uname,
            ocr_engine, bge_tokenizer, bge_model,
            scene_type_q, qdrant_search_similar, dinov2_extractor,
            precomputed_ocr_result=ocr_result_q,
            precomputed_bge_vec=precomputed_bge,
            qd_cache=user_qd,
        )

    if not cache_used:
        final_result, evaluation_logs, best_match = evaluate_scene_consistency(
            matcher, uploaded_paths, history_before,
            ocr_engine, bge_tokenizer, bge_model,
        )

    same_scene_paths = set()
    for log in evaluation_logs:
        if log['is_same_scene'] and log['query_path'] in uploaded_paths:
            same_scene_paths.add(log['query_path'])

    # ★ 先把需要可视化的图片数据读入内存，再删除文件 ★
    # 这样可视化构建不依赖磁盘文件，且不会被删除操作影响
    _image_cache = {}
    for log in evaluation_logs:
        if not log.get('is_same_scene'):
            continue
        for key in ('query_path', 'history_path'):
            p = log.get(key, '')
            if p and p not in _image_cache and os.path.exists(p):
                try:
                    import cv2, numpy as np
                    data = np.fromfile(p, dtype=np.uint8)
                    img = cv2.imdecode(data, cv2.IMREAD_COLOR)
                    if img is not None:
                        _image_cache[p] = img
                except Exception:
                    pass

    # 删除同一场景的文件
    for path in same_scene_paths:
        delete_file(path)

    # 构建可视化（使用内存中的图片数据）
    visualization_data = []
    for log in evaluation_logs:
        if not log.get('is_same_scene'):
            continue
        scene_type = log.get('scene_type', 'complex')
        try:
            if scene_type == 'complex':
                vis = _build_complex_vis_from_cache(log, _image_cache)
                if vis:
                    visualization_data.append(vis)
            elif scene_type == 'text':
                vis = _build_text_vis_from_cache(log, _image_cache)
                if vis:
                    visualization_data.append(vis)
        except Exception as e:
            print(f"  [可视化] 生成失败: {e}")

    per_image_results = {}
    for orig, path in [(os.path.basename(p), p) for p in uploaded_paths]:
        if path in same_scene_paths:
            per_image_results[path] = '未通过（同一场景，已删除）'
        else:
            per_image_results[path] = '通过' if final_result else '未通过'

    return final_result, evaluation_logs, best_match, scene_type_q, same_scene_paths, per_image_results, cache_used, precomputed_text, precomputed_bge, precomputed_dinov2, visualization_data


@app.post("/api/detect")
async def detect(
    username: str = Form(...),
    scene_mode: str = Form(default='auto'),
    images: list[UploadFile] = File(default=[]),
):
    global session_data, _qd_caches

    try:
        uname = validate_username(username)
    except ValueError as e:
        raise HTTPException(status_code=400, detail=str(e))

    user_dir = get_user_dir(uname)

    if len(images) == 0:
        raise HTTPException(status_code=400, detail="请上传至少一张图片")

    history_before = get_history_images(user_dir)

    saved_info = []
    for f in images:
        data = await f.read()
        if not data:
            continue
        path = save_uploaded_file(data, f.filename, user_dir)
        saved_info.append((f.filename, path))

    if not saved_info:
        raise HTTPException(status_code=400, detail="没有有效的上传图片")

    uploaded_paths = [p for _, p in saved_info]

    if len(history_before) == 0:
        # 清理可能的 Qdrant 残留数据（用户可能通过历史管理删除了所有图片）
        try:
            from module.qdrant_manager import qdrant_delete_user_collections
            qdrant_delete_user_collections(uname)
        except Exception:
            pass
        _qd_caches.pop(uname, None)

        per_image_results = {p: '通过（首次上传，跳过比对）' for _, p in saved_info}
        update_metadata(user_dir, saved_info, per_image_results)
        try:
            for _, p in saved_info:
                add_to_qdrant(uname, p, ocr_engine, dinov2_extractor,
                              bge_tokenizer, bge_model)
            _qd_caches[uname] = {
                'text': qdrant_query_scene(uname, 'text'),
                'complex': qdrant_query_scene(uname, 'complex'),
            }
        except Exception:
            pass
        return {
            "result": "pass",
            "reason": "first_upload",
            "saved_files": [
                {"filename": os.path.basename(p), "original_name": orig, "deleted": False}
                for orig, p in saved_info
            ],
            "evaluation_logs": [],
            "best_match": None,
            "visualizations": [],
        }

    pre_stats = _snapshot()
    detect_start_time = asyncio.get_event_loop().time()
    resource_monitor.start_monitoring()

    loop = asyncio.get_event_loop()
    final_result, evaluation_logs, best_match, scene_type_q, same_scene_paths, per_image_results, cache_used, precomp_text, precomp_bge, precomp_dino, visualization_data = \
        await loop.run_in_executor(
            executor, _run_detect_pipeline,
            uploaded_paths, uname, ocr_engine, bge_tokenizer,
            bge_model, dinov2_extractor, matcher, history_before,
            scene_mode,
        )

    peak = resource_monitor.stop_monitoring()
    detect_duration = asyncio.get_event_loop().time() - detect_start_time

    kept_saved_info = [(orig, p) for orig, p in saved_info if p not in same_scene_paths]
    update_metadata(user_dir, kept_saved_info, per_image_results)

    kept_paths_qd = [p for _, p in kept_saved_info]
    removed_paths_qd = list(same_scene_paths)
    st_q = scene_type_q

    def _post_update_qdrant():
        for path in kept_paths_qd:
            try:
                add_to_qdrant(uname, path, ocr_engine, dinov2_extractor,
                              bge_tokenizer, bge_model, scene_type=st_q,
                              precomputed_text=precomp_text,
                              precomputed_bge=precomp_bge,
                              precomputed_dinov2=precomp_dino)
                _qd_caches[uname] = {
                    'text': qdrant_query_scene(uname, 'text'),
                    'complex': qdrant_query_scene(uname, 'complex'),
                }
            except Exception:
                pass
        for path in removed_paths_qd:
            try:
                remove_from_qdrant(uname, path, scene_type=None)
            except Exception:
                pass

    loop.run_in_executor(executor, _post_update_qdrant)

    best_match_response = None
    if best_match:
        best_match_response = {
            'query_image': best_match['query_image'],
            'history_image': best_match['history_image'],
            'similarity_score': best_match['similarity_score'],
            'match_count': best_match['match_count'],
            'inlier_ratio': best_match['inlier_ratio'],
            'avg_confidence': best_match['avg_confidence'],
        }

    return {
        "result": "pass" if final_result else "fail",
        "saved_files": [
            {"filename": os.path.basename(p), "original_name": orig, "deleted": p in same_scene_paths}
            for orig, p in saved_info
        ],
        "evaluation_logs": [
            {
                "query_image": log['query_image'],
                "history_image": log['history_image'],
                "scene_type": log.get('scene_type', 'complex'),
                "similarity_score": log['similarity_score'],
                "match_count": log['match_count'],
                "dinov2_similarity": log.get('dinov2_similarity'),
                "text_similarity": log.get('text_similarity'),
                "bge_search_score": log.get('bge_search_score'),
                "is_same_scene": log['is_same_scene'],
                "error": log.get('error'),
            }
            for log in evaluation_logs
        ],
        "best_match": best_match_response,
        "visualizations": visualization_data,
    }


def _build_complex_vis_from_cache(log_entry, image_cache):
    """复杂场景可视化：DINOv2 patch 匹配连线

    图片数据从 image_cache 中获取（内存），不依赖磁盘文件。
    """
    try:
        import cv2
        import numpy as np

        q_path = log_entry.get('query_path', '')
        h_path = log_entry.get('history_path', '')
        patch_info = log_entry.get('patch_match_info')

        img1 = image_cache.get(q_path)
        img2 = image_cache.get(h_path)
        if img1 is None or img2 is None:
            return None
        if patch_info is None:
            return None

        orig_h1, orig_w1 = img1.shape[:2]
        orig_h2, orig_w2 = img2.shape[:2]

        max_h = 600
        scale1 = max_h / orig_h1
        scale2 = max_h / orig_h2
        img1_disp = cv2.resize(img1, (int(orig_w1*scale1), max_h))
        img2_disp = cv2.resize(img2, (int(orig_w2*scale2), max_h))

        h1, w1 = img1_disp.shape[:2]
        h2, w2 = img2_disp.shape[:2]

        matches = patch_info.get('matches', [])
        q_grid = patch_info.get('query_grid', (37, 37))
        h_grid = patch_info.get('hist_grid', (37, 37))
        q_orig_size = patch_info.get('query_image_size', (orig_w1, orig_h1))
        h_orig_size = patch_info.get('hist_image_size', (orig_w2, orig_h2))

        q_crop = patch_info.get('query_crop_offset', (0, 0))
        h_crop = patch_info.get('hist_crop_offset', (0, 0))
        q_resize = patch_info.get('query_resize_size', (518, 518))
        h_resize = patch_info.get('hist_resize_size', (518, 518))

        q_n_h, q_n_w = q_grid
        h_n_h, h_n_w = h_grid
        patch_size = 14

        q_scale_resize = q_resize[1] / q_orig_size[1]
        h_scale_resize = h_resize[1] / h_orig_size[1]

        gap = 20
        canvas = np.ones((max_h, w1 + gap + w2, 3), dtype=np.uint8) * 240
        canvas[:h1, :w1] = img1_disp
        canvas[:h2, w1+gap:] = img2_disp

        def _patch_center(idx, n_w, crop_off, scale_r, disp_s, is_right, gap_off):
            row, col = idx // n_w, idx % n_w
            cx = int(((col + 0.5) * patch_size + crop_off[0]) / scale_r * disp_s)
            cy = int(((row + 0.5) * patch_size + crop_off[1]) / scale_r * disp_s)
            if is_right:
                cx += gap_off
            return cx, cy

        for qi, hi, sim in matches:
            q_cx, q_cy = _patch_center(qi, q_n_w, q_crop, q_scale_resize, scale1, False, 0)
            h_cx, h_cy = _patch_center(hi, h_n_w, h_crop, h_scale_resize, scale2, True, w1 + gap)

            intensity = min(1.0, max(0.0, (sim - 0.3) / 0.7))
            color = (0, int(200 * intensity + 55), int(255 * (1 - intensity)))
            cv2.line(canvas, (q_cx, q_cy), (h_cx, h_cy), color, 1, cv2.LINE_AA)

        n_matches = len(matches)
        avg_sim = sum(s for _, _, s in matches) / max(n_matches, 1)

        _, buffer = cv2.imencode('.jpg', canvas, [cv2.IMWRITE_JPEG_QUALITY, 95])
        img_b64 = base64.b64encode(buffer).decode('utf-8')

        return {
            'scene_type': 'complex',
            'query_image': log_entry['query_image'],
            'history_image': log_entry['history_image'],
            'dinov2_similarity': patch_info.get('cls_similarity', 0),
            'patch_match_count': n_matches,
            'avg_patch_similarity': round(avg_sim, 4),
            'image_base64': f'data:image/jpeg;base64,{img_b64}',
        }
    except Exception as e:
        print(f"  [可视化-复杂] 生成失败: {e}")
        return None


def _build_text_vis_from_cache(log_entry, image_cache):
    """文本场景可视化：两张图并排 + 标注

    图片数据从 image_cache 中获取（内存），不依赖磁盘文件。
    """
    try:
        import cv2
        import numpy as np
        from PIL import Image, ImageDraw, ImageFont
        import re

        q_path = log_entry.get('query_path', '')
        h_path = log_entry.get('history_path', '')

        img1 = image_cache.get(q_path)
        img2 = image_cache.get(h_path)
        if img1 is None or img2 is None:
            return None

        text1 = log_entry.get('query_text', '')
        text2 = log_entry.get('history_text', '')
        text_sim = log_entry.get('text_similarity', 0)

        words1 = set(re.findall(r'[a-zA-Z\u4e00-\u9fff]{2,}', text1))
        words2 = set(re.findall(r'[a-zA-Z\u4e00-\u9fff]{2,}', text2))
        common_words = words1 & words2
        diff_words1 = words1 - words2
        diff_words2 = words2 - words1

        max_h = 600
        scale1 = max_h / img1.shape[0]
        scale2 = max_h / img2.shape[0]
        img1_s = cv2.resize(img1, (int(img1.shape[1]*scale1), max_h))
        img2_s = cv2.resize(img2, (int(img2.shape[1]*scale2), max_h))

        h1, w1 = img1_s.shape[:2]
        h2, w2 = img2_s.shape[:2]

        gap = 10
        total_w = w1 + gap + w2
        canvas = np.ones((max_h, total_w, 3), dtype=np.uint8) * 240
        canvas[:max_h, :w1] = img1_s
        canvas[:max_h, w1+gap:w1+gap+w2] = img2_s

        canvas_rgb = cv2.cvtColor(canvas, cv2.COLOR_BGR2RGB)
        pil_img = Image.fromarray(canvas_rgb)
        draw = ImageDraw.Draw(pil_img)

        try:
            font_large = ImageFont.truetype("/usr/share/fonts/truetype/dejavu/DejaVuSans.ttf", 16)
        except Exception:
            try:
                font_large = ImageFont.truetype("DejaVuSans.ttf", 16)
            except Exception:
                font_large = ImageFont.load_default()

        draw.text((10, 8), "Query (OCR text)", fill=(0, 0, 0), font=font_large)
        draw.text((w1+gap+10, 8), "History (OCR text)", fill=(0, 0, 0), font=font_large)

        common_list = sorted(common_words)[:20]

        canvas = cv2.cvtColor(np.array(pil_img), cv2.COLOR_RGB2BGR)
        _, buffer = cv2.imencode('.jpg', canvas, [cv2.IMWRITE_JPEG_QUALITY, 95])
        img_b64 = base64.b64encode(buffer).decode('utf-8')

        return {
            'scene_type': 'text',
            'query_image': log_entry['query_image'],
            'history_image': log_entry['history_image'],
            'text_similarity': round(text_sim, 4),
            'common_keyword_count': len(common_words),
            'diff_keyword_count_query': len(diff_words1),
            'diff_keyword_count_history': len(diff_words2),
            'common_keywords': common_list,
            'image_base64': f'data:image/jpeg;base64,{img_b64}',
        }
    except Exception as e:
        print(f"  [可视化-文本] 生成失败: {e}")
        return None


@app.get("/api/status")
async def status():
    return {
        "model_loaded": matcher is not None,
        "device": matcher.device if matcher else None,
        "ocr_available": ocr_engine is not None,
        "bge_available": bge_tokenizer is not None,
        "qdrant_cached_users": len(_qd_caches),
    }


@app.post("/api/demo")
async def demo(scene_type: str = Form("auto")):
    """示例场景检测：自动读取 test_imgs 中的示例图片执行检测

    Args:
        scene_type: 'text' / 'complex' / 'auto'（auto 表示自动选择）
    """
    import shutil
    demo_map = ensure_demo_images(BASE_DIR)
    if not demo_map:
        return JSONResponse({"error": "示例图片生成失败"}, status_code=500)

    # 选择场景类型
    if scene_type == 'auto':
        scene_type = 'text'
    if scene_type not in demo_map:
        return JSONResponse({"error": f"不支持的场景类型: {scene_type}"}, status_code=400)

    img_paths = demo_map[scene_type]
    if len(img_paths) < 2:
        return JSONResponse({"error": "示例图片不足"}, status_code=500)

    # 复制到临时目录执行检测（不污染任何用户的历史数据）
    import tempfile
    tmp_dir = tempfile.mkdtemp(prefix='demo_')
    tmp_imgs = []
    for i, p in enumerate(img_paths[:2]):
        dst = os.path.join(tmp_dir, f"demo_{i}.jpg")
        shutil.copy(p, dst)
        tmp_imgs.append(dst)

    try:
        loop = asyncio.get_event_loop()
        result = await loop.run_in_executor(
            executor,
            _run_demo_pipeline,
            tmp_imgs, scene_type
        )
        return result
    except Exception as e:
        import traceback
        traceback.print_exc()
        return JSONResponse({"error": f"示例检测失败: {str(e)}"}, status_code=500)
    finally:
        # 清理临时目录
        try:
            shutil.rmtree(tmp_dir, ignore_errors=True)
        except Exception:
            pass


def _run_demo_pipeline(img_paths, expected_scene):
    """执行示例检测流程（不涉及用户历史数据，纯两张图比对）

    expected_scene 直接指定场景类型，跳过 classify_scene 节省时间
    """
    import cv2
    import numpy as np

    path1, path2 = img_paths[0], img_paths[1]

    # 跳过场景分类，直接使用 expected_scene
    scene_type = expected_scene
    ocr_result = None
    full_text = ''
    print(f"[demo] 跳过场景分类，直接使用指定场景: {scene_type}")

    is_same = False
    similarity = 0.0
    log_entry = {
        'query_path': path1,
        'history_path': path2,
        'scene_type': scene_type,
        'is_same_scene': False,
        'query_image': os.path.basename(path1),
        'history_image': os.path.basename(path2),
        'query_text': full_text,
    }
    visualization_data = []

    if scene_type == 'text':
        # 文本场景：OCR + BGE 比对
        from module.text_matcher import compute_text_similarity
        try:
            # 直接 OCR 提取两张图文本，跳过 classify_scene
            from module.text_classifier import _ocr_image
            full_text = _ocr_image(path1, ocr_engine) or ''
            text2 = _ocr_image(path2, ocr_engine) or ''
            log_entry['query_text'] = full_text
            log_entry['history_text'] = text2

            # BGE 编码比对
            text_sim = compute_text_similarity(full_text, text2, bge_tokenizer, bge_model)
            similarity = text_sim
            is_same = text_sim >= 0.85
            log_entry['text_similarity'] = text_sim
            log_entry['is_same_scene'] = is_same
        except Exception as e:
            print(f"[demo] 文本场景比对失败: {e}")

    else:
        # 复杂场景：DINOv2 比对
        try:
            if dinov2_extractor is not None:
                dino_sim = dinov2_extractor.compute_similarity(path1, path2)
                similarity = dino_sim if dino_sim else 0.0
                is_same = similarity >= 0.5
                log_entry['dinov2_similarity'] = round(similarity, 4)
                log_entry['is_same_scene'] = is_same

                # 计算 patch 匹配（用于可视化）
                if is_same:
                    patch_info = dinov2_extractor.compute_patch_matches(path1, path2, top_k=50)
                    log_entry['patch_match_info'] = patch_info
        except Exception as e:
            print(f"[demo] 复杂场景比对失败: {e}")

    # 构建可视化（两张图都还在，直接构建）
    _image_cache = {}
    for p in [path1, path2]:
        try:
            data = np.fromfile(p, dtype=np.uint8)
            img = cv2.imdecode(data, cv2.IMREAD_COLOR)
            if img is not None:
                _image_cache[p] = img
        except Exception:
            pass

    try:
        if scene_type == 'text':
            vis = _build_text_vis_from_cache(log_entry, _image_cache)
            if vis:
                visualization_data.append(vis)
        else:
            vis = _build_complex_vis_from_cache(log_entry, _image_cache)
            if vis:
                visualization_data.append(vis)
    except Exception as e:
        print(f"[demo] 可视化生成失败: {e}")

    # 将原图也转为 base64 返回（前端无需再请求 /api/image）
    def _img_to_b64(path):
        try:
            data = np.fromfile(path, dtype=np.uint8)
            img = cv2.imdecode(data, cv2.IMREAD_COLOR)
            if img is None:
                return None
            _, buf = cv2.imencode('.jpg', img, [cv2.IMWRITE_JPEG_QUALITY, 95])
            return f'data:image/jpeg;base64,{base64.b64encode(buf).decode("utf-8")}'
        except Exception:
            return None

    return {
        "scene_type": scene_type,
        "is_same_scene": is_same,
        "similarity": round(similarity, 4),
        "query_image": os.path.basename(path1),
        "history_image": os.path.basename(path2),
        "query_image_b64": _img_to_b64(path1),
        "history_image_b64": _img_to_b64(path2),
        "visualizations": visualization_data,
        "evaluation_logs": [{
            "query_image": os.path.basename(path1),
            "history_image": os.path.basename(path2),
            "scene_type": scene_type,
            "is_same_scene": is_same,
            "similarity": round(similarity, 4),
        }],
    }


if __name__ == "__main__":
    import uvicorn
    port = int(os.environ.get("PORT", 7860))
    uvicorn.run("app:app", host="0.0.0.0", port=port, reload=False)