comfyui_o1key/utils/config.py

"""
配置管理模块
处理环境变量和 API 密钥管理
"""

import os
import threading
from typing import Dict, Optional


# 获取插件根目录
PLUGIN_ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
CONFIG_FILE = os.path.join(PLUGIN_ROOT, ".config")
_CONFIG_LOCK = threading.RLock()


# ============ API 基础配置 ============
# 所有 API 客户端的统一基础 URL
# 可通过环境变量 O1KEY_API_BASE_URL 覆盖
DEFAULT_API_BASE_URL = "https://api.o1key.cn"

# 异步 API 基础 URL（用于异步提交+轮询模式）
# 可通过环境变量 O1KEY_ASYNC_API_BASE_URL 覆盖
DEFAULT_ASYNC_API_BASE_URL = "https://api.o1key.cn"

# ============ 网络线路配置 ============
NETWORK_ROUTES = {
    "全球加速": "https://api.o1key.cn",
    "CF加速": "https://cf-api.o1key.com",
    "美国直连": "https://api.o1key.com",
}
NETWORK_ROUTE_OPTIONS = ["全球加速", "CF加速", "美国直连"]
DEFAULT_NETWORK_ROUTE = "全球加速"
NETWORK_ROUTE_CONFIG_KEY = "O1KEY_NETWORK_ROUTE"


def load_config(config_path: Optional[str] = None) -> Dict[str, str]:
    """
    从配置文件加载所有配置项
    
    Args:
        config_path: 配置文件路径，默认为插件目录下的 .config
    
    Returns:
        配置字典 {key: value}
    
    Example:
        >>> config = load_config()
        >>> api_key = config.get('O1KEY_API_KEY')
    """
    if config_path is None:
        config_path = CONFIG_FILE
    
    config = {}
    
    if not os.path.exists(config_path):
        return config
    
    try:
        with open(config_path, 'r', encoding='utf-8') as f:
            for line in f:
                line = line.strip()
                
                # 跳过空行和注释
                if not line or line.startswith('#'):
                    continue
                
                # 解析 KEY=VALUE 格式
                if '=' in line:
                    key, value = line.split('=', 1)
                    key = key.strip()
                    value = value.strip().strip('"').strip("'")
                    
                    if key and value:
                        config[key] = value
    
    except Exception as e:
        print(f"⚠️ 读取配置文件失败: {e}")
    
    return config


def save_config(config: Dict[str, str], config_path: Optional[str] = None) -> None:
    """原子写入配置，避免界面快速切换时产生半写入文件。"""
    if config_path is None:
        config_path = CONFIG_FILE

    temp_path = f"{config_path}.tmp"
    with _CONFIG_LOCK:
        with open(temp_path, "w", encoding="utf-8", newline="\n") as config_file:
            for key, value in config.items():
                config_file.write(f"{key}={value}\n")
        os.replace(temp_path, config_path)


def update_config(
    updates: Optional[Dict[str, str]] = None,
    remove: Optional[list[str]] = None,
) -> Dict[str, str]:
    """以一次原子写入更新配置，并返回更新后的配置。"""
    with _CONFIG_LOCK:
        config = load_config()
        for key, value in (updates or {}).items():
            config[key] = value
        for key in remove or []:
            config.pop(key, None)
        save_config(config)
        return config


def get_runtime_config_signature() -> tuple[str, str, str, str]:
    """用于复用客户端；仅当运行时 API 配置变化时才重建客户端。"""
    config = load_config()
    return (
        config.get("O1KEY_API_KEY", ""),
        config.get(NETWORK_ROUTE_CONFIG_KEY, DEFAULT_NETWORK_ROUTE),
        config.get("O1KEY_API_BASE_URL", ""),
        config.get("O1KEY_ASYNC_API_BASE_URL", ""),
    )


def get_api_key(key_name: str = "O1KEY_API_KEY") -> Optional[str]:
    """
    获取 API 密钥
    从 .config 文件读取

    Args:
        key_name: 密钥名称，默认为 O1KEY_API_KEY

    Returns:
        API 密钥字符串，如果未找到则返回 None
    """
    config = load_config()
    return config.get(key_name)


def get_api_key_or_raise(key_name: str = "O1KEY_API_KEY") -> str:
    """
    获取 API 密钥，如果未找到则抛出异常
    
    Args:
        key_name: 密钥名称
    
    Returns:
        API 密钥字符串
    
    Raises:
        ValueError: 如果未找到 API 密钥
    """
    api_key = get_api_key(key_name)
    
    if not api_key:
        raise ValueError("未授权！")
    
    return api_key


def get_network_route() -> str:
    """Return the globally configured network route."""
    route = load_config().get(NETWORK_ROUTE_CONFIG_KEY, DEFAULT_NETWORK_ROUTE)
    if isinstance(route, (list, tuple)):
        route = route[0] if route else DEFAULT_NETWORK_ROUTE
    return route if route in NETWORK_ROUTES else DEFAULT_NETWORK_ROUTE


def get_api_base_url() -> str:
    """
    获取 API 基础 URL
    从 .config 文件读取，如果未配置则使用默认值

    Returns:
        API 基础 URL 字符串
    """
    config = load_config()
    route = config.get(NETWORK_ROUTE_CONFIG_KEY)
    if route in NETWORK_ROUTES:
        return NETWORK_ROUTES[route].rstrip('/')
    base_url = config.get("O1KEY_API_BASE_URL")

    if base_url:
        return base_url.rstrip('/')

    return DEFAULT_API_BASE_URL


def get_async_api_base_url() -> str:
    """
    获取异步 API 基础 URL
    从 .config 文件读取，如果未配置则使用默认值

    Returns:
        异步 API 基础 URL 字符串
    """
    config = load_config()
    route = config.get(NETWORK_ROUTE_CONFIG_KEY)
    if route in NETWORK_ROUTES:
        return NETWORK_ROUTES[route].rstrip('/')
    base_url = config.get("O1KEY_ASYNC_API_BASE_URL")

    if base_url:
        return base_url.rstrip('/')

    return DEFAULT_ASYNC_API_BASE_URL


def get_base_url_by_route(route: Optional[str] = None) -> str:
    """Resolve an explicit route or fall back to the global route setting."""
    if isinstance(route, (list, tuple)):
        route = route[0] if route else None
    if route is None:
        route = get_network_route()
    return NETWORK_ROUTES.get(route, get_api_base_url())