"""Phase 3: audit trail columns and data classification fields.

Revision ID: b029phase3
Revises: b028phase0
"""

from typing import Sequence, Union

import sqlalchemy as sa
from alembic import op

revision: str = "b029phase3"
down_revision: Union[str, None] = "b028phase0"
branch_labels: Union[str, Sequence[str], None] = None
depends_on: Union[str, Sequence[str], None] = None


def _column_names(table: str) -> set[str]:
    bind = op.get_bind()
    insp = sa.inspect(bind)
    return {c["name"] for c in insp.get_columns(table)}


def upgrade() -> None:
    audit_cols = _column_names("audit_logs")
    if "ip_address" not in audit_cols:
        op.add_column("audit_logs", sa.Column("ip_address", sa.String(45), nullable=True))
    if "user_agent" not in audit_cols:
        op.add_column("audit_logs", sa.Column("user_agent", sa.String(500), nullable=True))
    if "integrity_hash" not in audit_cols:
        op.add_column("audit_logs", sa.Column("integrity_hash", sa.String(64), nullable=True))
    if "session_id" not in audit_cols:
        op.add_column("audit_logs", sa.Column("session_id", sa.String(100), nullable=True))

    for table in ("tests", "evidences", "campaigns"):
        cols = _column_names(table)
        if "data_classification" not in cols:
            op.add_column(
                table,
                sa.Column(
                    "data_classification",
                    sa.String(20),
                    nullable=False,
                    server_default="internal",
                ),
            )


def downgrade() -> None:
    for table in ("campaigns", "evidences", "tests"):
        cols = _column_names(table)
        if "data_classification" in cols:
            op.drop_column(table, "data_classification")

    audit_cols = _column_names("audit_logs")
    for col in ("session_id", "integrity_hash", "user_agent", "ip_address"):
        if col in audit_cols:
            op.drop_column("audit_logs", col)