YG_FT/backend/app/core/config.py

from __future__ import annotations

from dataclasses import dataclass, field
from functools import lru_cache
import os
from pathlib import Path
from typing import Any
from urllib.parse import quote, urlsplit, urlunsplit

import yaml


BACKEND_ROOT = Path(__file__).resolve().parents[2]
DEFAULT_CONFIG_PATH = BACKEND_ROOT / "config.yaml"
DEFAULT_DATABASE_BASE_URL = "postgresql+psycopg://localhost:5432/yg_ft"
DEFAULT_DATABASE_USERNAME = "yg_ft"
DEFAULT_DATABASE_PASSWORD = "change_me"


class ConfigurationError(ValueError):
    """Raised when the backend configuration cannot be parsed or validated."""


def _resolve_config_path(config_path: str | Path | None = None) -> tuple[Path, bool]:
    configured_path = config_path or os.getenv("BACKEND_CONFIG_FILE")
    is_explicit = configured_path is not None
    path = Path(configured_path).expanduser() if configured_path else DEFAULT_CONFIG_PATH
    if not path.is_absolute():
        path = BACKEND_ROOT / path
    return path.resolve(), is_explicit


def _load_yaml(config_path: str | Path | None = None) -> tuple[dict[str, Any], Path]:
    path, is_explicit = _resolve_config_path(config_path)
    if not path.exists():
        if is_explicit:
            raise ConfigurationError(f"Backend config file does not exist: {path}")
        return {}, path

    try:
        loaded = yaml.safe_load(path.read_text(encoding="utf-8"))
    except yaml.YAMLError as exc:
        raise ConfigurationError(f"Invalid YAML in backend config file {path}: {exc}") from exc

    if loaded is None:
        return {}, path
    if not isinstance(loaded, dict):
        raise ConfigurationError(f"Backend config root must be a mapping: {path}")
    return loaded, path


def _yaml_value(config: dict[str, Any], section: str, key: str, default: Any) -> Any:
    section_value = config.get(section, {})
    if section_value is None:
        return default
    if not isinstance(section_value, dict):
        raise ConfigurationError(f"Config section '{section}' must be a mapping")
    return section_value.get(key, default)


def _string_setting(
    env_name: str,
    config: dict[str, Any],
    section: str,
    key: str,
    default: str,
) -> str:
    env_value = os.getenv(env_name)
    if env_value is not None:
        return env_value
    value = _yaml_value(config, section, key, default)
    return default if value is None else str(value)


def _int_setting(
    env_name: str,
    config: dict[str, Any],
    section: str,
    key: str,
    default: int,
) -> int:
    env_value = os.getenv(env_name)
    value = (
        env_value
        if env_value not in (None, "")
        else _yaml_value(config, section, key, default)
    )
    if isinstance(value, (bool, float)):
        raise ConfigurationError(f"Config value {section}.{key} must be an integer")
    try:
        return int(value)
    except (TypeError, ValueError) as exc:
        raise ConfigurationError(f"Config value {section}.{key} must be an integer") from exc


def _port_setting(
    env_name: str,
    config: dict[str, Any],
    key: str,
    default: int,
) -> int:
    port = _int_setting(env_name, config, "server", key, default)
    if not 1 <= port <= 65535:
        raise ConfigurationError(
            f"Config value server.{key} must be between 1 and 65535"
        )
    return port


def _list_setting(
    env_name: str,
    config: dict[str, Any],
    section: str,
    key: str,
    default: list[str],
) -> list[str]:
    env_value = os.getenv(env_name)
    value = (
        env_value
        if env_value and env_value.strip()
        else _yaml_value(config, section, key, default)
    )
    if isinstance(value, str):
        return [item.strip() for item in value.split(",") if item.strip()]
    if isinstance(value, list):
        return [str(item).strip() for item in value if str(item).strip()]
    raise ConfigurationError(
        f"Config value {section}.{key} must be a list or comma-separated string"
    )


def _path_setting(
    env_name: str,
    config: dict[str, Any],
    section: str,
    key: str,
    default: str,
    config_dir: Path,
) -> str:
    env_value = os.getenv(env_name)
    if env_value is not None:
        return env_value
    value = Path(str(_yaml_value(config, section, key, default))).expanduser()
    if not value.is_absolute():
        value = config_dir / value
    return str(value.resolve())


def _build_database_url(base_url: str, username: str, password: str) -> str:
    try:
        parsed = urlsplit(base_url)
        host = parsed.hostname
        port = parsed.port
    except ValueError as exc:
        raise ConfigurationError("Config value database.url is not a valid URL") from exc

    if not parsed.scheme or not host:
        raise ConfigurationError(
            "Config value database.url must include a scheme and host"
        )
    if parsed.scheme not in {"postgresql", "postgresql+psycopg"}:
        raise ConfigurationError(
            "Config value database.url must use postgresql or postgresql+psycopg"
        )
    if parsed.path in {"", "/"}:
        raise ConfigurationError("Config value database.url must include a database name")
    if not username:
        raise ConfigurationError("Config value database.username cannot be empty")
    if not password:
        raise ConfigurationError("Config value database.password cannot be empty")

    formatted_host = f"[{host}]" if ":" in host and not host.startswith("[") else host
    host_and_port = f"{formatted_host}:{port}" if port is not None else formatted_host
    credentials = f"{quote(username, safe='')}:{quote(password, safe='')}"
    parts = (
        parsed.scheme,
        f"{credentials}@{host_and_port}",
        parsed.path,
        parsed.query,
        parsed.fragment,
    )
    return urlunsplit(parts)


def _database_url(config: dict[str, Any]) -> str:
    # 保留原有完整连接串环境变量，便于 Docker 和生产环境注入密钥。
    complete_url_override = os.getenv("DATABASE_URL")
    if complete_url_override is not None:
        return complete_url_override

    base_url = _string_setting(
        "DATABASE_BASE_URL",
        config,
        "database",
        "url",
        DEFAULT_DATABASE_BASE_URL,
    )
    username = _string_setting(
        "DATABASE_USERNAME",
        config,
        "database",
        "username",
        DEFAULT_DATABASE_USERNAME,
    )
    password = _string_setting(
        "DATABASE_PASSWORD",
        config,
        "database",
        "password",
        DEFAULT_DATABASE_PASSWORD,
    )
    return _build_database_url(base_url, username, password)


@dataclass(frozen=True)
class Settings:
    app_name: str = "YG Fine-Tune Platform API"
    app_env: str = "local"
    route_prefix: str = "/modelTF"
    app_mode: str = "local"
    frontend_port: int = 16801
    backend_port: int = 17861
    database_url: str = field(
        default="postgresql+psycopg://yg_ft:change_me@localhost:5432/yg_ft",
        repr=False,
    )
    cors_allow_origins: list[str] = field(
        default_factory=lambda: [
            "http://localhost:16801",
            "http://127.0.0.1:16801",
        ]
    )
    compute_mode: str = "real"
    compute_status_sync_mode: str = "polling"
    compute_poll_interval_seconds: int = 3
    log_level: str = "INFO"
    log_dir: str = "./logs"
    log_file_prefix: str = "backend"
    log_error_file_prefix: str = "error"
    log_max_bytes: int = 20 * 1024 * 1024
    log_retention_days: int = 10


def load_settings(config_path: str | Path | None = None) -> Settings:
    config, resolved_config_path = _load_yaml(config_path)
    frontend_port = _port_setting(
        "FRONTEND_PORT", config, "frontend_port", Settings.frontend_port
    )
    backend_port = _port_setting(
        "BACKEND_PORT", config, "backend_port", Settings.backend_port
    )
    if frontend_port == backend_port:
        raise ConfigurationError(
            "Config values server.frontend_port and server.backend_port must be different"
        )
    default_cors = [
        f"http://localhost:{frontend_port}",
        f"http://127.0.0.1:{frontend_port}",
    ]
    return Settings(
        app_name=_string_setting("APP_NAME", config, "app", "name", Settings.app_name),
        app_env=_string_setting("APP_ENV", config, "app", "env", Settings.app_env),
        route_prefix=_string_setting(
            "MODELTF_ROUTE_PREFIX", config, "app", "route_prefix", Settings.route_prefix
        ),
        app_mode=_string_setting("APP_MODE", config, "app", "mode", Settings.app_mode),
        frontend_port=frontend_port,
        backend_port=backend_port,
        database_url=_database_url(config),
        cors_allow_origins=_list_setting(
            "CORS_ALLOW_ORIGINS", config, "app", "cors_allow_origins", default_cors
        ),
        compute_mode=_string_setting(
            "COMPUTE_MODE", config, "compute", "mode", Settings.compute_mode
        ),
        compute_status_sync_mode=_string_setting(
            "COMPUTE_STATUS_SYNC_MODE",
            config,
            "compute",
            "status_sync_mode",
            Settings.compute_status_sync_mode,
        ),
        compute_poll_interval_seconds=_int_setting(
            "COMPUTE_POLL_INTERVAL_SECONDS",
            config,
            "compute",
            "poll_interval_seconds",
            Settings.compute_poll_interval_seconds,
        ),
        log_level=_string_setting(
            "LOG_LEVEL", config, "logging", "level", Settings.log_level
        ),
        log_dir=_path_setting(
            "LOG_DIR",
            config,
            "logging",
            "directory",
            Settings.log_dir,
            resolved_config_path.parent,
        ),
        log_file_prefix=_string_setting(
            "LOG_FILE_PREFIX", config, "logging", "file_prefix", Settings.log_file_prefix
        ),
        log_error_file_prefix=_string_setting(
            "LOG_ERROR_FILE_PREFIX",
            config,
            "logging",
            "error_file_prefix",
            Settings.log_error_file_prefix,
        ),
        log_max_bytes=_int_setting(
            "LOG_MAX_BYTES", config, "logging", "max_bytes", Settings.log_max_bytes
        ),
        log_retention_days=_int_setting(
            "LOG_RETENTION_DAYS",
            config,
            "logging",
            "retention_days",
            Settings.log_retention_days,
        ),
    )


@lru_cache
def get_settings() -> Settings:
    return load_settings()