"""Exercise PostgreSQL-only foundation invariants after Alembic migration."""

import asyncio
import os
import uuid

from sqlalchemy import text
from sqlalchemy.exc import DBAPIError
from sqlalchemy.ext.asyncio import create_async_engine


async def expect_append_only_rejection(engine, statement: str, parameters: dict) -> None:
    async with engine.connect() as connection:
        transaction = await connection.begin()
        try:
            await connection.execute(text(statement), parameters)
        except DBAPIError as error:
            await transaction.rollback()
            if "append-only" not in str(error.orig).lower():
                raise AssertionError("Mutation failed for an unexpected reason") from error
        else:
            await transaction.rollback()
            raise AssertionError("Append-only mutation unexpectedly succeeded")


async def verify() -> None:
    engine = create_async_engine(os.environ["DATABASE_URL"])
    workspace_id = uuid.uuid4()
    audit_id = uuid.uuid4()
    login_id = uuid.uuid4()
    try:
        async with engine.begin() as connection:
            installation = await connection.scalar(text(
                "SELECT owner_setup_completed FROM installation_state WHERE id = 1"
            ))
            if installation is not False:
                raise AssertionError("The one-time installation state was not seeded safely")
            await connection.execute(
                text(
                    "INSERT INTO workspaces "
                    "(id, name, timezone, audit_sequence) VALUES "
                    "(:id, 'Migration probe', 'UTC', 0)"
                ),
                {"id": workspace_id},
            )
            await connection.execute(
                text(
                    "INSERT INTO audit_events "
                    "(id, workspace_id, sequence, event_type, actor_type, actor_id, data, "
                    "event_hash) VALUES "
                    "(:id, :workspace_id, 1, 'migration.probe', 'system', :actor_id, "
                    "CAST(:data AS jsonb), :event_hash)"
                ),
                {
                    "id": audit_id,
                    "workspace_id": workspace_id,
                    "actor_id": uuid.uuid4(),
                    "data": "{}",
                    "event_hash": os.urandom(32),
                },
            )
            await connection.execute(
                text(
                    "INSERT INTO login_events "
                    "(id, email_hash, successful, reason) VALUES "
                    "(:id, :email_hash, false, 'migration_probe')"
                ),
                {"id": login_id, "email_hash": "0" * 64},
            )

        for table_name, record_id in (
            ("audit_events", audit_id),
            ("login_events", login_id),
        ):
            await expect_append_only_rejection(
                engine,
                f"UPDATE {table_name} SET created_at = now() WHERE id = :id",
                {"id": record_id},
            )
            await expect_append_only_rejection(
                engine,
                f"DELETE FROM {table_name} WHERE id = :id",
                {"id": record_id},
            )
    finally:
        await engine.dispose()


if __name__ == "__main__":
    asyncio.run(verify())
