"""
下载 BGE embedding 和 reranker 模型到本地 models/ 目录

首次配置时运行：
    python scripts/download_models.py
"""

import os
import sys

PROJECT_ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
MODELS_DIR = os.path.join(PROJECT_ROOT, "models")
EMBEDDING_MODEL_PATH = os.path.join(MODELS_DIR, "bge-base-zh-v1.5")
RERANK_MODEL_PATH = os.path.join(MODELS_DIR, "bge-reranker-base")

os.makedirs(MODELS_DIR, exist_ok=True)


def download_embedding_model():
    """下载 BGE 中文向量模型"""
    if os.path.exists(os.path.join(EMBEDDING_MODEL_PATH, "config.json")):
        print(f"[跳过] Embedding 模型已存在: {EMBEDDING_MODEL_PATH}")
        return True

    print(f"[下载] BGE 中文向量模型 (bge-base-zh-v1.5)...")
    print(f"  目标路径: {EMBEDDING_MODEL_PATH}")
    try:
        from sentence_transformers import SentenceTransformer
        model = SentenceTransformer("BAAI/bge-base-zh-v1.5")
        model.save(EMBEDDING_MODEL_PATH)
        print(f"[完成] Embedding 模型下载成功!")
        return True
    except Exception as e:
        print(f"[失败] Embedding 模型下载失败: {e}")
        return False


def download_reranker_model():
    """下载 BGE reranker 模型"""
    if os.path.exists(os.path.join(RERANK_MODEL_PATH, "config.json")):
        print(f"[跳过] Reranker 模型已存在: {RERANK_MODEL_PATH}")
        return True

    print(f"[下载] BGE Reranker 模型 (bge-reranker-base)...")
    print(f"  目标路径: {RERANK_MODEL_PATH}")
    try:
        from transformers import AutoModelForSequenceClassification, AutoTokenizer
        os.makedirs(RERANK_MODEL_PATH, exist_ok=True)
        model = AutoModelForSequenceClassification.from_pretrained("BAAI/bge-reranker-base")
        tokenizer = AutoTokenizer.from_pretrained("BAAI/bge-reranker-base")
        model.save_pretrained(RERANK_MODEL_PATH)
        tokenizer.save_pretrained(RERANK_MODEL_PATH)
        print(f"[完成] Reranker 模型下载成功!")
        return True
    except Exception as e:
        print(f"[失败] Reranker 模型下载失败: {e}")
        print(f"  (项目启动时会自动尝试下载，或可稍后手动重试)")
        return False


if __name__ == "__main__":
    print("=" * 60)
    print("RAG 项目模型下载脚本")
    print("=" * 60)

    ok1 = download_embedding_model()
    print()
    ok2 = download_reranker_model()

    print()
    print("=" * 60)
    if ok1:
        print("Embedding 模型: OK (项目可以启动)")
    else:
        print("Embedding 模型: FAILED (项目无法启动，请检查网络)")

    if ok2:
        print("Reranker 模型:  OK")
    else:
        print("Reranker 模型:  FAILED (不影响启动，首次使用 rerank 时自动下载)")

    print("=" * 60)