"""Lab module tables + component nutrients/external_no + recipe.ration_type.

Revision ID: lab_module
Revises: daily_component_norm_dm_locked
Create Date: 2026-06-08
"""

from __future__ import annotations

import sqlalchemy as sa
from alembic import op

from app.lab.models import (
    LabAnimalProfile,
    LabRationLine,
    LabRecipeRation,
)


revision = "lab_module"
down_revision = "daily_component_norm_dm_locked"
branch_labels = None
depends_on = None


def _add_column_if_missing(table: str, column: sa.Column) -> None:
    bind = op.get_bind()
    insp = sa.inspect(bind)
    if table not in insp.get_table_names():
        return
    cols = {c["name"] for c in insp.get_columns(table)}
    if column.name not in cols:
        with op.batch_alter_table(table) as batch:
            batch.add_column(column)


def upgrade(**kw) -> None:
    if kw.get("tag") == "reports":
        return
    bind = op.get_bind()

    for model in (
        LabAnimalProfile,
        LabRecipeRation,
        LabRationLine,
    ):
        model.__table__.create(bind, checkfirst=True)

    _add_column_if_missing("component", sa.Column("external_no", sa.Integer(), nullable=True))

    _add_column_if_missing("recipe", sa.Column("ration_type", sa.String(length=10), nullable=True))


def downgrade(**kw) -> None:
    if kw.get("tag") == "reports":
        return
    bind = op.get_bind()
    insp = sa.inspect(bind)

    for table in (
        "lab_ration_line",
        "lab_recipe_ration",
        "lab_animal_profile",
    ):
        if table in insp.get_table_names():
            op.drop_table(table)

    if "component" in insp.get_table_names():
        cols = {c["name"] for c in insp.get_columns("component")}
        with op.batch_alter_table("component") as batch:
            if "external_no" in cols:
                batch.drop_column("external_no")

    if "recipe" in insp.get_table_names():
        cols = {c["name"] for c in insp.get_columns("recipe")}
        if "ration_type" in cols:
            with op.batch_alter_table("recipe") as batch:
                batch.drop_column("ration_type")