# -*- coding: utf-8 -*-
"""
vector_builder.py — 向量预计算脚本

读取 搜索索引.json，调用 OpenAI 兼容 /v1/embeddings API 生成向量，
写入 搜索向量.json，供 SearchEngine 运行时加载。

用法：
  # 默认路径（自动查找索引、输出同目录）
  python -m utils.vector_builder

  # 指定路径
  python -m utils.vector_builder --index-path /path/to/搜索索引.json --output /path/to/搜索向量.json

  # 指定批大小与 API 参数
  python -m utils.vector_builder --batch-size 20 --model text-embedding-3-small
"""

import argparse
import json
import os
import sys
import time
from pathlib import Path

import requests

# ============================================================
# 路径定位（与 search_engine.py 保持一致）
# ============================================================

SCRIPT_DIR = Path(__file__).resolve().parent.parent  # .../web
DATA_DIR = SCRIPT_DIR.parent  # 部署时即 /opt/troubleshoot


def _find_index_path(explicit_path=None):
    """查找搜索索引文件"""
    if explicit_path:
        p = Path(explicit_path)
        if p.exists():
            return p
        print(f"[vector_builder] 指定的索引文件不存在：{p}")

    # 复用 search_engine.py 的路径查找逻辑
    candidates = [
        DATA_DIR / "搜索索引.json",
        SCRIPT_DIR.parent.parent.parent.parent / "Docs" / "PRD" / "问题知识库" / "搜索索引.json",
    ]
    for path in candidates:
        if path.exists():
            return path

    print("[vector_builder] 错误：未找到搜索索引文件")
    sys.exit(1)


def _get_api_config():
    """读取 API 配置（从 container.get_config 或环境变量）"""
    try:
        import container
        config = container.get_config()
        return {
            "api_base": config.get("claude_api_base", os.environ.get("CLAUDE_API_BASE", "")),
            "api_key": config.get("claude_api_key", os.environ.get("CLAUDE_API_KEY", "")),
            "model": config.get("embedding_model", "text-embedding-3-small"),
            "timeout": config.get("embedding_api_timeout", 10),
        }
    except Exception:
        return {
            "api_base": os.environ.get("CLAUDE_API_BASE", ""),
            "api_key": os.environ.get("CLAUDE_API_KEY", ""),
            "model": "text-embedding-3-small",
            "timeout": 10,
        }


def _call_embedding_api(texts, api_base, api_key, model, timeout):
    """调用 /v1/embeddings API，返回向量列表"""
    url = f"{api_base.rstrip('/')}/v1/embeddings"
    headers = {
        "Content-Type": "application/json",
        "Authorization": f"Bearer {api_key}",
    }
    payload = {
        "model": model,
        "input": texts,
    }

    resp = requests.post(url, json=payload, headers=headers, timeout=timeout)
    resp.raise_for_status()
    data = resp.json()

    # 按 index 排序保证顺序
    embeddings = sorted(data["data"], key=lambda x: x["index"])
    return [item["embedding"] for item in embeddings]


def build_vectors(index_path=None, output_path=None, batch_size=20, model_override=None):
    """
    读取索引，批量生成向量，写入文件。

    参数：
      index_path:  搜索索引文件路径（None 则自动查找）
      output_path: 输出向量文件路径（None 则与索引同目录）
      batch_size:  每批调 API 的记录数
      model_override: 覆盖配置中的模型名

    返回：
      (total, success, failed) 统计
    """
    index_file = _find_index_path(index_path)
    print(f"[vector_builder] 加载索引：{index_file}")

    with open(index_file, 'r', encoding='utf-8') as f:
        index_data = json.load(f)

    records = index_data.get('records', [])
    if not records:
        print("[vector_builder] 错误：索引中无记录")
        return 0, 0, 0

    # 构建待编码文本：title + full_text
    texts = []
    for rec in records:
        title = rec.get('title', '')
        full_text = rec.get('full_text', '')
        texts.append(f"{title} {full_text}".strip())

    # API 配置
    api_config = _get_api_config()
    model = model_override or api_config["model"]
    api_base = api_config["api_base"]
    api_key = api_config["api_key"]
    timeout = api_config["timeout"]

    if not api_base or not api_key:
        print("[vector_builder] 错误：API 地址或密钥未配置（CLAUDE_API_BASE / CLAUDE_API_KEY）")
        return len(records), 0, len(records)

    print(f"[vector_builder] 开始生成向量：{len(records)} 条记录，模型={model}，批大小={batch_size}")

    # 分批调用
    vectors = {}
    total = len(texts)
    success = 0
    failed = 0
    dimensions = None

    for batch_start in range(0, total, batch_size):
        batch_end = min(batch_start + batch_size, total)
        batch_texts = texts[batch_start:batch_end]

        try:
            batch_vectors = _call_embedding_api(batch_texts, api_base, api_key, model, timeout)

            if dimensions is None and batch_vectors:
                dimensions = len(batch_vectors[0])

            for i, vec in enumerate(batch_vectors):
                record_idx = batch_start + i
                record_key = f"rec_{record_idx:03d}"
                vectors[record_key] = vec
                success += 1

            print(f"  批次 {batch_start//batch_size + 1}：{batch_start+1}-{batch_end}/{total} 完成")

        except Exception as e:
            failed += len(batch_texts)
            print(f"  批次 {batch_start//batch_size + 1}：失败 — {e}")

        # 批次间间隔，避免限流
        if batch_end < total:
            time.sleep(0.5)

    if not vectors:
        print("[vector_builder] 错误：无向量生成成功，不写入文件")
        return total, 0, total

    # 确定输出路径
    if output_path:
        out_file = Path(output_path)
    else:
        out_file = index_file.parent / "搜索向量.json"

    # 写入向量文件
    output_data = {
        "model": model,
        "dimensions": dimensions or 0,
        "generated_at": time.strftime("%Y-%m-%dT%H:%M:%S"),
        "record_count": len(vectors),
        "vectors": vectors,
    }

    out_file.parent.mkdir(parents=True, exist_ok=True)
    with open(out_file, 'w', encoding='utf-8') as f:
        json.dump(output_data, f, ensure_ascii=False)

    print(f"[vector_builder] 完成：成功 {success}，失败 {failed}，总计 {total}")
    print(f"[vector_builder] 向量文件已写入：{out_file}")

    return total, success, failed


if __name__ == '__main__':
    parser = argparse.ArgumentParser(description="生成搜索向量索引")
    parser.add_argument('--index-path', help='搜索索引文件路径')
    parser.add_argument('--output', help='输出向量文件路径')
    parser.add_argument('--batch-size', type=int, default=20, help='每批调 API 的记录数（默认 20）')
    parser.add_argument('--model', help='覆盖配置中的 embedding 模型名')
    args = parser.parse_args()

    build_vectors(
        index_path=args.index_path,
        output_path=args.output,
        batch_size=args.batch_size,
        model_override=args.model,
    )
