from sqlalchemy.ext.asyncio import AsyncSession, create_async_engine, async_sessionmaker
from sqlalchemy.orm import DeclarativeBase
from app.config import settings

engine = create_async_engine(
    settings.DATABASE_URL,
    echo=settings.DEBUG,
    pool_size=10,
    max_overflow=20,
)

AsyncSessionLocal = async_sessionmaker(
    engine, class_=AsyncSession, expire_on_commit=False
)


class Base(DeclarativeBase):
    pass


async def get_db():
    async with AsyncSessionLocal() as session:
        try:
            yield session
            await session.commit()
        except Exception:
            await session.rollback()
            raise


async def init_db():
    async with engine.begin() as conn:
        await conn.run_sync(Base.metadata.create_all)
        await _run_migrations(conn)


async def _run_migrations(conn):
    """Postgres 专用：为已存在表追加新列（幂等）"""
    from sqlalchemy import text
    is_pg = engine.dialect.name == "postgresql"
    if not is_pg:
        return
    statements = [
        # production_plans 新字段
        "ALTER TABLE production_plans ADD COLUMN IF NOT EXISTS cold_coil_no VARCHAR(30)",
        "ALTER TABLE production_plans ADD COLUMN IF NOT EXISTS hot_coil_no VARCHAR(30)",
        "ALTER TABLE production_plans ADD COLUMN IF NOT EXISTS incoming_thickness DOUBLE PRECISION",
        "ALTER TABLE production_plans ADD COLUMN IF NOT EXISTS product_thickness DOUBLE PRECISION",
        "ALTER TABLE production_plans ADD COLUMN IF NOT EXISTS deviation_upper DOUBLE PRECISION",
        "ALTER TABLE production_plans ADD COLUMN IF NOT EXISTS deviation_lower DOUBLE PRECISION",
        "ALTER TABLE production_plans ADD COLUMN IF NOT EXISTS incoming_width DOUBLE PRECISION",
        "ALTER TABLE production_plans ADD COLUMN IF NOT EXISTS product_width DOUBLE PRECISION",
        "ALTER TABLE production_plans ADD COLUMN IF NOT EXISTS packaging_req VARCHAR(30)",
        "ALTER TABLE production_plans ADD COLUMN IF NOT EXISTS trim_req VARCHAR(30)",
        "ALTER TABLE production_plans ADD COLUMN IF NOT EXISTS rolling_mode VARCHAR(30)",
        "ALTER TABLE production_plans ADD COLUMN IF NOT EXISTS split_count INTEGER DEFAULT 1",
        "ALTER TABLE production_plans ADD COLUMN IF NOT EXISTS coil_diameter DOUBLE PRECISION",
        "ALTER TABLE production_plans ADD COLUMN IF NOT EXISTS next_process VARCHAR(30)",
        # 状态列改为 VARCHAR 以适配新值
        "ALTER TABLE production_plans ALTER COLUMN status TYPE VARCHAR(20) USING status::text",
        # production_records 新字段
        "ALTER TABLE production_records ADD COLUMN IF NOT EXISTS sub_coil_no VARCHAR(30)",
        "ALTER TABLE production_records ADD COLUMN IF NOT EXISTS hot_coil_no VARCHAR(30)",
        "ALTER TABLE production_records ADD COLUMN IF NOT EXISTS team VARCHAR(10)",
        "ALTER TABLE production_records ADD COLUMN IF NOT EXISTS steel_grade VARCHAR(30)",
        "ALTER TABLE production_records ADD COLUMN IF NOT EXISTS incoming_thickness DOUBLE PRECISION",
        "ALTER TABLE production_records ADD COLUMN IF NOT EXISTS deviation_upper DOUBLE PRECISION",
        "ALTER TABLE production_records ADD COLUMN IF NOT EXISTS deviation_lower DOUBLE PRECISION",
        "ALTER TABLE production_records ADD COLUMN IF NOT EXISTS incoming_width DOUBLE PRECISION",
        "ALTER TABLE production_records ADD COLUMN IF NOT EXISTS outlet_width DOUBLE PRECISION",
        "ALTER TABLE production_records ADD COLUMN IF NOT EXISTS incoming_weight DOUBLE PRECISION",
        "ALTER TABLE production_records ADD COLUMN IF NOT EXISTS weighed_weight DOUBLE PRECISION",
        "ALTER TABLE production_records ADD COLUMN IF NOT EXISTS packaging_req VARCHAR(30)",
        "ALTER TABLE production_records ADD COLUMN IF NOT EXISTS trim_req VARCHAR(30)",
        "ALTER TABLE production_records ADD COLUMN IF NOT EXISTS surface_quality VARCHAR(30)",
        "ALTER TABLE production_records ADD COLUMN IF NOT EXISTS product_quality DOUBLE PRECISION",
        "ALTER TABLE production_records ADD COLUMN IF NOT EXISTS product_length DOUBLE PRECISION",
        "ALTER TABLE production_records ADD COLUMN IF NOT EXISTS length_per_ton DOUBLE PRECISION",
        "ALTER TABLE production_records ADD COLUMN IF NOT EXISTS offline_time TIMESTAMP",
        "ALTER TABLE production_records ADD COLUMN IF NOT EXISTS status VARCHAR(20) DEFAULT 'UNWEIGH'",
        # qc_defect 新字段
        "ALTER TABLE qc_defect ADD COLUMN IF NOT EXISTS seq_no INTEGER",
        "ALTER TABLE qc_defect ADD COLUMN IF NOT EXISTS defect_desc VARCHAR(200)",
        "ALTER TABLE qc_defect ADD COLUMN IF NOT EXISTS start_position DOUBLE PRECISION",
        "ALTER TABLE qc_defect ADD COLUMN IF NOT EXISTS end_position DOUBLE PRECISION",
        "ALTER TABLE qc_defect ADD COLUMN IF NOT EXISTS length_val DOUBLE PRECISION",
        "ALTER TABLE qc_defect ADD COLUMN IF NOT EXISTS upper_surface BOOLEAN DEFAULT FALSE",
        "ALTER TABLE qc_defect ADD COLUMN IF NOT EXISTS lower_surface BOOLEAN DEFAULT FALSE",
        "ALTER TABLE qc_defect ADD COLUMN IF NOT EXISTS side_op BOOLEAN DEFAULT FALSE",
        "ALTER TABLE qc_defect ADD COLUMN IF NOT EXISTS side_middle BOOLEAN DEFAULT FALSE",
        "ALTER TABLE qc_defect ADD COLUMN IF NOT EXISTS side_drive BOOLEAN DEFAULT FALSE",
        "ALTER TABLE qc_defect ADD COLUMN IF NOT EXISTS is_main BOOLEAN DEFAULT FALSE",
        "ALTER TABLE qc_defect ADD COLUMN IF NOT EXISTS image_url VARCHAR(255)",
    ]
    for s in statements:
        try:
            await conn.execute(text(s))
        except Exception:
            pass