"""Plan E — URS schema: impacts table, risk snapshots, asset criticality

Revision ID: 017
Revises: 016
Create Date: 2026-05-19 14:00:00.000000

Three additions backing the Unified Risk Score (URS) feature:

1. compliance_impacts  — per-check impact ratings imported from CIS
   benchmark CSVs (cis_id, benchmark, impact 0-100). Used to weight
   the SCA score → Asset Security Score (ASS).
2. asset_risk_snapshots — daily URS persistence for trend tracking.
3. assets.criticality + assets.compliance_frameworks — criticality
   multiplier for URS, mandatory frameworks list for per-framework
   breakdown.

Also adds compliance_results.weighted_score for the impact-weighted
variant of the existing simple score.

Idempotent via column-existence + table-existence guards.
"""
from alembic import op
import sqlalchemy as sa


revision = '017'
down_revision = '016'
branch_labels = None
depends_on = None


def upgrade() -> None:
    bind = op.get_bind()
    inspector = sa.inspect(bind)
    asset_cols = {c["name"] for c in inspector.get_columns("assets")}
    cr_cols = {c["name"] for c in inspector.get_columns("compliance_results")}
    tables = set(inspector.get_table_names())

    # --- assets ---
    if "criticality" not in asset_cols:
        op.execute(
            "ALTER TABLE assets "
            "ADD COLUMN IF NOT EXISTS criticality VARCHAR(16) "
            "NOT NULL DEFAULT 'normal'"
        )
    if "compliance_frameworks" not in asset_cols:
        op.execute(
            "ALTER TABLE assets "
            "ADD COLUMN IF NOT EXISTS compliance_frameworks TEXT"
        )

    # --- compliance_results.weighted_score ---
    if "weighted_score" not in cr_cols:
        op.execute(
            "ALTER TABLE compliance_results "
            "ADD COLUMN IF NOT EXISTS weighted_score DOUBLE PRECISION"
        )
        op.execute(
            "CREATE INDEX IF NOT EXISTS ix_compliance_results_weighted_score "
            "ON compliance_results (weighted_score)"
        )

    # --- compliance_impacts ---
    if "compliance_impacts" not in tables:
        op.create_table(
            "compliance_impacts",
            sa.Column("id", sa.Integer(), primary_key=True),
            sa.Column("cis_id", sa.String(64), nullable=False),
            sa.Column("benchmark", sa.String(128), nullable=False),
            sa.Column("level", sa.String(16), nullable=True),
            sa.Column("title", sa.String(500), nullable=True),
            sa.Column("impact", sa.Integer(), nullable=False, server_default="50"),
            sa.Column("created_at", sa.DateTime(), nullable=False,
                      server_default=sa.text("now()")),
            sa.Column("updated_at", sa.DateTime(), nullable=False,
                      server_default=sa.text("now()")),
            sa.UniqueConstraint("cis_id", "benchmark",
                                name="uq_compliance_impact_cis_benchmark"),
        )
        op.create_index("ix_compliance_impacts_cis_id",
                        "compliance_impacts", ["cis_id"])
        op.create_index("ix_compliance_impacts_benchmark",
                        "compliance_impacts", ["benchmark"])
        op.create_index("ix_compliance_impacts_benchmark_cis",
                        "compliance_impacts", ["benchmark", "cis_id"])

    # --- asset_risk_snapshots ---
    if "asset_risk_snapshots" not in tables:
        op.create_table(
            "asset_risk_snapshots",
            sa.Column("id", sa.Integer(), primary_key=True),
            sa.Column("asset_id", sa.Integer(),
                      sa.ForeignKey("assets.id", ondelete="CASCADE"),
                      nullable=False),
            sa.Column("snapshot_date", sa.DateTime(), nullable=False,
                      server_default=sa.text("now()")),
            sa.Column("avs", sa.Float(), nullable=True),
            sa.Column("ass", sa.Float(), nullable=True),
            sa.Column("urs", sa.Float(), nullable=True),
            sa.Column("severity", sa.String(16), nullable=True),
            sa.Column("created_at", sa.DateTime(), nullable=False,
                      server_default=sa.text("now()")),
            sa.Column("updated_at", sa.DateTime(), nullable=False,
                      server_default=sa.text("now()")),
            sa.UniqueConstraint("asset_id", "snapshot_date",
                                name="uq_asset_risk_per_day"),
        )
        op.create_index("ix_asset_risk_snapshots_asset_id",
                        "asset_risk_snapshots", ["asset_id"])
        op.create_index("ix_asset_risk_snapshots_snapshot_date",
                        "asset_risk_snapshots", ["snapshot_date"])


def downgrade() -> None:
    op.execute("DROP TABLE IF EXISTS asset_risk_snapshots")
    op.execute("DROP TABLE IF EXISTS compliance_impacts")
    op.execute("DROP INDEX IF EXISTS ix_compliance_results_weighted_score")
    op.execute("ALTER TABLE compliance_results DROP COLUMN IF EXISTS weighted_score")
    op.execute("ALTER TABLE assets DROP COLUMN IF EXISTS compliance_frameworks")
    op.execute("ALTER TABLE assets DROP COLUMN IF EXISTS criticality")
