"""常驻检索进程。标准输出只用于 JSON 协议，初始化和检索诊断写标准错误。"""
import sys
import json
import time
from contextlib import redirect_stdout
from pathlib import Path

# 从任意 cwd 启动都能找到 app 包（uvicorn 工作目录不保证是 server/）
_SERVER_ROOT = Path(__file__).resolve().parents[2]
if str(_SERVER_ROOT) not in sys.path:
    sys.path.insert(0, str(_SERVER_ROOT))


def initialize():
    """每个 worker 仅初始化一次，避免每次查询重复执行迁移和历史表扫描。"""
    try:
        import fcntl
    except ImportError:
        fcntl = None
    from app import db as db_mod
    from app.config import DATA_DIR
    # 多个首次查询同时启动进程时，不并发执行相同的数据库迁移。
    with (DATA_DIR / 'rag-worker-init.lock').open('a') as lock:
        if fcntl:
            fcntl.flock(lock, fcntl.LOCK_EX)
        db_mod.init_db()


def handle(payload: dict) -> list:
    from sqlmodel import Session
    from app import db as db_mod
    from app.rag import pipeline as rag_pipeline

    with Session(db_mod.engine) as session:
        results = rag_pipeline.search_kbs(
            session,
            payload["kb_ids"],
            payload["query"],
            payload.get("k", 6),
            excluded_chunk_ids={int(k): set(v) for k, v in (payload.get("excluded") or {}).items()},
            campaign_id=payload.get("campaign_id", 0),
            trace_id=payload.get("trace_id", ""),
            user_id=payload.get("user_id", 0),
            organization_id=payload.get("organization_id", 1),
            file_whitelists={int(k): set(v) for k, v in (payload.get("file_whitelists") or {}).items()},
        )
    # 透传库名、召回来源、重排分数和各阶段耗时，不裁掉检索诊断信息。
    return results


def handle_warmup() -> dict:
    """Warm the resident worker without issuing a billable retrieval."""
    from app.rag import embed, store, rerank
    started = time.monotonic()
    embed.prewarm()
    # Optional local Cross-Encoder is loaded in the resident worker so the
    # first real retrieval does not pay model initialization latency.
    rerank.prewarm()
    # Construct the persistent vector client once. Collections remain opened
    # lazily per KB, while the client and embedding runtime stay resident.
    store.get_client()
    return {"warm_state": "warm", "warmup_ms": round((time.monotonic() - started) * 1000, 2)}


def main():
    protocol_out = sys.stdout
    try:
        with redirect_stdout(sys.stderr):
            initialize()
    except Exception as exc:
        from app.system_log import redact
        protocol_out.write(json.dumps({'ok': False, 'error': '检索进程初始化失败：' + redact(exc, 400)}, ensure_ascii=False) + '\n')
        protocol_out.flush()
        return
    for line in sys.stdin:
        line = line.strip()
        if not line:
            continue
        try:
            started = time.monotonic()
            with redirect_stdout(sys.stderr):
                payload = json.loads(line)
                if payload.get("op") == "warmup":
                    warm = handle_warmup()
                    protocol_out.write(json.dumps({"ok": True, **warm}, ensure_ascii=False) + "\n")
                    protocol_out.flush()
                    continue
                out = handle(payload)
            from app.rag import pipeline as _pipeline
            metric = _pipeline.last_retrieval_metric()
            envelope = {"ok": True, "results": out,
                        "worker_timings": {"search_ms": round((time.monotonic() - started) * 1000, 2)},
                        "retrieval_metric": metric}
            serialization_started = time.perf_counter()
            json.dumps(envelope, ensure_ascii=False)
            metric["serialization_ms"] = round((time.perf_counter() - serialization_started) * 1000, 2)
            envelope["retrieval_metric"] = metric
            protocol_out.write(json.dumps(envelope, ensure_ascii=False) + "\n")
        except Exception as exc:
            from app.system_log import redact
            metric = {}
            try:
                from app.rag import pipeline as _pipeline
                metric = _pipeline.last_retrieval_metric()
            except Exception:
                pass
            if not metric:
                metric = {"error": "retrieval_failed"}
            protocol_out.write(json.dumps({"ok": False, "error": redact(exc, 400),
                "retrieval_metric": metric}, ensure_ascii=False) + "\n")
        protocol_out.flush()


if __name__ == "__main__":
    main()