how/lib/data_loader.py

"""
通用数据加载模块

提供项目中常用的数据加载函数
"""
import os
import sys
from typing import List
from lib.utils import read_json


def load_persona_data(persona_dir: str) -> dict:
    """加载人设数据

    Args:
        persona_dir: 人设目录路径

    Returns:
        人设数据字典

    Raises:
        SystemExit: 文件不存在时退出
    """
    persona_data_path = os.path.join(persona_dir, "人设.json")
    try:
        return read_json(persona_data_path)
    except FileNotFoundError:
        print(f"❌ 找不到人设数据文件: {persona_data_path}")
        print(f"请检查路径是否正确: {persona_dir}")
        sys.exit(1)


def load_inspiration_list(persona_dir: str) -> List[str]:
    """加载灵感点列表

    Args:
        persona_dir: 人设目录路径

    Returns:
        灵感点文本列表

    Raises:
        SystemExit: 文件不存在或格式错误时退出
    """
    inspiration_list_path = os.path.join(persona_dir, "灵感点.json")
    try:
        inspiration_data = read_json(inspiration_list_path)
        if not isinstance(inspiration_data, list) or len(inspiration_data) == 0:
            print(f"❌ 灵感文件格式错误或为空: {inspiration_list_path}")
            sys.exit(1)
        return [item["灵感点"] for item in inspiration_data]
    except FileNotFoundError:
        print(f"❌ 找不到灵感文件: {inspiration_list_path}")
        print("请先运行 extract_inspirations.py 生成灵感点文件")
        sys.exit(1)


def load_inspiration_data(persona_dir: str) -> List[dict]:
    """加载完整的灵感点数据（包含 meta 信息）

    Args:
        persona_dir: 人设目录路径

    Returns:
        灵感点数据列表，每项包含 {"灵感点": str, "meta": dict}

    Raises:
        SystemExit: 文件不存在或格式错误时退出
    """
    inspiration_list_path = os.path.join(persona_dir, "灵感点.json")
    try:
        inspiration_data = read_json(inspiration_list_path)
        if not isinstance(inspiration_data, list) or len(inspiration_data) == 0:
            print(f"❌ 灵感文件格式错误或为空: {inspiration_list_path}")
            sys.exit(1)
        return inspiration_data
    except FileNotFoundError:
        print(f"❌ 找不到灵感文件: {inspiration_list_path}")
        print("请先运行 extract_inspirations.py 生成灵感点文件")
        sys.exit(1)


def select_inspiration(inspiration_arg: str, inspiration_list: List[str]) -> str:
    """根据参数选择灵感

    Args:
        inspiration_arg: 灵感参数（数字索引或灵感名称）
        inspiration_list: 灵感点文本列表

    Returns:
        选中的灵感点文本

    Raises:
        SystemExit: 选择失败时退出
    """
    try:
        # 尝试作为索引解析
        inspiration_index = int(inspiration_arg)
        if 0 <= inspiration_index < len(inspiration_list):
            inspiration = inspiration_list[inspiration_index]
            print(f"使用灵感[{inspiration_index}]: {inspiration}")
            return inspiration
        else:
            print(f"❌ 灵感索引超出范围: {inspiration_index} (有效范围: 0-{len(inspiration_list)-1})")
    except ValueError:
        # 不是数字，当作灵感名称
        if inspiration_arg in inspiration_list:
            print(f"使用灵感: {inspiration_arg}")
            return inspiration_arg
        else:
            print(f"❌ 找不到灵感: {inspiration_arg}")

    # 显示可用灵感列表后退出
    print(f"可用灵感列表:")
    for i, insp in enumerate(inspiration_list[:10]):
        print(f"  {i}: {insp}")
    if len(inspiration_list) > 10:
        print(f"  ... 还有 {len(inspiration_list) - 10} 个")
    sys.exit(1)


def load_step1_result(persona_dir: str, inspiration: str, model_name: str, scope: str = "all") -> dict:
    """加载 step1 匹配结果

    Args:
        persona_dir: 人设目录路径
        inspiration: 灵感点名称
        model_name: 模型名称（如 "google/gemini-2.5-pro"）
        scope: 范围标识（"all" 或 "top10" 等）

    Returns:
        step1 结果字典

    Raises:
        SystemExit: 文件不存在时退出
    """
    # 提取模型简称
    model_name_short = model_name.replace("google/", "").replace("/", "_")

    # 构建文件路径
    step1_file = os.path.join(
        persona_dir,
        "how",
        "灵感点",
        inspiration,
        f"{scope}_step1_灵感人设匹配_{model_name_short}.json"
    )

    try:
        return read_json(step1_file)
    except FileNotFoundError:
        print(f"❌ 找不到 step1 结果文件: {step1_file}")
        print(f"请先运行 step1_inspiration_match.py 生成结果")
        sys.exit(1)