import os

from sqlalchemy import create_engine
from sqlalchemy.orm import sessionmaker, declarative_base, Session

BASE_DIR = os.path.dirname(os.path.abspath(__file__))
DB_FILE = os.path.join(BASE_DIR, "avatar.db")

engine = create_engine(
    f"sqlite:///{DB_FILE}",
    connect_args={"check_same_thread": False},
)
SessionLocal = sessionmaker(bind=engine, autoflush=False, expire_on_commit=False)
Base = declarative_base()


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


def init_db():
    import models

    Base.metadata.create_all(bind=engine)

    # 轻量迁移：为已存在的表补充新列（SQLite 不支持自动 ALTER，逐列尝试）
    _try_add_columns(
        ("qa_pairs", "enabled", "BOOLEAN DEFAULT 1"),
        ("knowledge_docs", "vectorized", "BOOLEAN DEFAULT 0"),
        ("knowledge_docs", "embedding_model", "VARCHAR DEFAULT ''"),
        ("knowledge_docs", "chunk_count", "INTEGER DEFAULT 0"),
        ("knowledge_docs", "vectorized_at", "TIMESTAMP"),
        ("avatars", "owner_id", "VARCHAR DEFAULT ''"),
    )


def _try_add_columns(*cols):
    with engine.connect() as conn:
        for table, col, ddl in cols:
            try:
                conn.exec_driver_sql(f"ALTER TABLE {table} ADD COLUMN {col} {ddl}")
                conn.commit()
            except Exception:
                # 列已存在（或全新库由 create_all 建好）则忽略
                pass