from supabase_mcp.exceptions import OperationNotAllowedError
from supabase_mcp.logger import logger
from supabase_mcp.services.database.migration_manager import MigrationManager
from supabase_mcp.services.database.postgres_client import PostgresClient, QueryResult
from supabase_mcp.services.database.sql.loader import SQLLoader
from supabase_mcp.services.database.sql.models import QueryValidationResults
from supabase_mcp.services.database.sql.validator import SQLValidator
from supabase_mcp.services.safety.models import ClientType, SafetyMode
from supabase_mcp.services.safety.safety_manager import SafetyManager


class QueryManager:
    """
    Manages SQL query execution with validation and migration handling.

    This class is responsible for:
    1. Validating SQL queries for safety
    2. Executing queries through the database client
    3. Managing migrations for queries that require them
    4. Loading SQL queries from files

    It acts as a central point for all SQL operations, ensuring consistent
    validation and execution patterns.
    """

    def __init__(
        self,
        postgres_client: PostgresClient,
        safety_manager: SafetyManager,
        sql_validator: SQLValidator | None = None,
        migration_manager: MigrationManager | None = None,
        sql_loader: SQLLoader | None = None,
    ):
        """
        Initialize the QueryManager.

        Args:
            postgres_client: The database client to use for executing queries
            safety_manager: The safety manager to use for validating operations
            sql_validator: Optional SQL validator to use
            migration_manager: Optional migration manager to use
            sql_loader: Optional SQL loader to use
        """
        self.db_client = postgres_client
        self.safety_manager = safety_manager
        self.validator = sql_validator or SQLValidator()
        self.sql_loader = sql_loader or SQLLoader()
        self.migration_manager = migration_manager or MigrationManager(loader=self.sql_loader)

    def check_readonly(self) -> bool:
        """Returns true if current safety mode is SAFE."""
        result = self.safety_manager.get_safety_mode(ClientType.DATABASE) == SafetyMode.SAFE
        logger.debug(f"Check readonly result: {result}")
        return result

    async def handle_query(self, query: str, has_confirmation: bool = False, migration_name: str = "") -> QueryResult:
        """
        Handle a SQL query with validation and potential migration. Uses migration name, if provided.

        This method:
        1. Validates the query for safety
        2. Checks if the query requires migration
        3. Handles migration if needed
        4. Executes the query

        Args:
            query: SQL query to execute
            params: Query parameters
            has_confirmation: Whether the operation has been confirmed by the user

        Returns:
            QueryResult: The result of the query execution

        Raises:
            OperationNotAllowedError: If the query is not allowed in the current safety mode
            ConfirmationRequiredError: If the query requires confirmation and has_confirmation is False
        """
        # 1. Run through the validator
        validated_query = self.validator.validate_query(query)

        # 2. Ensure execution is allowed
        self.safety_manager.validate_operation(ClientType.DATABASE, validated_query, has_confirmation)
        logger.debug(f"Operation with risk level {validated_query.highest_risk_level} validated successfully")

        # 3. Handle migration if needed
        await self.handle_migration(validated_query, query, migration_name)

        # 4. Execute the query
        return await self.handle_query_execution(validated_query)

    async def handle_query_execution(self, validated_query: QueryValidationResults) -> QueryResult:
        """
        Handle query execution with validation and potential migration.

        This method:
        1. Checks the readonly mode
        2. Executes the query
        3. Returns the result

        Args:
            validated_query: The validation result
            query: The original query

        Returns:
            QueryResult: The result of the query execution
        """
        readonly = self.check_readonly()
        result = await self.db_client.execute_query(validated_query, readonly)
        logger.debug(f"Query result: {result}")
        return result

    async def handle_migration(
        self, validation_result: QueryValidationResults, original_query: str, migration_name: str = ""
    ) -> None:
        """
        Handle migration for a query that requires it.

        Args:
            validation_result: The validation result
            query: The original query
            migration_name: Migration name to use, if provided
        """
        # 1. Check if migration is needed
        if not validation_result.needs_migration():
            logger.debug("No migration needed for this query")
            return

        # 2. Prepare migration query
        migration_query, name = self.migration_manager.prepare_migration_query(
            validation_result, original_query, migration_name
        )
        logger.debug("Migration query prepared")

        # 3. Execute migration query
        try:
            # First, ensure the migration schema exists
            await self.init_migration_schema()

            # Then execute the migration query
            migration_validation = self.validator.validate_query(migration_query)
            await self.db_client.execute_query(migration_validation, readonly=False)
            logger.info(f"Migration '{name}' executed successfully")
        except Exception as e:
            logger.debug(f"Migration failure details: {str(e)}")
            # We don't want to fail the main query if migration fails
            # Just log the error and continue
            logger.warning(f"Failed to record migration '{name}': {e}")

    async def init_migration_schema(self) -> None:
        """Initialize the migrations schema and table if they don't exist."""
        try:
            # Get the initialization query
            init_query = self.sql_loader.get_init_migrations_query()

            # Validate and execute it
            init_validation = self.validator.validate_query(init_query)
            await self.db_client.execute_query(init_validation, readonly=False)
            logger.debug("Migrations schema initialized successfully")
        except Exception as e:
            logger.warning(f"Failed to initialize migrations schema: {e}")

    async def handle_confirmation(self, confirmation_id: str) -> QueryResult:
        """
        Handle a confirmed operation using its confirmation ID.

        This method retrieves the stored operation and passes it to handle_query.

        Args:
            confirmation_id: The unique ID of the confirmation to process

        Returns:
            QueryResult: The result of the query execution
        """
        # Get the stored operation
        operation = self.safety_manager.get_stored_operation(confirmation_id)
        if not operation:
            raise OperationNotAllowedError(f"Invalid or expired confirmation ID: {confirmation_id}")

        # Get the query from the operation
        query = operation.original_query
        logger.debug(f"Processing confirmed operation with ID {confirmation_id}")

        # Call handle_query with the query and has_confirmation=True
        return await self.handle_query(query, has_confirmation=True)

    def get_schemas_query(self) -> str:
        """Get a query to list all schemas."""
        return self.sql_loader.get_schemas_query()

    def get_tables_query(self, schema_name: str) -> str:
        """Get a query to list all tables in a schema."""
        return self.sql_loader.get_tables_query(schema_name)

    def get_table_schema_query(self, schema_name: str, table: str) -> str:
        """Get a query to get the schema of a table."""
        return self.sql_loader.get_table_schema_query(schema_name, table)

    def get_migrations_query(
        self, limit: int = 50, offset: int = 0, name_pattern: str = "", include_full_queries: bool = False
    ) -> str:
        """Get a query to list migrations."""
        return self.sql_loader.get_migrations_query(
            limit=limit, offset=offset, name_pattern=name_pattern, include_full_queries=include_full_queries
        )
