"""Introduce BoefjeConfig model

Revision ID: f9de6eb7824b
Revises: 6f99834a4a5a
Create Date: 2024-05-31 10:45:16.474714

"""

import logging

import sqlalchemy as sa
from alembic import op
from psycopg2._json import Json
from psycopg2.extensions import register_adapter
from psycopg2.extras import execute_values

from boefjes.local_repository import get_local_repository
from boefjes.models import Boefje, Normalizer

# revision identifiers, used by Alembic.
revision = "f9de6eb7824b"
down_revision = "6f99834a4a5a"
branch_labels = None
depends_on = None


logger = logging.getLogger(__name__)


def upgrade() -> None:
    # ### commands auto generated by Alembic - please adjust! ###
    op.create_table(
        "boefje_config",
        sa.Column("id", sa.Integer(), autoincrement=True, nullable=False),
        sa.Column("settings", sa.String(length=512), nullable=False, server_default="{}"),
        sa.Column("enabled", sa.Boolean(), nullable=False, server_default="false"),
        sa.Column("boefje_id", sa.Integer(), nullable=False),
        sa.Column("organisation_pk", sa.Integer(), nullable=False),
        sa.ForeignKeyConstraint(["boefje_id"], ["boefje.id"], ondelete="CASCADE"),
        sa.ForeignKeyConstraint(["organisation_pk"], ["organisation.pk"], ondelete="CASCADE"),
        sa.PrimaryKeyConstraint("id"),
        sa.UniqueConstraint("organisation_pk", "boefje_id", name="unique_boefje_config_per_organisation_per_boefje"),
    )
    op.create_table(
        "normalizer_config",
        sa.Column("id", sa.Integer(), autoincrement=True, nullable=False),
        sa.Column("enabled", sa.Boolean(), server_default="false", nullable=False),
        sa.Column("normalizer_id", sa.Integer(), nullable=False),
        sa.Column("organisation_pk", sa.Integer(), nullable=False),
        sa.ForeignKeyConstraint(["normalizer_id"], ["normalizer.id"], ondelete="CASCADE"),
        sa.ForeignKeyConstraint(["organisation_pk"], ["organisation.pk"], ondelete="CASCADE"),
        sa.PrimaryKeyConstraint("id"),
        sa.UniqueConstraint(
            "organisation_pk", "normalizer_id", name="unique_normalizer_config_per_organisation_per_normalizer"
        ),
    )

    op.add_column("boefje", sa.Column("static", sa.Boolean(), server_default="false", nullable=False))
    op.add_column("normalizer", sa.Column("static", sa.Boolean(), server_default="false", nullable=False))

    register_adapter(dict, Json)

    local_plugins = {plugin.id: plugin for plugin in get_local_repository().get_all()}
    connection = op.get_bind()

    # Get unique plugin_ids from the settings table for boefjes that do not exist yet in the database
    query = """
    SELECT DISTINCT s.plugin_id FROM settings s left join boefje b on b.plugin_id = s.plugin_id
        where b.plugin_id IS NULL
    """  # noqa: S608

    to_insert: list[Boefje] = []

    for plugin_id_output in connection.execute(query).fetchall():
        plugin_id = plugin_id_output[0]
        if plugin_id not in local_plugins:
            raise ValueError(f"Invalid plugin id found: {plugin_id}")

        # Since settings are boefje-only at this moment
        if local_plugins[plugin_id].type != "boefje":
            raise ValueError(f"Settings for normalizer or bit found: {plugin_id}. Remove these entries first.")

        res = connection.execute(f"SELECT id FROM boefje where plugin_id = '{plugin_id}'")  # noqa: S608
        if res.fetchone() is not None:
            continue  # The Boefje already exists

        if local_plugins[plugin_id].type == "boefje":
            to_insert.append(local_plugins[plugin_id])

    entries = [
        (
            boefje.id,
            boefje.name,
            boefje.description,
            str(boefje.scan_level),
            list(boefje.consumes),
            list(boefje.produces),
            ["TEST_KEY"],
            boefje.oci_image,
            boefje.oci_arguments,
            boefje.version,
        )
        for boefje in to_insert
    ]
    query = """INSERT INTO boefje (plugin_id, name, description, scan_level, consumes, produces, environment_keys,
        oci_image, oci_arguments, version) values %s"""

    with connection.begin():
        cursor = connection.connection.cursor()
        execute_values(cursor, query, entries)

    to_insert = []

    query = """
    SELECT DISTINCT p.plugin_id FROM plugin_state p left join boefje b on b.plugin_id = p.plugin_id
        where b.plugin_id IS NULL
    """

    for plugin_id_output in connection.execute(query).fetchall():
        plugin_id = plugin_id_output[0]
        if plugin_id not in local_plugins:
            logger.warning("Unknown plugin id found: %s. You might have to re-enable the plugin!", plugin_id)
            continue

        res = connection.execute(f"SELECT id FROM boefje where plugin_id = '{plugin_id}'")  # noqa: S608
        if res.fetchone() is not None:
            continue  # The Boefje already exists

        if local_plugins[plugin_id].type == "boefje":
            to_insert.append(local_plugins[plugin_id])

    entries = [
        (
            boefje.id,
            boefje.name,
            boefje.description,
            str(boefje.scan_level),
            list(boefje.consumes),
            list(boefje.produces),
            ["TEST_KEY"],
            boefje.oci_image,
            boefje.oci_arguments,
            boefje.version,
        )
        for boefje in to_insert
    ]
    query = """INSERT INTO boefje (plugin_id, name, description, scan_level, consumes, produces, environment_keys,
        oci_image, oci_arguments, version) values %s"""  # noqa: S608

    with connection.begin():
        cursor = connection.connection.cursor()
        execute_values(cursor, query, entries)

    normalizers_to_insert: list[Normalizer] = []
    query = """
    SELECT DISTINCT p.plugin_id FROM plugin_state p left join normalizer n on n.plugin_id = p.plugin_id
        where n.plugin_id IS NULL
    """  # noqa: S608

    for plugin_id_output in connection.execute(query).fetchall():
        plugin_id = plugin_id_output[0]
        if plugin_id not in local_plugins:
            logger.warning("Unknown plugin id found: %s. You might have to re-enable the plugin!", plugin_id)
            continue

        res = connection.execute(f"SELECT id FROM normalizer where plugin_id = '{plugin_id}'")  # noqa: S608
        if res.fetchone() is not None:
            continue  # The Normalizer already exists

        if local_plugins[plugin_id].type == "normalizer":
            normalizers_to_insert.append(local_plugins[plugin_id])

    normalizer_entries = [
        (
            normalizer.id,
            normalizer.name,
            normalizer.description,
            normalizer.consumes,
            normalizer.produces,
            ["TEST_KEY"],
            normalizer.version,
        )
        for normalizer in normalizers_to_insert
    ]
    query = """INSERT INTO normalizer (plugin_id, name, description, consumes, produces, environment_keys, version)
        values %s"""  # noqa: S608

    with connection.begin():
        cursor = connection.connection.cursor()
        execute_values(cursor, query, normalizer_entries)

    with connection.begin():
        connection.execute("""
            INSERT INTO boefje_config (settings, boefje_id, organisation_pk)
            SELECT s.values, b.id, s.organisation_pk from settings s
            join boefje b on s.plugin_id = b.plugin_id
        """)  # Add boefjes and set the settings for boefjes

    with connection.begin():
        connection.execute("""
            INSERT INTO boefje_config (enabled, boefje_id, organisation_pk)
            SELECT p.enabled, b.id, p.organisation_pk FROM plugin_state p
            JOIN boefje b ON p.plugin_id = b.plugin_id
            LEFT JOIN boefje_config bc ON bc.boefje_id = b.id WHERE bc.boefje_id IS NULL
        """)  # Add boefjes and set the enabled field for boefjes that to not exist yet
        connection.execute("""
            UPDATE boefje_config bc SET enabled = p.enabled from plugin_state p
            JOIN boefje b ON p.plugin_id = b.plugin_id
            where b.id = bc.boefje_id and p.organisation_pk = bc.organisation_pk
        """)  # Set the enabled field for boefjes
        connection.execute("""
            UPDATE normalizer_config nc SET enabled = p.enabled from plugin_state p
            JOIN normalizer n ON p.plugin_id = n.plugin_id
            where n.id = nc.normalizer_id and p.organisation_pk = nc.organisation_pk
        """)  # Set the enabled field for normalizers

    op.drop_table("settings")
    op.drop_table("plugin_state")
    # ### end Alembic commands ###


def downgrade() -> None:
    # ### commands auto generated by Alembic - please adjust! ###
    op.drop_column("normalizer", "static")
    op.drop_column("boefje", "static")

    op.create_table(
        "settings",
        sa.Column("id", sa.INTEGER(), autoincrement=True, nullable=False),
        sa.Column("values", sa.VARCHAR(length=512), autoincrement=False, nullable=False),
        sa.Column("plugin_id", sa.VARCHAR(length=64), autoincrement=False, nullable=False),
        sa.Column("organisation_pk", sa.INTEGER(), autoincrement=False, nullable=False),
        sa.ForeignKeyConstraint(
            ["organisation_pk"], ["organisation.pk"], name="settings_organisation_pk_fkey", ondelete="CASCADE"
        ),
        sa.PrimaryKeyConstraint("id", name="settings_pkey"),
        sa.UniqueConstraint("organisation_pk", "plugin_id", name="unique_settings_per_organisation_per_plugin"),
    )
    op.create_table(
        "plugin_state",
        sa.Column("id", sa.INTEGER(), autoincrement=True, nullable=False),
        sa.Column("plugin_id", sa.VARCHAR(length=64), autoincrement=False, nullable=False),
        sa.Column("enabled", sa.BOOLEAN(), autoincrement=False, nullable=False),
        sa.Column("organisation_pk", sa.INTEGER(), autoincrement=False, nullable=False),
        sa.ForeignKeyConstraint(
            ["organisation_pk"], ["organisation.pk"], name="plugin_state_organisation_pk_fkey", ondelete="CASCADE"
        ),
        sa.PrimaryKeyConstraint("id", name="plugin_state_pkey"),
        sa.UniqueConstraint("plugin_id", "organisation_pk", name="unique_plugin_id_per_org"),
    )

    connection = op.get_bind()
    with connection.begin():
        connection.execute("""
            INSERT INTO settings (values, plugin_id, organisation_pk)
            SELECT bc.settings, b.plugin_id, bc.organisation_pk from boefje_config bc
            join boefje b on bc.boefje_id = b.id
        """)

    with connection.begin():
        connection.execute("""
            INSERT INTO plugin_state (enabled, plugin_id, organisation_pk)
            SELECT bc.enabled, b.plugin_id, bc.organisation_pk from boefje_config bc
            join boefje b on bc.boefje_id = b.id
        """)

    with connection.begin():
        connection.execute("""
            INSERT INTO plugin_state (enabled, plugin_id, organisation_pk)
            SELECT nc.enabled, n.plugin_id, nc.organisation_pk from normalizer_config nc
            join normalizer n on nc.normalizer_id = n.id
        """)

    op.drop_table("boefje_config")
    op.drop_table("normalizer_config")

    # ### end Alembic commands ###
