import os
from sqlalchemy import create_engine
from sqlalchemy.orm import sessionmaker
from dotenv import load_dotenv

# 加载 .env 文件
load_dotenv()

DATABASE_URL = os.getenv("DATABASE_URL")

#  自动将旧版的 postgres:// 替换为 SQLAlchemy 强要求的 postgresql://
if DATABASE_URL and DATABASE_URL.startswith("postgres://"):
    DATABASE_URL = DATABASE_URL.replace("postgres://", "postgresql://", 1)
elif DATABASE_URL and DATABASE_URL.startswith("postgres+"):
    DATABASE_URL = DATABASE_URL.replace("postgres+", "postgresql+", 1)

# 如果环境变量彻底缺失，做一个安全的报错提示
if not DATABASE_URL:
    raise ValueError("DATABASE_URL environment variable is not set or empty.")

engine = create_engine(
    DATABASE_URL,
    pool_recycle=3600,
    pool_pre_ping=True,
    connect_args={
        "keepalives": 1,
        "keepalives_idle": 30,
        "keepalives_interval": 10,
        "keepalives_count": 5
    }
)
SessionLocal = sessionmaker(bind=engine)

def get_db():
    db = SessionLocal()
    try:
        yield db
    finally:
        db.close()