from sqlalchemy import create_engine
from sqlalchemy.ext.declarative import declarative_base
from sqlalchemy.orm import sessionmaker
import os

# 优先从环境变量读取，否则使用生成的默认配置
SQLALCHEMY_DATABASE_URL = os.getenv("DATABASE_URL")
if not SQLALCHEMY_DATABASE_URL:
    {% if db.url %}
    SQLALCHEMY_DATABASE_URL = "{{ db.url }}"
    {% else %}
    # 根据生成的详细参数构建
    DB_TYPE = "{{ db.type }}"
    DB_USER = "{{ db.user }}"
    DB_PASS = "{{ db.password }}"
    DB_HOST = "{{ db.host }}"
    DB_PORT = {{ db.port }}
    DB_NAME = "{{ db.name }}"
    
    driver = "postgresql+psycopg2" if DB_TYPE == "postgresql" else DB_TYPE
    if DB_TYPE == "mysql":
        driver = "mysql+pymysql"
        
    SQLALCHEMY_DATABASE_URL = f"{driver}://{DB_USER}:{DB_PASS}@{DB_HOST}:{DB_PORT}/{DB_NAME}"
    {% endif %}

connect_args = {}
if SQLALCHEMY_DATABASE_URL.startswith("sqlite"):
    connect_args = {"check_same_thread": False}

# 如果是 PostgreSQL 且指定了 Schema
execution_options = {}
{% if db.schema_name and db.type == 'postgresql' %}
execution_options = {"schema_translate_map": {None: "{{ db.schema_name }}"}}
{% endif %}

engine = create_engine(
    SQLALCHEMY_DATABASE_URL, 
    connect_args=connect_args,
    execution_options=execution_options
)
SessionLocal = sessionmaker(autocommit=False, autoflush=False, bind=engine)

Base = declarative_base()

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