import signal
import pytest
import json
import os
import pg8000.dbapi as pg  # type: ignore
import psycopg

from typing import Generator
from conftest import (
    MAX_LOGIN_ATTEMPTS,
    Instance,
    Cluster,
    Retriable,
)


# Auth methods that don't require requests to remote servers (e.g. LDAP).
# TODO: maybe we should add chap-sha1 to pgproto via plain text auth.
LOCAL_AUTH_METHODS = ["scram-sha256", "chap-sha1", "md5"]
LOCAL_AUTH_METHODS_PGPROTO = [x for x in LOCAL_AUTH_METHODS if x != "chap-sha1"]


@pytest.fixture
def instance_with_audit_file(unstarted_instance: Instance):
    unstarted_instance.audit = os.path.join(unstarted_instance.instance_dir, "audit.log")
    unstarted_instance.start()
    unstarted_instance.wait_online()
    yield unstarted_instance


class AuditFile:
    def __init__(self, path):
        self.path = path
        self._f = open(path)
        self.parsed_bytes = 0

    def events(self) -> Generator[dict, None, None]:
        if self.parsed_bytes > 0:
            self._f.close()
            self._f = open(self.path)
            self._f.seek(self.parsed_bytes)

        while True:
            # NOTE: cannot use `for line in self._f:` because that breaks the
            # calls to `self._f.tell()` which we use below
            line = self._f.readline()

            if not line:
                # EOF
                break

            if not line.endswith("\n"):
                # Partial read. Possible in case of concurrent writes from picodata.
                raise Exception(f"Partial read: `{line}`")

            try:
                data = json.loads(line)
            except json.decoder.JSONDecodeError as e:
                e.args = (f"{e.args[0]} when parsing line: `{line}`", *e.args[1:])
                raise e from e

            self.parsed_bytes = self._f.tell()
            yield data


def take_until_title(events, title: str) -> dict | None:
    for event in events:
        if event["title"] == title:
            return event
    return None


def test_pytest_audit_file_events_api(unstarted_instance: Instance):
    """This test checks the `AuditFile` object itself.
    Checks that `AuditFile.events()` is reenterable, and handles partial reads.
    """

    filepath = os.path.join(unstarted_instance.instance_dir, "audit.log")
    writer = open(filepath, "w")

    audit_file = AuditFile(filepath)
    assert list(audit_file.events()) == []
    assert list(audit_file.events()) == []

    writer.write('{"key1":"value1"}\n')
    writer.flush()
    assert list(audit_file.events()) == [dict(key1="value1")]
    assert list(audit_file.events()) == []

    writer.write('{"key2":"value2"}\n{"key3":"value3"}\n')
    writer.flush()
    assert list(audit_file.events()) == [dict(key2="value2"), dict(key3="value3")]
    assert list(audit_file.events()) == []

    writer.write('{"partial":')
    writer.flush()
    with pytest.raises(Exception) as e:
        list(audit_file.events())
    assert e.value.args == ('Partial read: `{"partial":`',)

    writer.write('"read"}')
    writer.flush()
    with pytest.raises(Exception) as e:
        list(audit_file.events())
    assert e.value.args == ('Partial read: `{"partial":"read"}`',)

    writer.write("\n")
    writer.flush()
    assert list(audit_file.events()) == [dict(partial="read")]

    writer.write("GARBAGE\n")
    writer.flush()
    with pytest.raises(json.decoder.JSONDecodeError) as e:
        list(audit_file.events())
    assert e.value.args[0].endswith("when parsing line: `GARBAGE\n`")


def test_startup(instance_with_audit_file: Instance):
    instance = instance_with_audit_file
    instance.start()
    instance.terminate()

    events = list(AuditFile(instance.audit_flag_value).events())
    assert len(events) > 0

    # Check identifiers
    i = 1
    event: dict | None
    for event in events:
        assert event["id"] == f"1.0.{i}"
        i += 1

    # These should be the first two events
    assert events[0]["title"] == "init_audit"
    assert events[0]["message"] == "audit log is ready"
    assert events[0]["severity"] == "low"
    assert events[1]["title"] == "local_startup"
    assert events[1]["message"] == "instance is starting"
    assert events[1]["severity"] == "low"

    event = take_until_title(iter(events), "join_instance")
    assert event is not None
    assert event["instance_name"] == "default_1_1"
    assert event["raft_id"] == "1"
    assert event["initiator"] == "admin"

    event = take_until_title(iter(events), "change_target_state")
    assert event is not None
    assert event["new_state"] == "Offline(0)"
    assert event["instance_name"] == "default_1_1"
    assert event["raft_id"] == "1"
    assert event["message"] == f"target state of instance `{event['instance_name']}` changed to {event['new_state']}"
    assert event["severity"] == "low"
    assert event["initiator"] == "admin"

    event = take_until_title(iter(events), "change_current_state")
    assert event is not None
    assert event["new_state"] == "Offline(0)"
    assert event["instance_name"] == "default_1_1"
    assert event["raft_id"] == "1"
    assert event["message"] == f"current state of instance `{event['instance_name']}` changed to {event['new_state']}"
    assert event["severity"] == "medium"
    assert event["initiator"] == "admin"

    create_db = take_until_title(events, "create_local_db")
    assert create_db is not None
    assert create_db["initiator"] == "admin"
    assert create_db["instance_name"] == "default_1_1"
    assert create_db["raft_id"] == "1"

    event = take_until_title(events, "connect_local_db")
    assert event is not None
    assert event["initiator"] == "admin"
    assert event["instance_name"] == "default_1_1"
    assert event["raft_id"] == "1"


def test_integrity_violation(instance_with_audit_file: Instance):
    instance = instance_with_audit_file
    # Instance was up for some time
    instance.start()
    instance.terminate()

    # Then data files were modified and instance restarted
    snap_name = f"{instance.instance_dir}/00000000000000000000.snap"
    with open(snap_name, "w") as snap:
        snap.write("abc")

    # Instance should fail to start
    instance.fail_to_start()

    # Instance should  detect integrity violation
    events = list(AuditFile(instance.audit_flag_value).events())
    event = take_until_title(iter(events), "integrity_violation")
    assert event is not None
    assert event["message"] == "integrity violation detected"
    assert event["severity"] == "high"
    assert event["error"] == "XlogError: Unexpected end of file, run with 'force_recovery = true'"


def test_recover_database(instance_with_audit_file: Instance):
    instance = instance_with_audit_file
    instance.start()
    instance.wait_online()

    instance.restart()
    instance.wait_online()
    instance.terminate()

    events = AuditFile(instance.audit_flag_value).events()

    # At first local db (data storage) is initialized
    create_db = take_until_title(events, "create_local_db")
    assert create_db is not None
    assert create_db["initiator"] == "admin"
    assert create_db["instance_name"] == "default_1_1"
    assert create_db["raft_id"] == "1"

    # On restart instance recovers it's local data
    event = take_until_title(iter(events), "recover_local_db")
    assert event is not None
    assert event["initiator"] == "admin"
    assert event["instance_name"] == "default_1_1"
    assert event["raft_id"] == "1"

    event = take_until_title(events, "connect_local_db")
    assert event is not None
    assert event["initiator"] == "admin"
    assert event["instance_name"] == "default_1_1"
    assert event["raft_id"] == "1"


def test_create_drop_table(instance_with_audit_file: Instance):
    instance = instance_with_audit_file
    instance.start()
    instance.sql(
        """
        create table "foo" ("val" int not null, primary key ("val"))
        distributed by ("val")
        """
    )
    instance.sql(
        """
        drop table "foo"
        """
    )
    instance.terminate()

    events = AuditFile(instance.audit_flag_value).events()

    create_table = take_until_title(events, "create_table")
    assert create_table is not None
    assert create_table["name"] == "foo"
    assert create_table["message"] == "created table `foo`"
    assert create_table["severity"] == "medium"
    assert create_table["initiator"] == "pico_service"

    drop_table = take_until_title(events, "drop_table")
    assert drop_table is not None
    assert drop_table["name"] == "foo"
    assert drop_table["message"] == "dropped table `foo`"
    assert drop_table["severity"] == "medium"
    assert drop_table["initiator"] == "pico_service"


@pytest.mark.parametrize("auth_method", LOCAL_AUTH_METHODS_PGPROTO)
def test_user(auth_method: str, instance_with_audit_file: Instance):
    instance = instance_with_audit_file
    instance.start()
    instance.sql(
        f"""
        create user "ymir" with password 'T0psecret' using {auth_method}
        """
    )
    # TODO user cant change password without access to _pico_property
    # https://git.picodata.io/picodata/picodata/picodata/-/issues/449
    instance.sql(
        f"""
        alter user "ymir" password 'Topsecre1' using {auth_method}
        """,
        sudo=True,
    )
    instance.sql(
        """
        drop user "ymir"
        """
    )
    instance.terminate()

    events = AuditFile(instance.audit_flag_value).events()

    create_user = take_until_title(events, "create_user")
    assert create_user is not None
    assert create_user["user"] == "ymir"
    assert create_user["auth_type"] == auth_method
    assert create_user["message"] == f"created user `{create_user['user']}`"
    assert create_user["severity"] == "high"
    assert create_user["initiator"] == "pico_service"

    change_password = take_until_title(events, "change_password")
    assert change_password is not None
    assert change_password["user"] == "ymir"
    assert change_password["auth_type"] == auth_method
    assert change_password["message"] == f"password of user `{change_password['user']}` was changed"
    assert change_password["severity"] == "high"
    assert change_password["initiator"] == "admin"

    drop_user = take_until_title(events, "drop_user")
    assert drop_user is not None
    assert drop_user["user"] == "ymir"
    assert drop_user["message"] == f"dropped user `{drop_user['user']}`"
    assert drop_user["severity"] == "medium"
    assert drop_user["initiator"] == "pico_service"


@pytest.mark.parametrize("auth_method", LOCAL_AUTH_METHODS_PGPROTO)
def test_role(auth_method: str, instance_with_audit_file: Instance):
    instance = instance_with_audit_file
    instance.start()

    setup = [
        f"""
        create user "bubba" with password 'T0psecret' using {auth_method}
        """,
        """
        grant create role to "bubba"
        """,
    ]
    for query in setup:
        instance.sql(query, sudo=True)

    with instance.connect_via_pgproto(user="bubba", password="T0psecret") as c:
        c.execute(
            """
            create role "skibidi"
            """
        )
        c.execute(
            """
            create role "dummy"
            """
        )
        c.execute(
            """
            grant "dummy" to "skibidi"
            """
        )
        c.execute(
            """
            revoke "dummy" from "skibidi"
            """
        )
        c.execute(
            """
            drop role "skibidi"
            """
        )
    instance.terminate()

    events = AuditFile(instance.audit_flag_value).events()

    create_role = take_until_title(events, "create_role")
    assert create_role is not None
    assert create_role["role"] == "skibidi"
    assert create_role["message"] == f"created role `{create_role['role']}`"
    assert create_role["severity"] == "high"
    assert create_role["initiator"] == "bubba"

    grant_role = take_until_title(events, "grant_role")
    assert grant_role is not None
    assert grant_role["role"] == "dummy"
    assert grant_role["grantee"] == "skibidi"
    assert grant_role["grantee_type"] == "role"
    assert grant_role["message"] == f"granted role `{grant_role['role']}` to role `{grant_role['grantee']}`"
    assert grant_role["severity"] == "high"
    assert grant_role["initiator"] == "bubba"

    revoke_role = take_until_title(events, "revoke_role")
    assert revoke_role is not None
    assert revoke_role["role"] == "dummy"
    assert revoke_role["grantee"] == "skibidi"
    assert revoke_role["grantee_type"] == "role"
    assert revoke_role["message"] == f"revoked role `{grant_role['role']}` from role `{revoke_role['grantee']}`"
    assert revoke_role["severity"] == "high"
    assert revoke_role["initiator"] == "bubba"

    drop_role = take_until_title(events, "drop_role")
    assert drop_role is not None
    assert drop_role["role"] == "skibidi"
    assert drop_role["message"] == f"dropped role `{drop_role['role']}`"
    assert drop_role["severity"] == "medium"
    assert drop_role["initiator"] == "bubba"


def test_index(instance_with_audit_file: Instance):
    instance = instance_with_audit_file
    instance.start()
    instance.sql(
        """
        create table "foo" ("val" int not null, primary key ("val"))
        distributed by ("val")
        """
    )
    instance.sql(""" create index "foo_idx" on "foo" ("val") """)
    instance.sql(""" drop index "foo_idx" """)
    instance.terminate()

    events = AuditFile(instance.audit_flag_value).events()

    create_index = take_until_title(events, "create_index")
    assert create_index is not None
    assert create_index["name"] == "foo_idx"
    assert create_index["message"] == f"created index `{create_index['name']}`"
    assert create_index["severity"] == "medium"
    assert create_index["initiator"] == "pico_service"

    drop_index = take_until_title(events, "drop_index")
    assert drop_index is not None
    assert drop_index["name"] == "foo_idx"
    assert drop_index["message"] == f"dropped index `{drop_index['name']}`"
    assert drop_index["severity"] == "medium"
    assert drop_index["initiator"] == "pico_service"


def test_join_expel_instance(cluster: Cluster):
    cluster.deploy(instance_count=0)
    audit = os.path.join(cluster.data_dir, "audit.log")
    i1 = cluster.add_instance(audit=audit, replicaset_name="default_1")

    audit_i1 = AuditFile(i1.audit_flag_value)
    for _ in audit_i1.events():
        pass
    events = audit_i1.events()

    audit = os.path.join(cluster.data_dir, "audit.log")
    i2 = cluster.add_instance(name="i2", audit=audit, replicaset_name="default_1")

    join_instance = take_until_title(events, "join_instance")
    assert join_instance is not None
    assert join_instance["instance_name"] == "i2"
    assert join_instance["raft_id"] == str(i2.raft_id)
    assert join_instance["severity"] == "low"
    assert join_instance["initiator"] == "admin"

    # Joining of a new instance is treated as creation of a local database on it
    create_db = take_until_title(events, "create_local_db")
    assert create_db is not None
    assert create_db["instance_name"] == "i2"
    assert create_db["raft_id"] == str(i2.raft_id)
    assert create_db["severity"] == "low"
    assert create_db["initiator"] == "admin"

    cluster.expel(i2, force=True)
    cluster.wait_has_states(i2, "Expelled", "Expelled")

    expel_instance = take_until_title(events, "expel_instance")
    assert expel_instance is not None
    assert expel_instance["instance_name"] == "i2"
    assert expel_instance["raft_id"] == str(i2.raft_id)
    assert expel_instance["severity"] == "low"
    assert expel_instance["initiator"] == "admin"

    # Expelling an instance is treated as removal of a local database on it
    drop_db = take_until_title(events, "drop_local_db")
    assert drop_db is not None
    assert drop_db["instance_name"] == "i2"
    assert drop_db["raft_id"] == str(i2.raft_id)
    assert drop_db["severity"] == "low"
    assert drop_db["initiator"] == "admin"


def test_join_connect_instance(cluster: Cluster):
    cluster.deploy(instance_count=0)
    audit = os.path.join(cluster.data_dir, "audit.log")
    i1 = cluster.add_instance(audit=audit)

    audit = os.path.join(cluster.data_dir, "audit.log")
    i2 = cluster.add_instance(name="i2", audit=audit)
    i2.terminate()

    events = AuditFile(i2.audit_flag_value).events()

    create_db = take_until_title(events, "create_local_db")
    assert create_db is not None
    assert create_db["instance_name"] == "default_1_1"
    assert create_db["raft_id"] == str(i1.raft_id)
    assert create_db["severity"] == "low"
    assert create_db["initiator"] == "admin"

    create_db = take_until_title(events, "create_local_db")
    assert create_db is not None
    assert create_db["instance_name"] == "i2"
    assert create_db["raft_id"] == str(i2.raft_id)
    assert create_db["severity"] == "low"
    assert create_db["initiator"] == "admin"

    # Joining instance should emit 'connect_local_db' audit event.
    # connect_local_db is a local event as only joining instance emits it
    # in contrast to create_local_db
    connect_db = take_until_title(events, "connect_local_db")
    assert connect_db is not None
    assert connect_db["instance_name"] == "i2"
    assert connect_db["raft_id"] == str(i2.raft_id)
    assert connect_db["severity"] == "low"
    assert connect_db["initiator"] == "admin"


@pytest.mark.parametrize("auth_method", LOCAL_AUTH_METHODS_PGPROTO)
def test_auth(auth_method: str, instance_with_audit_file: Instance):
    instance = instance_with_audit_file
    instance.start()

    instance.sql(
        f"""
        create user "ymir" with password 'T0psecret' using {auth_method}
        """,
        sudo=True,
    )
    instance.sql(
        """
        alter user "ymir" login
        """,
        sudo=True,
    )

    audit = AuditFile(instance.audit_flag_value)
    for _ in audit.events():
        pass
    events = audit.events()

    with instance.connect_via_pgproto(timeout=4, user="ymir", password="T0psecret"):
        pass

    auth_ok = take_until_title(events, "auth_ok")
    assert auth_ok is not None
    assert auth_ok["user"] == "ymir"
    assert auth_ok["severity"] == "high"
    assert auth_ok["initiator"] == "ymir"
    assert auth_ok["verdict"] == "user is not blocked"

    for _ in range(MAX_LOGIN_ATTEMPTS):
        with pytest.raises(psycopg.Error, match="authentication failed for user"):
            with instance.connect_via_pgproto(timeout=4, user="ymir", password="wrong_pwd"):
                pass

        auth_fail = take_until_title(events, "auth_fail")
        assert auth_fail is not None
        assert auth_fail["message"] == "failed to authenticate user `ymir`"
        assert auth_fail["severity"] == "high"
        assert auth_fail["verdict"] == "user is not blocked"
        assert auth_fail["user"] == "ymir"
        assert auth_fail["initiator"] == "ymir"

    with pytest.raises(psycopg.Error, match="authentication failed for user"):
        with instance.connect_via_pgproto(timeout=4, user="ymir", password="Wr0ng_pwd"):
            pass

    auth_fail = take_until_title(events, "auth_fail")
    assert auth_fail is not None
    assert auth_fail["message"] == "failed to authenticate user `ymir`"
    assert auth_fail["severity"] == "high"
    assert auth_fail["verdict"] == ("Maximum number of login attempts exceeded; user blocked")
    assert auth_fail["user"] == "ymir"
    assert auth_fail["initiator"] == "ymir"


@pytest.mark.parametrize("auth_method", LOCAL_AUTH_METHODS_PGPROTO)
def test_access_denied(auth_method: str, instance_with_audit_file: Instance):
    instance = instance_with_audit_file
    instance.start()

    instance.create_user(with_name="ymir", with_password="T0psecret", with_auth=auth_method)

    audit = AuditFile(instance.audit_flag_value)
    for _ in audit.events():
        pass

    events = audit.events()

    expected_error = "Create access to role 'R' is denied for user 'ymir'"
    expected_audit = "create access denied on role `R` for user `ymir`"

    with pytest.raises(
        Exception,
        match=expected_error,
    ):
        with instance.connect_via_pgproto(user="ymir", password="T0psecret") as c:
            c.execute('CREATE ROLE "R"')

    access_denied = take_until_title(events, "access_denied")
    assert access_denied is not None
    assert access_denied["message"] == expected_audit
    assert access_denied["severity"] == "medium"
    assert access_denied["privilege"] == "create"
    assert access_denied["object_type"] == "role"
    assert access_denied["object"] == "R"
    assert access_denied["initiator"] == "ymir"


@pytest.mark.parametrize("auth_method", LOCAL_AUTH_METHODS_PGPROTO)
def test_grant_revoke(auth_method: str, instance_with_audit_file: Instance):
    instance = instance_with_audit_file
    instance.start()

    user = "ymir"
    password = "T0psecret"

    instance.create_user(with_name=user, with_password=password, with_auth=auth_method)

    instance.sql(f'GRANT CREATE ROLE TO "{user}"', sudo=True)

    audit = AuditFile(instance.audit_flag_value)
    for _ in audit.events():
        pass

    events = audit.events()

    # wildcard privilege to user
    instance.sql(f'GRANT CREATE TABLE TO "{user}"', sudo=True)

    grant_privilege = take_until_title(events, "grant_privilege")

    assert grant_privilege is not None
    assert grant_privilege["message"] == "granted privilege create on table `*` to user `ymir`"
    assert grant_privilege["severity"] == "high"
    assert grant_privilege["privilege"] == "create"
    assert grant_privilege["object"] == "*"
    assert grant_privilege["object_type"] == "table"
    assert grant_privilege["grantee"] == user
    assert grant_privilege["grantee_type"] == "user"
    assert grant_privilege["initiator"] == "admin"

    instance.sql(f'REVOKE CREATE TABLE FROM "{user}"', sudo=True)

    revoke_privilege = take_until_title(events, "revoke_privilege")

    assert revoke_privilege is not None
    assert revoke_privilege["message"] == f"revoked privilege create on table `*` from user `{user}`"
    assert revoke_privilege["severity"] == "high"
    assert revoke_privilege["privilege"] == "create"
    assert revoke_privilege["object"] == "*"
    assert revoke_privilege["object_type"] == "table"
    assert revoke_privilege["grantee"] == user
    assert revoke_privilege["grantee_type"] == "user"
    assert revoke_privilege["initiator"] == "admin"

    # specific privilege to user
    instance.sql(f'GRANT READ ON TABLE "_pico_tier" TO "{user}"', sudo=True)

    grant_privilege = take_until_title(events, "grant_privilege")

    assert grant_privilege is not None
    assert grant_privilege["message"] == f"granted privilege read on table `_pico_tier` to user `{user}`"
    assert grant_privilege["severity"] == "high"
    assert grant_privilege["privilege"] == "read"
    assert grant_privilege["object_type"] == "table"
    assert grant_privilege["grantee"] == user
    assert grant_privilege["grantee_type"] == "user"
    assert grant_privilege["initiator"] == "admin"
    assert grant_privilege["object"] == "_pico_tier"

    instance.sql(f'REVOKE READ ON TABLE "_pico_tier" FROM "{user}"', sudo=True)

    revoke_privilege = take_until_title(events, "revoke_privilege")

    assert revoke_privilege is not None
    assert revoke_privilege["message"] == f"revoked privilege read on table `_pico_tier` from user `{user}`"
    assert revoke_privilege["severity"] == "high"
    assert revoke_privilege["privilege"] == "read"
    assert revoke_privilege["object_type"] == "table"
    assert revoke_privilege["grantee"] == user
    assert revoke_privilege["grantee_type"] == "user"
    assert revoke_privilege["initiator"] == "admin"
    assert revoke_privilege["object"] == "_pico_tier"

    with instance.connect_via_pgproto(user=user, password=password) as c:
        c.execute('CREATE ROLE "R"')

    # wildcard privilege to role
    instance.sql('GRANT CREATE TABLE TO "R"', sudo=True)

    grant_privilege = take_until_title(events, "grant_privilege")

    assert grant_privilege is not None
    assert grant_privilege["message"] == "granted privilege create on table `*` to role `R`"
    assert grant_privilege["severity"] == "high"
    assert grant_privilege["privilege"] == "create"
    assert grant_privilege["object"] == "*"
    assert grant_privilege["object_type"] == "table"
    assert grant_privilege["grantee"] == "R"
    assert grant_privilege["grantee_type"] == "role"
    assert grant_privilege["initiator"] == "admin"

    instance.sql('REVOKE CREATE TABLE FROM "R"', sudo=True)

    revoke_privilege = take_until_title(events, "revoke_privilege")

    assert revoke_privilege is not None
    assert revoke_privilege["message"] == "revoked privilege create on table `*` from role `R`"
    assert revoke_privilege["severity"] == "high"
    assert revoke_privilege["privilege"] == "create"
    assert grant_privilege["object"] == "*"
    assert revoke_privilege["object_type"] == "table"
    assert revoke_privilege["grantee"] == "R"
    assert revoke_privilege["grantee_type"] == "role"
    assert revoke_privilege["initiator"] == "admin"

    # specific privilege to role
    instance.sql('GRANT READ ON TABLE "_pico_user" TO "R"', sudo=True)

    grant_privilege = take_until_title(events, "grant_privilege")

    assert grant_privilege is not None
    assert grant_privilege["message"] == "granted privilege read on table `_pico_user` to role `R`"
    assert grant_privilege["severity"] == "high"
    assert grant_privilege["privilege"] == "read"
    assert grant_privilege["object_type"] == "table"
    assert grant_privilege["grantee"] == "R"
    assert grant_privilege["grantee_type"] == "role"
    assert grant_privilege["initiator"] == "admin"
    assert grant_privilege["object"] == "_pico_user"

    instance.sql('REVOKE READ ON TABLE "_pico_user" FROM "R"', sudo=True)

    revoke_privilege = take_until_title(events, "revoke_privilege")

    assert revoke_privilege is not None
    assert revoke_privilege["message"] == "revoked privilege read on table `_pico_user` from role `R`"
    assert revoke_privilege["severity"] == "high"
    assert revoke_privilege["privilege"] == "read"
    assert revoke_privilege["object_type"] == "table"
    assert revoke_privilege["grantee"] == "R"
    assert revoke_privilege["grantee_type"] == "role"
    assert revoke_privilege["initiator"] == "admin"
    assert revoke_privilege["object"] == "_pico_user"

    # role to user
    with instance.connect_via_pgproto(user=user, password=password) as c:
        c.execute('GRANT "R" TO "ymir"')

    grant_role = take_until_title(events, "grant_role")
    assert grant_role is not None
    assert grant_role["message"] == "granted role `R` to user `ymir`"
    assert grant_role["severity"] == "high"
    assert grant_role["role"] == "R"
    assert grant_role["grantee"] == "ymir"
    assert grant_role["grantee_type"] == "user"
    assert grant_role["initiator"] == "ymir"

    with instance.connect_via_pgproto(user=user, password=password) as c:
        c.execute('REVOKE "R" FROM "ymir"')

    revoke_role = take_until_title(events, "revoke_role")

    assert revoke_role is not None
    assert revoke_role["message"] == f"revoked role `R` from user `{user}`"
    assert revoke_role["severity"] == "high"
    assert revoke_role["role"] == "R"
    assert revoke_role["grantee"] == user
    assert revoke_role["grantee_type"] == "user"
    assert revoke_role["initiator"] == user

    # one role to another role
    with instance.connect_via_pgproto(user=user, password=password) as c:
        c.execute('CREATE ROLE "R2"')
        c.execute('GRANT "R" TO "R2"')

    grant_role = take_until_title(events, "grant_role")
    assert grant_role is not None
    assert grant_role["message"] == "granted role `R` to role `R2`"
    assert grant_role["severity"] == "high"
    assert grant_role["role"] == "R"
    assert grant_role["grantee"] == "R2"
    assert grant_role["grantee_type"] == "role"
    assert grant_role["initiator"] == "ymir"

    with instance.connect_via_pgproto(user=user, password=password) as c:
        c.execute('REVOKE "R" FROM "R2"')

    revoke_role = take_until_title(events, "revoke_role")

    assert revoke_role is not None
    assert revoke_role["message"] == "revoked role `R` from role `R2`"
    assert revoke_role["severity"] == "high"
    assert revoke_role["role"] == "R"
    assert revoke_role["grantee"] == "R2"
    assert revoke_role["grantee_type"] == "role"
    assert revoke_role["initiator"] == user


def check_audit_log_rotate(audit: AuditFile):
    rotate = take_until_title(audit.events(), "audit_rotate")
    assert rotate is not None
    assert rotate["message"] == "log file has been reopened"
    assert rotate["severity"] == "low"


def test_rotation(cluster: Cluster):
    instance = cluster.add_instance(wait_online=False, log_to_file=True, log_to_console=False)
    instance.audit = os.path.join(instance.instance_dir, "audit.log")
    instance.start_and_wait()

    user = "ymir"
    password = "T0psecret"

    audit = AuditFile(instance.audit_flag_value)
    for _ in audit.events():
        pass

    instance.create_user(with_name=user, with_password=password)
    instance.sql(f'GRANT CREATE ROLE TO "{user}"', sudo=True)

    assert instance.process is not None
    instance.process.send_signal(sig=signal.SIGHUP)

    Retriable(fatal=json.decoder.JSONDecodeError).call(check_audit_log_rotate, audit)


def test_rename_user(instance_with_audit_file: Instance):
    instance = instance_with_audit_file
    instance.start()

    user = "ymir"
    password = "T0psecret"

    instance.create_user(with_name=user, with_password=password)

    instance.sql('ALTER USER "ymir" RENAME TO "TOM"', sudo=True)

    instance.terminate()

    events = AuditFile(instance.audit_flag_value).events()

    rename_user = take_until_title(events, "rename_user")
    assert rename_user is not None
    assert rename_user["message"] == "name of user `ymir` was changed to `TOM`"
    assert rename_user["severity"] == "high"
    assert rename_user["old_name"] == "ymir"
    assert rename_user["new_name"] == "TOM"


@pytest.mark.parametrize(
    "table_type",
    [("global"), ("sharded")],
)
@pytest.mark.parametrize(
    "protocol_type",
    [("pgproto"), ("iproto")],
)
def test_audit_via_sql_simple(instance_with_audit_file: Instance, table_type: str, protocol_type: str):
    instance = instance_with_audit_file
    instance.start()

    user = "pico_service"

    if protocol_type == "pgproto":
        password = "P@ssw0rd"
        instance.sql(f"ALTER USER \"{user}\" WITH PASSWORD '{password}'")

        os.environ["PGSSLMODE"] = "disable"
        conn = pg.Connection(user, password=password, host=instance.pg_host, port=instance.pg_port)
        conn.autocommit = True
        cur = conn.cursor()
        execute_func = cur.execute
    else:
        execute_func = instance.sql

    execute_func(
        f"""
        AUDIT POLICY dml_default BY {user};
        """
    )
    distributed_clause = "GLOBALLY" if table_type == "global" else "BY (id)"
    execute_func(
        f"""
        CREATE TABLE test (id UNSIGNED NOT NULL, value TEXT NOT NULL, PRIMARY KEY (id))
        DISTRIBUTED {distributed_clause}
        """
    )
    execute_func(
        """
        SELECT * FROM test
        """
    )
    execute_func("INSERT INTO test VALUES (42, 'my_value')")
    execute_func("UPDATE test SET value = 'my_value2' WHERE id = 42")
    execute_func("CREATE PROCEDURE test_proc() language sql AS $$ UPDATE test SET value = 'my_value3' WHERE id = 42 $$")
    execute_func("CALL test_proc()")
    execute_func("DELETE FROM test WHERE id = 42")
    execute_func(
        f"""
        AUDIT POLICY dml_default EXCEPT {user}
        """
    )
    execute_func("INSERT INTO test VALUES (42, 'my_value')")
    instance.terminate()

    events = AuditFile(instance.audit_flag_value).events()

    audit_policy = take_until_title(events, "audit_policy")
    assert audit_policy is not None
    assert audit_policy["message"] == f"audit policy `dml_default` for user `{user}` was turned on"
    assert audit_policy["severity"] == "high"
    assert audit_policy["initiator"] == user

    dml_insert = take_until_title(events, "dml")
    assert dml_insert is not None
    assert dml_insert["message"] == "apply `INSERT INTO test VALUES (42, 'my_value')`"
    assert dml_insert["severity"] == "medium"
    assert dml_insert["initiator"] == user

    dml_update = take_until_title(events, "dml")
    assert dml_update is not None
    assert dml_update["message"] == "apply `UPDATE test SET value = 'my_value2' WHERE id = 42`"
    assert dml_update["severity"] == "medium"
    assert dml_update["initiator"] == user

    dml_update = take_until_title(events, "create_procedure")
    assert dml_update is not None
    assert dml_update["message"] == "created procedure `test_proc`"
    assert dml_update["severity"] == "medium"
    assert dml_update["initiator"] == user

    dml_update = take_until_title(events, "dml")
    assert dml_update is not None
    assert dml_update["message"] == "apply `UPDATE test SET value = 'my_value3' WHERE id = 42`"
    assert dml_update["severity"] == "medium"
    assert dml_update["initiator"] == user

    dml_delete = take_until_title(events, "dml")
    assert dml_delete is not None
    assert dml_delete["message"] == "apply `DELETE FROM test WHERE id = 42`"
    assert dml_delete["severity"] == "medium"
    assert dml_delete["initiator"] == user

    audit_policy = take_until_title(events, "audit_policy")
    assert audit_policy is not None
    assert audit_policy["message"] == f"audit policy `dml_default` for user `{user}` was turned off"
    assert audit_policy["severity"] == "high"
    assert audit_policy["initiator"] == user

    dml_insert = take_until_title(events, "dml")
    assert dml_insert is None


@pytest.mark.parametrize(
    "table_type",
    [("global"), ("sharded")],
)
def test_audit_via_sql_pgproto_prepared(instance_with_audit_file: Instance, table_type: str):
    instance = instance_with_audit_file
    instance.start()

    user = "admin"
    password = "P@ssw0rd"
    instance.sql(f"ALTER USER \"{user}\" WITH PASSWORD '{password}'")
    host = instance.pg_host
    port = instance.pg_port

    conn = psycopg.connect(f"user = {user} password={password} host={host} port={port} sslmode=disable")
    conn.autocommit = True

    conn.execute(
        """
        AUDIT POLICY dml_default BY admin;
        """
    )
    distributed_clause = "GLOBALLY" if table_type == "global" else "BY (id)"
    conn.execute(
        f"""
        CREATE TABLE test (id UNSIGNED NOT NULL, value TEXT NOT NULL, PRIMARY KEY (id))
        DISTRIBUTED {distributed_clause}
        """
    )
    params = (42, "string")
    conn.execute(
        'INSERT INTO "test" VALUES (%t, %t::text);',
        params,
        binary=False,
    )
    instance.terminate()

    events = AuditFile(instance.audit_flag_value).events()

    audit_policy = take_until_title(events, "audit_policy")
    assert audit_policy is not None
    assert audit_policy["message"] == "audit policy `dml_default` for user `admin` was turned on"
    assert audit_policy["severity"] == "high"
    assert audit_policy["initiator"] == "admin"

    dml = take_until_title(events, "dml")
    assert dml is not None
    assert dml["message"] == "apply `INSERT INTO \"test\" VALUES ($1, $2::text);` with params [42,'string']"
    assert dml["severity"] == "medium"
    assert dml["initiator"] == "admin"
