from testcontainers.postgres import PostgresContainer


class PostgresContainerHelper:
    def __init__(self, container: PostgresContainer):
        self.container = container
        self.snapshot_name: str | None = "migrated_template"
        self.db_name = container.dbname
        self.user = container.username

    async def snapshot(self, snapshot_name: str = "migrated_template") -> None:
        """Create a template database snapshot"""
        if snapshot_name:
            self.snapshot_name = snapshot_name

        commands = [
            # Allow dropping the template database if it exists
            f"UPDATE pg_database SET datistemplate = FALSE WHERE datname = '{snapshot_name}'",
            f"DROP DATABASE IF EXISTS {snapshot_name}",
            # Create new template from current database
            f"CREATE DATABASE {snapshot_name} WITH TEMPLATE {self.db_name} OWNER {self.user}",
            # Mark as template
            f"ALTER DATABASE {snapshot_name} WITH is_template = TRUE",
        ]

        print(f"Creating snapshot: {snapshot_name}")

        await self._execute_system_commands(commands)

        print(f"Created snapshot: {snapshot_name}")

    async def restore(self, snapshot_name: str | None = None) -> None:
        """Restore database from snapshot"""
        if not snapshot_name:
            snapshot_name = self.snapshot_name

        if not snapshot_name:
            raise ValueError("No snapshot name provided and no default exists")

        commands = [
            f"DROP DATABASE {self.db_name} WITH (FORCE)",
            f"CREATE DATABASE {self.db_name} WITH TEMPLATE {snapshot_name} OWNER {self.user}",
        ]

        print(f"Restoring snapshot: {snapshot_name}")

        await self._execute_system_commands(commands)

        print(f"Restored snapshot: {snapshot_name}")

    async def _execute_system_commands(self, commands: list[str]) -> None:
        for cmd in commands:
            (exit_code, output) = self.container.exec(
                f"psql -U {self.container.username} -d {self.container.dbname} -c '{cmd}'"
            )

            if exit_code != 0:
                print(
                    f"\nExit code: {exit_code} for Command: {cmd} \nOutput: {output.decode('utf-8')}"
                )