# -*- coding: utf-8 -*-
"""图片名称检索测试台 —— 后端(FastAPI)。

三路检索,全部基于「图片名称」:
  1) BM25     —— 关键词匹配(jieba 分词)
  2) 向量     —— 本地 BAAI/bge-m3 余弦相似度
  3) 融合     —— RRF 合并上面两路排名

【写死目录 + 局域网模式】图片目录写死在 IMAGE_DIR,服务器启动即扫描并自动建索引、
自己通过 /images 把图片发出去。绑 0.0.0.0,mentor 用局域网 IP 打开即可直接检索,
无需选文件夹、无需建索引。

跑法:python server.py  → 控制台会打印可发给 mentor 的局域网地址。
"""
from __future__ import annotations

import os
import re
import socket
import sys
import threading

import numpy as np
from fastapi import Body, FastAPI, Query
from fastapi.responses import JSONResponse
from fastapi.staticfiles import StaticFiles

from bm25 import BM25Index
from embed_local import LocalEmbedder
from fuse import rrf

for _s in (sys.stdout, sys.stderr):
    try:
        _s.reconfigure(encoding="utf-8", errors="replace")
    except Exception:
        pass

HERE = os.path.dirname(os.path.abspath(__file__))
# 图片目录:默认 = 本仓库的 crawler/images(爬虫默认就把图下到那里)。
# 想指向别处已爬好的图,启动前设环境变量 IMAGE_DIR=<绝对路径> 覆盖即可,无需改代码。
IMAGE_DIR = os.environ.get("IMAGE_DIR") or os.path.abspath(
    os.path.join(HERE, os.pardir, "crawler", "images"))
IMG_EXT = (".jpg", ".jpeg", ".png", ".webp", ".gif", ".bmp")
PORT = 8799

app = FastAPI(title="图片名称检索测试台")


def scan_images(img_dir: str) -> dict[str, list[str]]:
    """扫描图片目录 → {菜名: [文件名,...]}。去掉尾部 _NN 把同名多图归并为一条。"""
    name2files: dict[str, list[str]] = {}
    if not os.path.isdir(img_dir):
        return name2files
    for fn in os.listdir(img_dir):
        if not fn.lower().endswith(IMG_EXT):
            continue
        stem = os.path.splitext(fn)[0]
        name = re.sub(r"_\d+$", "", stem) or stem
        name2files.setdefault(name, []).append(fn)
    for files in name2files.values():
        files.sort()
    return name2files


class State:
    """全局索引状态(单实例工具,够用)。"""

    def __init__(self):
        self.lock = threading.Lock()
        self.names: list[str] = []
        self.name2files: dict[str, list[str]] = {}
        self.bm25: BM25Index | None = None
        self.embedder: LocalEmbedder | None = None
        self.model = "BAAI/bge-m3"
        self.vec_names: list[str] = []
        self.vec_matrix: np.ndarray | None = None
        self.state = "idle"      # idle | building | ready | error
        self.phase = ""          # bm25 | loadmodel | embedding | done
        self.emb_done = 0
        self.emb_total = 0
        self.failed = 0
        self.message = ""
        self.device = ""


ST = State()


def build_index(names: list[str], model: str) -> None:
    try:
        with ST.lock:
            ST.state, ST.phase, ST.message = "building", "bm25", ""
            ST.names, ST.model = names, model
            ST.bm25 = None
            ST.vec_matrix, ST.vec_names = None, []
            ST.emb_done = ST.emb_total = ST.failed = 0

        # 1) BM25(秒级)
        bm = BM25Index(names)
        with ST.lock:
            ST.bm25, ST.phase = bm, "loadmodel"

        # 2) 本地 BGE-M3:先加载模型(首次会下载约 2.3G),再编码(有 CUDA 走 GPU)
        emb = LocalEmbedder(model=model, cache_dir=os.path.join(HERE, "cache"))
        dev = emb.load_model()
        with ST.lock:
            ST.device, ST.phase = dev, "embedding"

        def prog(done: int, total: int) -> None:
            with ST.lock:
                ST.emb_done, ST.emb_total = done, total

        emb.embed_corpus(names, progress=prog)
        vnames, mat = emb.matrix(names)
        with ST.lock:
            ST.embedder, ST.vec_names, ST.vec_matrix = emb, vnames, mat
            ST.failed = emb.last_failed
            ST.message = "" if vnames else ("向量未生成:" + emb.last_error if emb.last_error else "")
            ST.state, ST.phase = "ready", "done"
    except Exception as e:  # noqa: BLE001
        with ST.lock:
            ST.state, ST.message = "error", f"{type(e).__name__}: {e}"


@app.on_event("startup")
def _startup_autobuild() -> None:
    """启动即扫描写死目录并后台建索引,mentor 打开就能搜。"""
    n2f = scan_images(IMAGE_DIR)
    with ST.lock:
        ST.name2files = n2f
    names = sorted(n2f.keys())
    if not names:
        with ST.lock:
            ST.state = "error"
            ST.message = f"图片目录为空或不存在:{IMAGE_DIR}"
        return
    threading.Thread(target=build_index, args=(names, "BAAI/bge-m3"), daemon=True).start()


@app.post("/api/index")
def api_index(payload: dict = Body(...)):
    """(保留)手动指定名称重建索引;写死目录模式下一般用不到。"""
    names = [str(n).strip() for n in payload.get("names", []) if str(n).strip()]
    names = list(dict.fromkeys(names))
    model = (payload.get("model") or "BAAI/bge-m3").strip()
    if not names:
        return JSONResponse({"ok": False, "error": "没有菜名"}, status_code=400)
    if ST.state == "building":
        return JSONResponse({"ok": False, "error": "正在建索引,请稍候"}, status_code=409)
    threading.Thread(target=build_index, args=(names, model), daemon=True).start()
    return {"ok": True, "total": len(names), "model": model}


@app.get("/api/status")
def api_status():
    with ST.lock:
        return {
            "state": ST.state,
            "phase": ST.phase,
            "bm25_ready": ST.bm25 is not None,
            "vec_ready": ST.vec_matrix is not None and len(ST.vec_names) > 0,
            "total": len(ST.names),
            "n_images": sum(len(v) for v in ST.name2files.values()),
            "emb_done": ST.emb_done,
            "emb_total": ST.emb_total,
            "vec_count": len(ST.vec_names),
            "failed": ST.failed,
            "model": ST.model,
            "device": ST.device,
            "message": ST.message,
        }


def _bm25_search(q: str, topk: int) -> list[tuple[str, float]]:
    return ST.bm25.search(q, topk=topk) if ST.bm25 else []


def _vec_search(q: str, topk: int) -> list[tuple[str, float]]:
    mat, names, emb = ST.vec_matrix, ST.vec_names, ST.embedder
    if mat is None or len(names) == 0 or emb is None:
        return []
    qv = emb.embed_query(q)
    sims = mat @ qv
    order = np.argsort(sims)[::-1][:topk]
    return [(names[i], float(sims[i])) for i in order]


@app.get("/api/search")
def api_search(q: str = Query(...), topk: int = 30, rrf_k: int = 60,
               w_bm25: float = 1.0, w_vec: float = 1.0):
    q = q.strip()
    if not q:
        return {"q": q, "bm25": [], "vector": [], "fusion": []}

    bm = _bm25_search(q, topk)
    ve = _vec_search(q, topk)
    bm_rank = [n for n, _ in bm]
    ve_rank = [n for n, _ in ve]
    fused = rrf({"bm25": bm_rank, "vector": ve_rank},
                {"bm25": w_bm25, "vector": w_vec}, k=rrf_k)[:topk]

    bm_s, ve_s = dict(bm), dict(ve)
    bm_pos = {n: i + 1 for i, n in enumerate(bm_rank)}
    ve_pos = {n: i + 1 for i, n in enumerate(ve_rank)}

    def card(name: str, score: float) -> dict:
        return {
            "name": name,
            "score": round(score, 4),
            "files": ST.name2files.get(name, []),     # 服务器 serve 的图片文件名
            "bm25": round(bm_s[name], 3) if name in bm_s else None,
            "bm25_rank": bm_pos.get(name),
            "vec": round(ve_s[name], 4) if name in ve_s else None,
            "vec_rank": ve_pos.get(name),
        }

    return {
        "q": q,
        "bm25": [card(n, s) for n, s in bm],
        "vector": [card(n, s) for n, s in ve],
        "fusion": [card(n, s) for n, s in fused],
    }


# 图片由服务器 serve(局域网里 mentor 的浏览器据此显示);必须在 "/" 之前挂载。
if os.path.isdir(IMAGE_DIR):
    app.mount("/images", StaticFiles(directory=IMAGE_DIR), name="images")
# 静态前端(放在所有 /api 与 /images 之后挂载)
app.mount("/", StaticFiles(directory=os.path.join(HERE, "static"), html=True), name="static")


def _lan_ips() -> list[str]:
    ips: list[str] = []
    try:
        s = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
        s.connect(("8.8.8.8", 80))
        ips.append(s.getsockname()[0])
        s.close()
    except Exception:
        pass
    try:
        for info in socket.getaddrinfo(socket.gethostname(), None, socket.AF_INET):
            ip = info[4][0]
            if ip not in ips:
                ips.append(ip)
    except Exception:
        pass
    return [i for i in ips if not i.startswith("127.")]


if __name__ == "__main__":
    import uvicorn

    print("=" * 56, flush=True)
    print(f"图片目录:{IMAGE_DIR}", flush=True)
    print("发给 mentor 的局域网地址:", flush=True)
    for ip in _lan_ips() or ["<查不到,用 ipconfig 看本机 IPv4>"]:
        print(f"   http://{ip}:{PORT}", flush=True)
    print(f"本机自测:http://127.0.0.1:{PORT}", flush=True)
    print("(若 mentor 连不上:Windows 防火墙放行该端口入站)", flush=True)
    print("=" * 56, flush=True)
    uvicorn.run(app, host="0.0.0.0", port=PORT, log_level="info")