"""Lambda snowflake_query function module.

This lambda function is used to query Snowflake database using a raw SQL query
and return the results.
"""

import asyncio
import json
from typing import Literal, Optional
import brotli
import zstandard
from pydantic import BaseModel, ValidationError, Field

from common.src import logger
from common.src.enums import CompressionFormats, PlatformName
from . import config
from .auth import get_snowflake_private_key
from .connectors.snowflake_db import Client, staged
from .enums import OutputFormat, SFClientSettings

logger = logger.new_logger(__name__)


class QueryEvent(BaseModel):
    """Pydantic model for validating incoming Lambda event payloads."""

    sql: str = Field(..., description="SQL query to execute")
    staged: Optional[bool | Literal["S3", "s3"]] = Field(
        False,
        description=(
            # TODO: Update this description to reflect the actual behavior
            """
        Whether to run the query using Snowflake staging. This means
        the results of the query won't be returned directly, but will be
        first staged in Snowflake and then retrieved. This is useful for
        large result sets or when the query is expected to take a long time,
        as it leverages parallel processing in Snowflake and the results are
        returned compressed in ZSTD format.

        This mode won't allow SELECT * queries (columns must be explicitly defined).

        This mode is not recommended for simple queries or small result sets, as
        it has overhead and may be slower than running the query directly.

        If set to true, the following parameters have no effect:
            - output_format
            - compression

        For maximum efficiency, results will always be returned
        as a ZSTD compressed CSV string.
        """
        ),
    )
    output_format: Optional[Literal["json", "csv", "msgpack"]] = Field(
        "json", description="Output format: 'json' (default), 'csv', or 'msgpack'"
    )
    compression: Optional[Literal["brotli", "zstd"]] = Field(
        None, description="Compression algorithm: 'brotli' or 'zstd' (default: None)"
    )


def handler(event, *_):
    """
    AWS Lambda entry point to execute a SQL query against Snowflake.

    Args:
        event (dict): Input matching the QueryEvent schema.

    Returns:
        Query result as:
            - list[dict] if output_format is 'json' or 'msgpack'
            - str if output_format is 'csv'
            - bytes if compression is applied

    Raises:
        ValidationError: If the event doesn't match QueryEvent schema.
    """
    try:
        parsed = QueryEvent(**event)
    except ValidationError as ex:
        logger.error(f"Invalid event: {ex}")
        raise

    output_format = OutputFormat[parsed.output_format.upper()]

    _sf_client = sf_client or _sf_client_factory()

    if parsed.staged:
        stage_to_s3 = str(parsed.staged).upper() == "S3"
        logger.debug(
            "Executing staged SQL query using {}...",
            PlatformName.S3 if stage_to_s3 else PlatformName.SNOWFLAKE,
        )
        response = asyncio.run(
            staged.run_staged_query(_sf_client, parsed.sql, s3=stage_to_s3)
        )
        if not stage_to_s3:
            response.seek(0)  # Expecting a BytesIO object
            response = response.getvalue()

    else:
        logger.debug("Executing SQL query...")
        response = asyncio.run(
            _sf_client.afetch_all(parsed.sql, output_format=output_format)
        )
        response = _handle_response_non_staged(response, compression=parsed.compression)

    logger.debug("SQL query executed successfully.")

    return response


# Global persistent Snowflake client instance
# (for re-use across Lambda warm invocations)
sf_client: Optional[Client] = None


def _handle_response_non_staged(response, *, compression):
    """
    Handle the response from a non-staged SQL query execution.

    Args:
        response: The response from the SQL query.
        compression: The compression algorithm to apply to the result.

    Returns:
        The processed response, either as a string or bytes.
    """
    if compression:
        if not isinstance(response, bytes):
            if not isinstance(response, str):
                response = json.dumps(response)
            response = response.encode("utf-8")

        response_bytes = response.encode() if isinstance(response, str) else response
        if compression == CompressionFormats.BROTLI:
            return brotli.compress(response_bytes)
        elif compression == CompressionFormats.ZSTD:
            return zstandard.ZstdCompressor().compress(response_bytes)

    return response


def _sf_client_factory() -> Client:
    """Factory function to create a Snowflake client."""
    global sf_client
    if sf_client is None:
        client_settings = config.SNOWFLAKE_CLIENT_SETTINGS.copy()
        client_settings[SFClientSettings.PRIVATE_KEY] = get_snowflake_private_key()
        sf_client = Client(**client_settings)
    return sf_client