#!/bin/python3
import fcntl
import json
import argparse
import struct
import re
import shutil
import os
import sys
from os.path import isfile

from typing import Dict, Any


class TransformationError(Exception):
    """Exception raised when a transformation fails."""

    pass


class TransformationRule:
    """Base class for transformation rules."""

    def describe(self) -> str:
        raise NotImplementedError("Subclasses must implement describe method")

    def apply(self, catalog: Dict[str, Any]) -> Dict[str, Any]:
        raise NotImplementedError("Subclasses must implement apply method")


# Helpers

# Helper function to map CedarDB column type to Postgres column type
def mapToPostgresType(t):
    col_type = t['type']

    # Get the value of key 'array', otherwise this returns False
    if t.get('array', False):
        return 'array'

    type_mapping = {
        'bool': 'bool',
        'smallint': 'int2',
        'integer': 'int4',
        'bigint': 'int8',
        'numeric': 'numeric',
        'bignumeric': 'big_numeric',
        'double': 'float8',
        'char1': 'char',
        'char': 'bpchar',
        'text': 'text',
        'bit': 'bit',
        'bitvar': 'varbit',
        'date': 'date',
        'time': 'time',
        'timestamp': 'timestamp',
        'timestamptz': 'timestamptz',
        'interval': 'interval',
        'jsonb': 'jsonb',
        'uuid': 'uuid',
        'bytearray': 'bytea',
        'oid': 'oid',
        'tinyint': 'int1'
    }
    # Look up the type mapping
    if col_type in type_mapping:
        return type_mapping[col_type]

    raise TransformationError(f"Unknown type: {col_type}")

# V6 -> V7
class MoveRemoteStorageToRemoteServer(TransformationRule):
    """Use a previous remote storage as the first remote server ref and add the correct id to all cloud tables"""
    def describe(self) -> str:
        return "move from a single remote storage to a list of remote servers"

    def apply(self, catalog: Dict[str, Any]) -> Dict[str, Any]:
        # Move from the global space into list
        remoteStorage = catalog.get('remoteStorage')

        # Set up the default strcuture of the new catalog and sets all server refs to invalid if not cloud
        catalog['remoteServers'] = []
        for database in catalog['databases']:
            database['value']['remoteServers'] = []
        for table in catalog['tables']:
            if table['value']['storage'] != 'cloud':
                table['value']['storageOptions']['server'] = 4294967295

        # If there is no remote storage, we just make the schema valid again
        if remoteStorage is None:
            return catalog

        # If it has a bucket, we need to move and duplicate it
        del catalog['remoteStorage']
        remoteStorage['location'] = remoteStorage['bucket']
        del remoteStorage['bucket']


        # Add the remote server to all databases
        id = 0
        mapping = {}
        for database in catalog['databases']:
            storage = {"key": id, "value": remoteStorage.copy()}
            storage['value']['database'] = database['key']
            catalog['remoteServers'].append(storage)
            database['value']['remoteServers'] = [{"key": remoteStorage['name'], "value": id}]
            for namespace in database['value']['namespaces']:
                mapping[namespace['value']] = id
            id += 1
        # Update all remote tables
        for table in catalog['tables']:
            table_data = table['value']
            # As defaulted to invalid, we change only cloud tables to the server of the namespace
            if table_data['storage'] == 'column' and table_data['storageOptions']['backend'] == 'cloud':
                table_data['storageOptions']['server'] = mapping[table_data['namespace']]
        return catalog


# V5 -> v6
class QuoteSequences(TransformationRule):
    """We previously had an issue incrementing case-sensitive sequences"""
    def describe(self) -> str:
        return "correctly quote sequence references in column default values"

    def apply(self, catalog: Dict[str, Any]) -> Dict[str, Any]:
        for table in catalog['tables']:
            table_data = table['value']
            for col in table_data['physicalColumns']:
                # match arbitrary Unicode strings separated by a dot
                pattern = r"nextval\('([^']+)\.([^']+)'\)"
                replacement = r"""nextval('"\1"."\2"')"""
                prev = col["default"]
                col["default"] = re.sub(pattern, replacement, col["default"])
                if prev != col["default"]:
                    print(
                        f'\t\tReplacing "{prev}" with "{col["default"]}"'
                    )
        return catalog

# V4 -> V5
class AddIndexOpclassToDeletedIndexes(TransformationRule):
    """We need to add the default opclass to deleted indexes as well."""

    def describe(self) -> str:
        return "add default index opclass to deleted indexes as well"

    def apply(self, catalog: Dict[str, Any]) -> Dict[str, Any]:
        # Make a map which maps table id to a map of column id to column
        table_column_map = {t['key']:
                                {i: c for i, c in enumerate(t['value']['physicalColumns'])}
                            for t in catalog['tables']}
        # Iterate over all the indexes in the catalog
        for index in catalog['indexes']:
            index_data = index['value']

            # Iterate over the columns of each index
            for column in index_data['columns']:
                # We only have two methods in catalog version 2: btree and unchecked
                if index_data['method'] == 'btree':
                    # Get the table column and retrieve its corresponding opclass
                    tableCol = table_column_map[index_data['table']][column['column']['id']]
                    column['opclass'] = mapToPostgresType(tableCol['type']) + '_ops'
                else:
                    # For unchecked indexes, the default opclass for all types is unchecked_ops
                    column['opclass'] = 'unchecked_ops'
        return catalog


# V3 -> V4
class AddIndexOpclass(TransformationRule):
    """We need to add the default opclass to each column based on its type"""

    def describe(self) -> str:
        return "add default index opclass to each index column"

    def apply(self, catalog: Dict[str, Any]) -> Dict[str, Any]:
        # Make a map which maps table id to a map of column id to column
        table_column_map = {t['key']:
                                {i: c for i, c in enumerate(t['value']['physicalColumns'])}
                            for t in catalog['tables'] if t['value']['state'] != 'deleted'}
        # Iterate over all the indexes in the catalog
        for index in catalog['indexes']:
            index_data = index['value']

            if index_data['state'] == 'deleted':
                continue

            # Iterate over the columns of each index
            for column in index_data['columns']:
                # We only have two methods in catalog version 2: btree and unchecked
                if index_data['method'] == 'btree':
                    # Get the table column and retrieve its corresponding opclass
                    tableCol = table_column_map[index_data['table']][column['column']['id']]
                    column['opclass'] = mapToPostgresType(tableCol['type']) + '_ops'
                else:
                    # For unchecked indexes, the default opclass for all types is unchecked_ops
                    column['opclass'] = 'unchecked_ops'
        return catalog


# V2 -> V3
class DropCedarNamespace(TransformationRule):
    """Drop the cedar namespace and adjust all other namespace ids accordingly"""

    def describe(self) -> str:
        return "drop cedar namespace"

    def apply(self, catalog: Dict[str, Any]) -> Dict[str, Any]:

        # Adapt namespace ids in all tables
        for table in catalog["tables"]:
            if int(table["value"]["namespace"]) >= 1:
                table["value"]["namespace"] = (int(table["value"]["namespace"] - 1))

        # Adapt namespace ids in all sequences
        for seq in catalog["sequences"]:
            if int(seq["value"]["namespace"]) >= 1:
                seq["value"]["namespace"] = (int(seq["value"]["namespace"] - 1))

        # Adapt namespace ids in all indexes
        for ind in catalog["indexes"]:
            if int(ind["value"]["namespace"]) >= 1:
                ind["value"]["namespace"] = (int(ind["value"]["namespace"] - 1))

        # Adapt namespace ids in all views
        for view in catalog["views"]:
            if int(view["value"]["namespace"]) >= 1:
                view["value"]["namespace"] = (int(view["value"]["namespace"] - 1))

        # Adapt namespace ids in all foreignKeys
        for fk in catalog["foreignKeys"]:
            if int(fk["value"]["namespace"]) >= 1:
                fk["value"]["namespace"] = (int(fk["value"]["namespace"] - 1))

        # Adapt namespace ids in all functions
        for func in catalog["functions"]:
            if int(func["value"]["namespace"]) >= 1:
                func["value"]["namespace"] = (int(func["value"]["namespace"] - 1))

        # Adapt namespace ids in all databases, dropping references to the cedar namespace
        for db in catalog["databases"]:
            for namespace in db["value"]["namespaces"]:
                if int(namespace["value"]) >= 1:
                    namespace["value"] = (int(namespace["value"] - 1))
            # Drop the cedar namespace
            for i, namespace in enumerate(db["value"]["namespaces"]):
                if namespace["key"] == 'cedar':
                    db["value"]["namespaces"].pop(i)
                    break


        # Adapt all namespace keys
        for namespace in catalog["namespaces"]:
            if int(namespace["key"]) >= 1:
                namespace["key"] = (int(namespace["key"] - 1))

        # Finally drop the cedar namespace
        for i, namespace in enumerate(catalog["namespaces"]):
            if namespace["value"]["name"] == 'cedar':
                catalog["namespaces"].pop(i)
                break

        return catalog



# V1 -> V2
class FixPgExpandarray(TransformationRule):
    """We need to fix the built in _pg_expandarray to work with non 1-indexed arrays."""

    def describe(self) -> str:
        return "patch builtin function _pg_expandarray to work with non 1-indexed arrays"

    def apply(self, catalog: Dict[str, Any]) -> Dict[str, Any]:
        namespace_map = {ns["key"]: ns["value"]["name"] for ns in catalog.get("namespaces", [])}

        for function in catalog['functions']:
            function_data = function['value']
            if function_data['name'] == '_pg_expandarray' and namespace_map[function_data['namespace']] == 'information_catalog':
                function_data['properties']['source'] = 'select arr[s] as x, s::integer - array_lower(arr, 1) + 1 as n from generate_series(array_lower(arr, 1), array_upper(arr, 1), 1) as g(s)'

        return catalog

# V1 -> V2
class EnsureTableObjectNameUniqueness(TransformationRule):
    """In postgres you cannot create tables, indexes, sequences if any other of those with the same name already exists.
    We previously allowed this, but it may lead to issues with pg_dump so we rename conflicting objects"""

    def describe(self) -> str:
        return "rename index and sequence names conflicting with tables"


    def apply(self, catalog: Dict[str, Any]) -> Dict[str, Any]:

        existing_names = {}

        # Save the elements to change for updating the namespaces later
        indexes_to_change = {}
        sequences_to_change = {}

        # Remember all tables
        for table in catalog['tables']:
            table_data = table['value']
            if table_data['state'] == 'deleted':
                continue
            existing_names.setdefault(table_data['namespace'], set()).add(table_data['name'])

        # Remember all views
        for view in catalog['views']:
            view_data = view['value']
            if view_data['state'] == 'deleted':
                continue
            existing_names.setdefault(view_data['namespace'], set()).add(view_data['name'])


        # Rename conflicting indexes
        for index in catalog['indexes']:
            index_data = index['value']
            # Deleted indexes are invisible anyway, we don't care for naming conflicts
            if index_data['state'] == 'deleted':
                continue

            namespace_id = index_data['namespace']
            potential_conflicts = existing_names[namespace_id]
            if index_data['name'] in potential_conflicts:
                new_name = self.generate_unique_name(index_data['name'], potential_conflicts)
                print(f'\t\tRenaming conflicting index: "{index_data['name']}" to: "{new_name}"')
                index_data['name'] = new_name
                indexes_to_change[index['key']] = new_name
                potential_conflicts.add(new_name)


        # Rename conflicting sequences
        for sequence in catalog['sequences']:
            sequence_data = sequence['value']
            # Deleted indexes are invisible anyway, we don't care for naming conflicts
            if sequence_data['state'] == 'deleted':
                continue

            namespace_id = sequence_data['namespace']
            potential_conflicts = existing_names[namespace_id]
            if sequence_data['name'] in potential_conflicts:
                new_name = self.generate_unique_name(sequence_data['name'], potential_conflicts)
                print(f'\t\tRenaming conflicting sequence: "{sequence_data['name']}" to: "{new_name}"')
                sequence_data['name'] = new_name
                sequences_to_change[sequence['key']] = new_name
                potential_conflicts.add(new_name)

        # Update the references in the namespace. Annoyingly enough keys and values are flipped here
        for namespace in catalog["namespaces"]:
            namespace_data = namespace["value"]
            for index in namespace_data["indexes"]:
                if index["value"] in indexes_to_change:
                    index["key"] = indexes_to_change[index["value"]]

            for sequence in namespace_data["sequences"]:
                if sequence["value"] in sequences_to_change:
                    sequence["key"] = sequences_to_change[sequence["value"]]


        return catalog

    def generate_unique_name(self, base_name: str, existing_names: set) -> str:
        """Generate a unique name by appending an incrementing number."""
        counter = 1
        new_name = f"{base_name}{counter}"
        while new_name in existing_names:
            counter += 1
            new_name = f"{base_name}{counter}"
        return new_name

# V1 -> V2
class SequenceNameInNextval(TransformationRule):
    """To increment sequences, we used their internal ids. This is an issue when we change the id format.
    Reference the sequence by name instead so the catalog stays stable when we change sequence references."""
    def describe(self) -> str:
        return "Store sequence name instead of id in table default values"

    def apply(self, catalog: Dict[str, Any]) -> Dict[str, Any]:

        # Build a map of all sequences: key -> content
        sequence_map = {seq["key"]: seq["value"] for seq in catalog.get("sequences", [])}

        # Build a map of all namespaces: key -> name
        namespace_map = {ns["key"]: ns["value"]["name"] for ns in catalog.get("namespaces", [])}

        for table in catalog['tables']:
            table_data = table['value']
            for col in table_data['physicalColumns']:
                # Is the default of the form nextval(X)?
                default_value = col.get("default", "")
                match = re.match(r"nextval\((\d+)\)", default_value)
                if match:
                    seq_key = int(match.group(1))
                    seq_data = sequence_map.get(seq_key)
                    if seq_data:
                        namespace_name = namespace_map.get(seq_data["namespace"])
                        seq_name = seq_data["name"]
                        col["default"] = f"nextval('{namespace_name}.{seq_name}')"  # Replace default value
                        print(
                            f'\t\tReplacing "nextval({seq_key})" with "{col["default"]}" in column "{col["name"]}" of table "{namespace_name}.{table_data["name"]}" '
                        )
        return catalog

# VO -> V1
class AddSequenceStartParam(TransformationRule):
    """We didn't store the start value of a sequence which is relevant when resetting
    the sequence. Set it to the minValue as sane default"""

    def describe(self) -> str:
        return "add start parameters to sequences"

    def apply(self, catalog: Dict[str, Any]) -> Dict[str, Any]:
        for seq in catalog["sequences"]:
            if "minValue" not in seq["value"]:
                raise TransformationError("sequence didn't have a minimum value")
            seq["value"]["start"] = seq["value"]["minValue"]
            print(
                f'\t\tAdded start value : "{seq["value"]["start"]}" to sequence: "{seq["value"]["name"]}"'
            )
        return catalog


# VO -> V1
class DeleteSequencesOfDeletedTables(TransformationRule):
    """Since we didn't store a reference between tables and sequences between restarts,
    we didn't correctly mark sequences of deleted tables as deleted after restart"""

    def describe(self) -> str:
        return "mark sequences depending on deleted tables as deleted"

    def apply(self, catalog: Dict[str, Any]) -> Dict[str, Any]:

        # Find the sequences that should be deleted
        deleted_sequences = set()

        for table in catalog["tables"]:
            if table["value"]["state"] != "deleted":
                continue

            # Find out if the table references a sequence (i.e., it's default value is "nextval(seq_id)"
            for col in table["value"]["physicalColumns"]:
                if "default" not in col:
                    continue

                match = re.match(r"nextval\((\d+)\)", col["default"])
                if match:
                    deleted_sequences.add(int(match.group(1)))
                    print(
                        f'\t\tMarked sequence with id : "{match.group(1)}" of deleted table: "{table["value"]["name"]}" as deleted'
                    )

        # Iterate over all sequences and update those that should be deleted
        for seq in catalog["sequences"]:
            if seq["key"] in deleted_sequences:
                seq["value"]["state"] = "deleted"

        return catalog


# VO -> V1
class CreateSequenceMappings(TransformationRule):
    """Create all references between tables, sequences, and namespaces.
    Also drops explicitly created sequences as we didn't store them correctly in the catalog which caused them
    disappear after restart. This transformation mimics that behavior, but makes sure everything is in a consistent state.
    """

    def describe(self) -> str:
        return "create mappings between tables, sequences, and namespaces"

    def apply(self, catalog: Dict[str, Any]) -> Dict[str, Any]:
        # Find the sequences that are referenced by tables
        sequence_to_table = dict()
        # We also store the namespace the sequence should be in
        sequence_to_namespace = dict()
        namespace_to_sequence = dict()

        for table in catalog["tables"]:
            # All tables require a sequences entry now
            table["value"].setdefault("sequences", [])
            # Find out if the table *should* reference a sequence (i.e., it's default value is "nextval(seq_id)"
            for col in table["value"]["physicalColumns"]:
                if "default" not in col:
                    continue

                match = re.match(r"nextval\((\d+)\)", col["default"])
                if match:
                    seq_id = int(match.group(1))
                    # Directly link the sequence to the table itself
                    table["value"]["sequences"].append(seq_id)
                    sequence_to_table[seq_id] = table["key"]
                    # Also remember the namespace so we can update it later
                    ns_id = table["value"]["namespace"]
                    sequence_to_namespace[seq_id] = ns_id
                    namespace_to_sequence.setdefault(ns_id, []).append(seq_id)

        # Iterate over all sequences and update backreferences where applicable
        sequenceid_to_name = dict()
        for seq in catalog["sequences"]:
            seq_id = seq["key"]
            sequenceid_to_name[seq_id] = seq["value"]["name"]
            if seq_id in sequence_to_table.keys():
                seq["value"]["table"] = sequence_to_table[seq_id]
                seq["value"]["namespace"] = sequence_to_namespace[seq_id]
            else:
                # It was an explicitly created sequence which we can't recover as it was never
                # correctly added to the corresponding namespace and may even be a duplicate
                # Mark it as deleted.
                seq["value"]["state"] = "deleted"
                # Insert a dummy namespace so CedarDB doesn't complain
                seq["value"]["namespace"] = 3  # "public" namespace
                namespace_to_sequence.setdefault(3, []).append(seq_id)

        # Finally add references from namespace to sequences, which was also missing
        for namespace in catalog["namespaces"]:
            ns_id = namespace["key"]
            namespace["value"].setdefault("sequences", [])

            if ns_id in namespace_to_sequence:
                for seq_id in namespace_to_sequence[ns_id]:
                    seq = {"key": sequenceid_to_name[seq_id], "value": seq_id}
                    namespace["value"]["sequences"].append(seq)

        return catalog

class AddLogicalIds(TransformationRule):
    """Create logical ids to all catalog objects, such as tables, namespaces, functions, etc.
    These will be their current offset in the vector to start with
    """

    def describe(self) -> str:
        return "add logical ids to all catalog objects"

    def apply(self, catalog: Dict[str, Any]) -> Dict[str, Any]:
        keys = ["roles", "databases", "namespaces", "tables", "indexes", "views", "foreignKeys", "functions", "sequences", "subscriptions", "remoteServers"]

        for key in keys:
            for catalog_obj in catalog[key]:
                keyVal = catalog_obj["key"]
                catalog_obj["value"]["logicalId"] = keyVal

        return catalog

class AddInstanceReferences(TransformationRule):
    """
    Add instance references to tables and indexes to help identify their on-disk relations and indexes.
    Identical to the logical and physical id at time of upgrade
    """

    def describe(self) -> str:
        return "add instance references to tables and indexes"

    def apply(self, catalog: Dict[str, Any]) -> Dict[str, Any]:
        for table in catalog["tables"]:
            table["value"]["instanceRef"] = {"id": table["key"], "type": "table"}

        for index in catalog["indexes"]:
            index["value"]["instanceRef"] = index["key"]


        return catalog

class AddReferencedTableObjectsToView(TransformationRule):
    """
    Add a field to keep track of referenced table objects to views.
    Unfortunately we cannot easily regain the information without building a SQL parser,
    so we just initialize them emptily for existing views.
    """

    def describe(self) -> str:
        return "add field for references table objects to views"

    def apply(self, catalog: Dict[str, Any]) -> Dict[str, Any]:
        for catalog_obj in catalog["views"]:
            catalog_obj["value"]["referencedTableObjects"] = []
        return catalog

class AddMissingTypeAndPartitionInfoToIndexes(TransformationRule):
    """
    Add information about the storage column types and table-level collate of index columns to
    the index to avoid accessing the table during recovery. Also add information about partitioning directly
    to the index.
    """

    def describe(self) -> str:
        return "add partition, collate, and type information to indexes"

    def apply(self, catalog: Dict[str, Any]) -> Dict[str, Any]:
        table_dict = dict()

        # Collect information about the partition key and physical columns from tables
        for table in catalog["tables"]:
            table_val = dict()
            table_val["tablePartitionKey"] = table["value"]["partitionKey"]
            table_val["columns"] = dict()

            for key, column in enumerate(table["value"]["physicalColumns"]):
                table_val["columns"][key] = column
            table_dict[table["key"]] = table_val


        for index in catalog["indexes"]:
            table = table_dict[index["value"]["table"]]
            # Copy over the partition key from the table
            index["value"]["tablePartitionKey"] = table["tablePartitionKey"]

            # Copy over per-column type and table-side collate
            for column in index["value"]["columns"]:
                table_column = table["columns"][column["column"]["id"]]
                column["type"] = table_column["type"]
                column["tableCollate"] = table_column["collate"]


        return catalog

class DropViewStorageOptions(TransformationRule):
    """
    We stored storage information from views, even though that was meaningless.
    We have moved that information to tables only, so clean up in the views
    """

    def describe(self) -> str:
        return "Drop unused storage options from views"

    def apply(self, catalog: Dict[str, Any]) -> Dict[str, Any]:
        for catalog_obj in catalog["views"]:
            del catalog_obj["value"]["storage"]
            if "storageOptions" in catalog_obj["value"]:
                del catalog_obj["value"]["storageOptions"]
        return catalog

class AddDatabaseToSubscription(TransformationRule):
    """
    We did not keep track of the database to which a subscription belongs, but have this information through other paths
    """

    def describe(self) -> str:
        return "Add database reference to subscriptions"

    def apply(self, catalog: Dict[str, Any]) -> Dict[str, Any]:
        subscription_db_lookup = dict()

        # Build a lookup from subscription id to db id
        for database in catalog["databases"]:
            for subscription in database["value"]["subscriptions"]:
                subscription_db_lookup[subscription["key"]] = database["key"]


        # For each subscription, look up the containing database and store the reference
        for catalog_obj in catalog["subscriptions"]:
            catalog_obj["value"]["database"] = subscription_db_lookup[catalog_obj["key"]]

        return catalog

class DropRedundantRelationshipInformation(TransformationRule):
    """
    We kept track of containment of catalog objects in two ways, one with a list of members and
    one time with references to the parent. To avoid inconsistencies during changes, we only want to
    track this with parent references and build the other information at runtime.
    """

    def describe(self) -> str:
        return "Drop redundant relationship information"

    def apply(self, catalog: Dict[str, Any]) -> Dict[str, Any]:
        subscription_db_lookup = dict()

        # First, drop dependencies and reverse dependencies from all catalog objects
        keys = ["roles", "databases", "namespaces", "tables", "indexes", "views", "foreignKeys", "functions", "sequences", "subscriptions", "remoteServers"]

        for key in keys:
            for catalog_obj in catalog[key]:
                if "dependencies" in catalog_obj["value"]:
                    del catalog_obj["value"]["dependencies"]

                if "reverseDependencies" in catalog_obj["value"]:
                    del catalog_obj["value"]["reverseDependencies"]

        # Drop namespaced, subscriptions and remote servers from databases
        for database in catalog["databases"]:
            del database["value"]["namespaces"]
            del database["value"]["subscriptions"]
            del database["value"]["remoteServers"]

        # Drop table objects, indexes, functions and sequences from namespaces
        for namespace in catalog["namespaces"]:
            del namespace["value"]["tableObjects"]
            del namespace["value"]["indexes"]
            del namespace["value"]["functions"]
            del namespace["value"]["sequences"]

        # Drop index and foreign key information from tables
        for table in catalog["tables"]:
            del table["value"]["primaryKeys"]
            del table["value"]["indexes"]
            del table["value"]["foreignKeys"]
            del table["value"]["incomingForeignKeys"]

        return catalog

class DropUncheckedIndexes(TransformationRule):
    """
    Unchecked indexes were a benchmarking hack and should never have existed in the first place.
    As they don't do anything, we can just drop them everywhere
    """

    def describe(self) -> str:
        return "remove unchecked indexes"

    def apply(self, catalog: Dict[str, Any]) -> Dict[str, Any]:
        dropped_indexes = set()

        # Drop all indexes with method unchecked
        for index in catalog["indexes"]:
            if index["value"]["method"] == "unchecked":
                dropped_indexes.add(index["key"])
                del index

        # Drop all foreign keys referencing an unchecked index
        for fk in catalog["foreignKeys"]:
            if fk["value"]["referencedIndex"] in dropped_indexes or fk["value"]["referencingIndex"] in dropped_indexes:
                del fk
        return catalog

class PatchRoleOids(TransformationRule):
    """In order to be fully compatible with PostgreSQL, we need to create a set of roles with a fixed OID < 16384
    As this might conflict with current role ids, make sure to increase all user-defined ones to be above 16384 and set the postgres role to the expected value 10
    """

    def describe(self) -> str:
        return "make sure role ids are postgres compatible"
    def apply(self, catalog: Dict[str, Any]) -> Dict[str, Any]:
        role_id_mapping = dict[int, int]()
        for role in catalog["roles"]:
            if role["value"]["name"] == "postgres" and role["value"]["state"] == "visible":
                role_id_mapping[role["value"]["logicalId"]] = 10
            elif role["value"]["name"] != "":
                role_id_mapping[role["value"]["logicalId"]] = int(role["value"]["logicalId"]) + 16384

        # Handle public role as well, shouldn't exist anywhere but its better to be safe
        role_id_mapping[0] = 0

        # Patch owners in all catalog objects, skipping remoteServers as they currently do not have an owner set
        keys = ["databases", "namespaces", "tables", "indexes", "views", "foreignKeys", "functions", "sequences", "subscriptions"]
        # Patch owner in all non-role options
        for key in keys:
            for catalog_obj in catalog[key]:
                owner_val = role_id_mapping[catalog_obj["value"]["owner"]]
                catalog_obj["value"]["owner"] = owner_val

        # Patch membership and adminOf info in roles
        for role in catalog["roles"]:
            new_memberships = []
            for member in role["value"]["memberOf"]:
                new_memberships.append(role_id_mapping[member])
            role["value"]["memberOf"] = new_memberships


            new_admin_of = []
            for member in role["value"]["adminOf"]:
                new_admin_of.append(role_id_mapping[member])
            role["value"]["adminOf"] = new_admin_of

        # Finally patch the logicalIds of roles
        for role in catalog["roles"]:
            role["value"]["logicalId"] = role_id_mapping[role["value"]["logicalId"]]

        return catalog

class AddBuiltinPostgresRoles(TransformationRule):
    """We did not have all built-in postgresql roles, so add the missing roles with their respective oids
    """

    def describe(self) -> str:
        return "add postgresql's pre-defined roles"
    def apply(self, catalog: Dict[str, Any]) -> Dict[str, Any]:
        role_names = set[str]()
        max_key = 0
        for role in catalog["roles"]:
            role_names.add(role["value"]["name"])
            max_key = max(max_key, role["key"])

        predefined_roles = [
            [6171, "pg_database_owner"],
            [6181, "pg_read_all_data"],
            [6182, "pg_write_all_data"],
            [3373, "pg_monitor"],
            [3374, "pg_read_all_settings"],
            [3375, "pg_read_all_stats"],
            [3377, "pg_stat_scan_tables"],
            [4569, "pg_read_server_files"],
            [4570, "pg_write_server_files"],
            [4571, "pg_execute_server_program"],
            [4200, "pg_signal_backend"],
            [4544, "pg_checkpoint"],
            [6337, "pg_maintain"],
            [4550, "pg_use_reserved_connections"],
            [6304, "pg_create_subscription"],
        ]

        role_key = max_key + 1
        # Patch owner in all non-role options
        for oid, name in predefined_roles:
            if name in role_names:
                raise Exception(f"role with name '{name}' was already defined by the user, please reach out for instructions")

            # Create the new role to be inserted
            new_role = {
                "key": role_key,
                "value": {
                    "owner": 4294967295,
                    "name": name,
                    "state": "visible",
                    "logicalId": oid,
                    "createdAtReference": False,
                    "deletedAtReference": False,
                    "dependencies": [],
                    "reverseDependencies": [],
                    "permissions": [],
                    "connectionLimit": 0,
                    "createDb": False,
                    "createRole": False,
                    "superUser": False,
                    "replication": False,
                    "inherit": True,
                    "login": False,
                    "passwordState": "nopassword",
                    "password": "",
                    "passwordValidUntil": 18446744073709551615,
                    "memberOf": [],
                    "adminOf": []
                }
            }
            if name == "pg_monitor":
                new_role["value"]["memberOf"].append(3374)
                new_role["value"]["memberOf"].append(3375)
                new_role["value"]["memberOf"].append(3377)
            catalog["roles"].append(new_role)
            role_key += 1

        return catalog

class AddOwnerToRemoteServers(TransformationRule):
    """We didn't have owners for remote servers, so use postgres as default"""
    def describe(self) -> str:
        return "add owner to remote servers who were previously ownerless"


    def apply(self, catalog: Dict[str, Any]) -> Dict[str, Any]:
        # Find the postgres user
        pgUser = None
        for role in catalog['roles']:
            if role['value']['name'] == 'postgres':
                pgUser = role['value']['logicalId']
                break

        if pgUser is None:
            raise TransformationError("Could not find user 'postgres'")

        for remoteServer in catalog['remoteServers']:
            if remoteServer['value']['owner'] == 4294967295: # <=> UINT32_MAX <=> unset
                remoteServer['value']['owner'] = pgUser
        return catalog


class PatchLogicalColumns(TransformationRule):
    """Move column metadata from physical to logical columns"""
    def describe(self) -> str:
        return "Move column metadata from physical to logical columns"
    def apply(self, catalog: Dict[str, Any]) -> Dict[str, Any]:
        # Patch the columns in all table-like objects with column definitions
        keys = ["tables", "views"]
        for key in keys:
            for obj in catalog[key]:
                table_object = obj["value"]
                physical_columns = table_object["physicalColumns"]
                logical_columns = table_object["columns"]
                for column in logical_columns:
                    physical_column_id = column["phsicalColumns"][0]["id"]
                    physical_column = physical_columns[physical_column_id]
                    column["collate"] = physical_column["collate"]
                    del physical_column["collate"]
                    column["default"] = physical_column["default"]
                    del physical_column["default"]
        return catalog


class AddAdditionalRoleMembershipInformation(TransformationRule):
    """We tracked to little information about memberships, so add them with the PostgreSQL default values they would have anyway
    """

    def describe(self) -> str:
        return "add postgresql's membership options to roles"
    def apply(self, catalog: Dict[str, Any]) -> Dict[str, Any]:
        for role in catalog["roles"]:
            for member in role["value"]["memberOf"]:
                member_id = member
                member = {
                    "group": member_id,
                    "grantedBy": 10, # Granted by postgres
                    "admin": False,
                    "inherit": True,
                    "set": True,
                }
            del role["value"]["adminOf"] # Drop now unused adminOf field
        return catalog

# Define transformation rules for each catalog version upgrade
# The list is ordered (i.e., it's allowed that a later transformation depends on an earlier one)
TRANSFORMATION_RULES = {
    0: [
        AddSequenceStartParam(),
        DeleteSequencesOfDeletedTables(),
        CreateSequenceMappings(),
    ],
    1: [
        EnsureTableObjectNameUniqueness(),
        SequenceNameInNextval(),
        FixPgExpandarray()
    ],
    2: [
        DropCedarNamespace()
    ],
    3: [
        AddIndexOpclass()
    ],
    4: [
        AddIndexOpclassToDeletedIndexes()
    ],
    5: [
        QuoteSequences()
    ],
    6: [
        MoveRemoteStorageToRemoteServer()
    ],
    7: [
        AddLogicalIds(),
        AddInstanceReferences(),
        AddReferencedTableObjectsToView(),
        AddMissingTypeAndPartitionInfoToIndexes(),
        DropViewStorageOptions(),
        AddDatabaseToSubscription(),
        DropRedundantRelationshipInformation(),
        DropUncheckedIndexes(),
    ],
    8: [
        PatchRoleOids(),
        AddBuiltinPostgresRoles()
    ],
    9: [
        AddOwnerToRemoteServers()
    ],
    10: [
        PatchLogicalColumns()
    ],
    11: [
        AddAdditionalRoleMembershipInformation()
    ]
}

HEADER_FORMAT = "8s I I Q Q I"  # 8-byte magic, 4-byte pointer size, 4-byte bin fmt version, 8-byte GSN, 8-byte epoch, 4-byte catalog version
HEADER_SIZE = struct.calcsize(HEADER_FORMAT)


class CatalogUpgrader:
    def __init__(self, catalog: Dict[str, Any], version: int):
        self.catalog = catalog
        self.version = version

    def upgrade(self) -> bool:
        """Applies transformations until the latest version is reached."""
        while self.version in TRANSFORMATION_RULES:
            print(
                f"Upgrading catalog from version {self.version} to {self.version + 1}..."
            )
            try:
                for transformation in TRANSFORMATION_RULES[self.version]:
                    print(f"\tApplying transformation: {transformation.describe()}")
                    self.catalog = transformation.apply(self.catalog)
                self.version += 1
            except TransformationError as e:
                print(f"Upgrade failed at version {self.version}: {e}")
                return False
        return True

    def get_catalog(self) -> Dict[str, Any]:
        return self.catalog


def read_catalog(file_path: str):
    with open(file_path, "rb") as f:
        header_data = f.read(HEADER_SIZE)
        magic, ptr_size, bin_fmt_version, gsn, epoch, catalog_version = struct.unpack(
            HEADER_FORMAT, header_data
        )
        json_data = f.read().decode("utf-8")
        catalog = json.loads(json_data)
        return (magic, ptr_size, bin_fmt_version, gsn, epoch, catalog_version, catalog)


def write_catalog(file_path: str, header, catalog: Dict[str, Any]):
    with open(file_path, "wb") as f:
        header_data = struct.pack(HEADER_FORMAT, *header[:-1], header[-1])
        f.write(header_data)
        f.write(json.dumps(catalog, indent=4).encode("utf-8"))

def was_shutdown_cleanly(wal_dir):
    # wal files have the format db.wal.N
    wal_files = [f for f in os.listdir(wal_dir) if re.match(r"db\.wal\.\d+", f)]

    if len(wal_files) != 1:
        return False  # More than one WAL file exists, indicating non-truncated log -> unclean shutdown

    wal_file_path = os.path.join(wal_dir, wal_files[0])

    try:
        with open(wal_file_path, "rb") as f:
            content = f.read()
            return len(content) == 0  # The file may have a reported size, but actual content must be empty
    except Exception as e:
        print(f"Error reading WAL file: {e}")
        return False

if __name__ == "__main__":
    parser = argparse.ArgumentParser(
        description="Upgrade a binary-stored JSON catalog to the latest version."
    )
    parser.add_argument("db_dir", nargs="?", help="Path to the database directory.")
    args = parser.parse_args()

    # If no argument is given, ask the user if they want to use the default path
    defaultPath = os.path.expanduser("~/.cedardb/db")
    if not args.db_dir:
        response = input(f"No path provided. Use default path '{defaultPath}'? (y/n): ").strip().lower()
        if response == 'y':
            args.db_dir = defaultPath
        else:
            print("No database directory provided. Exiting.")
            sys.exit(1)

    # Check if the db exists
    catalog_file = os.path.join(args.db_dir, "db")
    if not os.path.exists(catalog_file):
        print(f"unable to open database catalog file \"{catalog_file}\": File not found")
        sys.exit(1)

    # Get a lock on the database to ensure it's not currently running
    lock_file = os.path.join(args.db_dir, ".db.lock")
    lock = open(lock_file, "w+b")
    try:
        fcntl.lockf(lock, fcntl.LOCK_EX | fcntl.LOCK_NB)
    except OSError:
        print("Could not acquire the database lock on existing database. Please shut down CedarDB and retry.")
        sys.exit(1)

    # Read the catalog
    magic, ptr_size, bin_fmt_version, gsn, epoch, catalog_version, catalog_data = (
        read_catalog(catalog_file)
    )

    # Check whether the catalog needs to be upgraded at all
    # The latest version + 1 for which we have a transformation is the target version of this script
    # This way of getting the latest version is fine, since dicts are guaranteed to be ordered since Python 3.7
    if catalog_version == list(TRANSFORMATION_RULES.keys())[-1] + 1:
        # We're at the newest version already, unlock the database and exit
        fcntl.lockf(lock, fcntl.LOCK_UN)
        sys.exit(0)

    # We only upgrade catalog files of databases that have been shut down cleanly to avoid any hard to track down
    # bugs where catalog upgrades might interact with recovery
    if not was_shutdown_cleanly(args.db_dir):
        print("Database is not in a clean state. Please re-start CedarDB, shut it down normally and retry.")
        # Unlock the database
        fcntl.lockf(lock, fcntl.LOCK_UN)
        sys.exit(1)

    # Backup the current catalog file
    backup_file = catalog_file + ".backup"
    shutil.copy(catalog_file, backup_file)

    # Upgrade the catalog
    upgrader = CatalogUpgrader(catalog_data, catalog_version)
    success = upgrader.upgrade()

    if not success:
        print(
            "Upgrade failed. Original catalog restored. Please contact support and continue using the old CedarDB version in the meantime."
        )
        # Unlock the database
        fcntl.lockf(lock, fcntl.LOCK_UN)
        # No need to restore the catalog file from backup as we haven't written anything yet
        sys.exit(1)

    # Write the new catalog inplace
    new_header = (magic, ptr_size, bin_fmt_version, gsn, epoch, upgrader.version)
    write_catalog(catalog_file, new_header, upgrader.get_catalog())
    print(
        f"Schema upgraded successfully to version {upgrader.version}. Output saved to {catalog_file}"
    )
    # Unlock the database
    fcntl.lockf(lock, fcntl.LOCK_UN)
