config.py

"""
配置文件 - 管理数据库连接、API Key 等敏感信息

数据源：
  .env  → 全部配置（唯一数据源）
"""

import os
from datetime import datetime
from dotenv import load_dotenv

load_dotenv()

BASE_DIR = os.path.dirname(os.path.abspath(__file__))

# ==================== 达梦 8 默认数据库配置 ====================
DB_CONFIG = {
    "user": os.getenv("DM_USER", "SYSDBA"),
    "password": os.getenv("DM_PASSWORD", "SYSDBA"),
    "host": os.getenv("DM_HOST", "localhost"),
    "port": int(os.getenv("DM_PORT", "5236")),
    "schema": os.getenv("DM_SCHEMA", ""),
}

# ==================== 大模型 API 配置 ====================

_LLM_PRESETS = {
    "zhipu": {
        "provider": "zhipu",
        "api_key": os.getenv("ZHIPU_API_KEY", ""),
        "api_base": os.getenv("ZHIPU_API_BASE", "https://open.bigmodel.cn/api/paas/v4"),
        "model": os.getenv("ZHIPU_MODEL", "glm-4.7-flash"),
    },
    "deepseek": {
        "provider": "deepseek",
        "api_key": os.getenv("DEEPSEEK_API_KEY", ""),
        "api_base": os.getenv("DEEPSEEK_API_BASE", "https://api.deepseek.com/v1"),
        "model": os.getenv("DEEPSEEK_MODEL", "deepseek-chat"),
    },
}

_shared_llm = {
    "temperature": float(os.getenv("LLM_TEMPERATURE", "0")),
    "max_tokens": int(os.getenv("LLM_MAX_TOKENS", "2048")),
}

LLM_DEFAULT = os.getenv("LLM_DEFAULT", "zhipu")

LLM_CONFIG = dict(_LLM_PRESETS.get(LLM_DEFAULT, _LLM_PRESETS["zhipu"]))
LLM_CONFIG.update(_shared_llm)


def list_llm_providers() -> list:
    """列出所有已配置的 LLM 提供商"""
    result = []
    for name, cfg in _LLM_PRESETS.items():
        if cfg.get("api_key"):
            result.append(name)
    return result


def set_llm(name: str):
    """切换当前使用的 LLM 配置（会更新 LLM_CONFIG 字典内容）"""
    preset = _LLM_PRESETS.get(name)
    if preset is None:
        available = list(_LLM_PRESETS.keys())
        raise ValueError(f"未找到 LLM 配置: {name}，可用: {available}")
    LLM_CONFIG.clear()
    LLM_CONFIG.update(preset)
    LLM_CONFIG.update(_shared_llm)

# ==================== 项目常量 ====================
SQL_DIR = os.path.join(BASE_DIR, "sql")
EXPORT_DIR = os.getenv("EXPORT_DIR", os.path.join(BASE_DIR, "export_dir"))
DEFAULT_DB_TYPE = "dm8"
DEFAULT_INSTANCE = "local"
DEFAULT_USER = DB_CONFIG["user"]

ENABLE_SCHEMA_CACHE = True
ENABLE_AUTO_FIX_SQL = True
SAFE_MODE_ONLY_SELECT = True

# ==================== .env 数据库配置解析 ====================

_instances = {}
_users = {}
_env_loaded = False


def _load_env_databases():
    """解析 .env 中 DB_* 变量，构建 _instances 和 _users 字典"""
    global _env_loaded, _instances, _users
    if _env_loaded:
        return
    _env_loaded = True

    import re

    prefix_pat = re.compile(r"^DB_([A-Za-z0-9]+)_([A-Za-z0-9]+)_(.+)$")

    for key, value in os.environ.items():
        if not key.startswith("DB_"):
            continue
        m = prefix_pat.match(key)
        if not m:
            continue
        db_type = m.group(1).lower()
        instance = m.group(2).lower()
        suffix = m.group(3)

        if suffix == "HOST":
            _instances.setdefault(f"{db_type}:{instance}", {})["host"] = value
        elif suffix == "PORT":
            _instances.setdefault(f"{db_type}:{instance}", {})["port"] = value
        elif suffix == "USERS":
            _instances.setdefault(f"{db_type}:{instance}", {})
            prefix = f"DB_{db_type.upper()}_{instance.upper()}"
            for user in [u.strip() for u in value.split(",") if u.strip()]:
                user_upper = user.upper()
                password = os.getenv(f"{prefix}_PASS_{user_upper}", "")
                _users[f"{db_type}:{instance}:{user}"] = {
                    "user": user,
                    "password": password,
                }


# ==================== 工具函数 ====================

def timestamp_str() -> str:
    return datetime.now().strftime("%Y%m%d_%H%M%S")


def print_rules():
    rules_path = os.path.join(BASE_DIR, "AGENTS.md")
    try:
        with open(rules_path, encoding="utf-8") as f:
            print(f.read())
    except FileNotFoundError:
        print(f"[config] 未找到 AGENTS.md ({rules_path})")


# ==================== 配置枚举 ====================

def list_db_types() -> list:
    _load_env_databases()
    types = set()
    for key in _instances:
        types.add(key.split(":")[0])
    if not types:
        types.add(DEFAULT_DB_TYPE)
    return sorted(types)


def list_instances(db_type: str) -> list:
    _load_env_databases()
    instances = []
    for key in _instances:
        t, inst = key.split(":")
        if t == db_type:
            instances.append(inst)
    if not instances:
        instances.append(DEFAULT_INSTANCE)
    return instances


def list_users(db_type: str, instance: str) -> list:
    _load_env_databases()
    users = []
    for key in _users:
        t, inst, user = key.split(":")
        if t == db_type and inst == instance:
            users.append(_users[key]["user"])
    if not users:
        users.append(DB_CONFIG["user"])
    return users


# ==================== 配置读取 ====================

def get_instance_config(db_type: str, instance: str) -> dict:
    _load_env_databases()
    key = f"{db_type}:{instance}"
    if key in _instances:
        return dict(_instances[key])
    return {
        "host": DB_CONFIG["host"],
        "port": DB_CONFIG["port"],
    }


def get_user_credential(db_type: str, instance: str, user: str) -> dict:
    _load_env_databases()
    key = f"{db_type}:{instance}:{user}"
    if key in _users:
        return dict(_users[key])
    return {
        "user": DB_CONFIG["user"],
        "password": DB_CONFIG["password"],
    }


# ==================== 连接工厂 ====================

def _connect_dm8(host, port, user, password):
    """达梦8 数据库连接"""
    import dmPython
    return dmPython.connect(user=user, password=password, server=host, port=port)


def _connect_mysql(host, port, user, password):
    """MySQL / MariaDB / PolarDB / OceanBase / TiDB / GBase 连接"""
    import pymysql
    return pymysql.connect(host=host, port=port, user=user, password=password)


def _connect_pgsql(host, port, user, password):
    """PostgreSQL / KingbaseES / GaussDB / OpenGauss 连接"""
    import psycopg2
    return psycopg2.connect(host=host, port=port, user=user, password=password)


def _connect_oracle(host, port, user, password):
    """Oracle 数据库连接"""
    import cx_Oracle
    dsn = cx_Oracle.makedsn(host, port)
    return cx_Oracle.connect(user=user, password=password, dsn=dsn)


def _connect_sqlserver(host, port, user, password):
    """SQL Server 数据库连接"""
    import pymssql
    return pymssql.connect(server=host, port=port, user=user, password=password)


def _connect_sqlite(host, port, user, password):
    """SQLite 数据库连接（host 作为文件路径，为空时使用 :memory:）"""
    import sqlite3
    db_path = host if host else ":memory:"
    return sqlite3.connect(db_path)


_CONNECTORS = {
    "dm8": _connect_dm8,
    "mysql": _connect_mysql,
    "mariadb": _connect_mysql,
    "pgsql": _connect_pgsql,
    "oracle": _connect_oracle,
    "sqlserver": _connect_sqlserver,
    "sqlite": _connect_sqlite,
    "kingbase": _connect_pgsql,
    "gaussdb": _connect_pgsql,
    "opengauss": _connect_pgsql,
    "polardb": _connect_mysql,
    "oceanbase": _connect_mysql,
    "tidb": _connect_mysql,
    "gbase": _connect_mysql,
}


def create_connection(db_type: str = None, instance: str = None, user: str = None):
    if db_type is None:
        db_type = DEFAULT_DB_TYPE
    if instance is None:
        instance = DEFAULT_INSTANCE
    if user is None:
        user = DEFAULT_USER

    inst_cfg = get_instance_config(db_type, instance)
    user_cred = get_user_credential(db_type, instance, user)

    connector = _CONNECTORS.get(db_type)
    if connector is None:
        raise ValueError(
            f"不支持的数据库类型: {db_type}，支持的: {', '.join(sorted(_CONNECTORS))}"
        )

    host = inst_cfg.get("host", "localhost")
    port_str = inst_cfg.get("port", "")
    port = int(port_str) if port_str else 0
    username = user_cred["user"]
    password = user_cred["password"]

    return connector(host, port, username, password)