"""
This module provides a high-level interface to Snowflake,
for use in other Lambda functions.
"""

import os
import json
from enum import StrEnum
from io import StringIO
from typing import Iterator
from typing import Literal
import brotli
import msgpack
import zstandard

from common.src.enums import AWSPayload, Formats, CompressionFormats
from common.src.models.sql import BaseColumns
from common.src.typings import SqlString
from common.src import aws, logger

logger = logger.new_logger(__name__)

SF_QUERY_FUNCTION_NAME: str = "qa-lambda-business-solutions-automations-sf_query"

ALLOWED_OUTPUT_FORMATS: set[Formats] = {Formats.JSON, Formats.CSV}

Staged = bool | Literal["S3"]


class PayloadKeys(StrEnum):
    """Keys for the payload sent to the Lambda function."""

    SQL = "sql"
    OUTPUT_FORMAT = "output_format"
    COMPRESSION = "compression"
    STAGED = "staged"


class _RunSQL:
    """Wrapper for running SQL queries on Snowflake via AWS Lambda.

    This class provides a high-level interface to run SQL queries on Snowflake
    using the `sf_query` Lambda function. It handles the preparation of inputs,
    invocation of the Lambda function, and processing of the response.
    The class is designed to be used as a singleton, and it initializes the
    AWS Lambda client on the first invocation.

    Attributes:
        default_region (str): Default AWS region to use for the Lambda client.
        lambda_client (None): AWS Lambda client instance. It will be initialized
          on the first invocation of the `run` method and cached for future use.
    """

    default_region: str = "us-east-1"
    lambda_client = None

    def run(
        self,
        sql: SqlString,
        *,
        staged: Staged = False,
        output_format: Formats = Formats.JSON,
        resp_model: BaseColumns = None,
    ) -> Iterator[dict] | str:
        """
        Run a SQL query using the `sf_query` Lambda function.

        Under the hood, this function uses the `sf_query` Lambda function to run
        the SQL query on Snowflake and return the results.

        Args:
            sql (str): SQL query string.
            staged (Staged): If True, the SQL query will be staged in Snowflake.
                The result will be a CSV file compressed with Zstandard (ZSTD),
                regardless of the output format specified. This option is for
                heavy and long-running queries that need to be parallelized in
                Snowflake, and is memory efficient. However, it's an overkill
                for simple, straightforward queries.
            output_format (str): Either "json" (default) or "csv".
            resp_model (BaseColumns): Only effective if output_format is "json".
                Optional Pydantic model to wrap the results. If provided,
                the results will be validated against this model.

        Returns:
            - List of dictionaries if format is "json" or iterator of
                pydantic models if resp_model is provided.
            - String if format is "csv"
        """
        if output_format not in ALLOWED_OUTPUT_FORMATS:
            raise ValueError("Invalid output_format: must be 'json' or 'csv'")
        output_format = "msgpack" if output_format == Formats.JSON else Formats.CSV

        pk = PayloadKeys
        payload = {
            pk.SQL.value: sql,
            pk.STAGED.value: staged,
            pk.OUTPUT_FORMAT.value: output_format.lower(),
            pk.COMPRESSION.value: CompressionFormats.BROTLI,
        }

        if self.lambda_client is None:
            self._set_lambda_client()

        logger.debug("Running SQL query on Snowflake via Lambda...")

        response = self.lambda_client.invoke(
            FunctionName=SF_QUERY_FUNCTION_NAME,
            InvocationType=AWSPayload.REQUEST_RESPONSE,
            Payload=json.dumps(payload),
        )

        logger.debug("SQL query executed.")

        return self._handle_response(
            response, staged=staged, output_format=output_format, model=resp_model
        )

    def _set_lambda_client(self):
        """Set the AWS Lambda client."""
        region_name = os.environ.get("AWS_REGION", self.default_region)
        self.lambda_client = aws.new_lambda_client(region_name=region_name)

    def _handle_response(
        self,
        response,
        *,
        staged: Staged = False,
        output_format: Formats = Formats.JSON,
        model: BaseColumns | None = None,
    ) -> Iterator[dict | str]:
        """Handle the response from the Lambda function, as obtained from
        the AWS SDK.

        Args:
            response: The response from the Lambda function.
            output_format (Formats): The format of the output.
            model (BaseColumns | None): Pydantic model to wrap the results (if
                applicable).

        Returns:
            Iterator[dict | str]: An iterator of dictionaries or strings,
                depending on the output format.
        """
        if isinstance(staged, str) and staged.upper() == "S3":
            s3_file_names = json.loads(
                response[AWSPayload.PAYLOAD].read().decode("utf-8")
            )
            return iter(s3_file_names)

        payload_content = self._decompress_payload_content(
            response[AWSPayload.PAYLOAD].read()
        )

        if staged:
            return self._handle_zstd_csv(payload_content)

        if output_format == Formats.CSV:
            return self._handle_csv(payload_content)

        # It's a brotli compressed, msgpack serialized, JSON response
        return self._handle_decompressed_json_payload_content(payload_content, model)

    @staticmethod
    def _handle_zstd_csv(payload_content: bytes):
        r"""Decompresses Zstandard-compressed CSV content and yields each line
        as a UTF-8 string.
        The header row is lowercased; all other rows are returned as-is.
        Preserves original newline behavior (only adds '\n' where originally present).

        Args:
            payload_content (bytes): Zstandard-compressed CSV data.

        Yields:
            str: Each line of the decompressed CSV, with header lowercased.
        """
        decompressor = zstandard.ZstdDecompressor()
        with decompressor.stream_reader(payload_content) as reader:
            buffer = b""
            first = True
            while True:
                chunk = reader.read(8192)
                if not chunk:
                    break
                buffer += chunk
                while b"\n" in buffer:
                    line, buffer = buffer.split(b"\n", 1)
                    line = line.decode("utf-8")
                    yield line.lower() + "\n" if first else line + "\n"
                    first = False
            if buffer:
                line = buffer.decode("utf-8")
                yield line.lower() if first else line

    @staticmethod
    def _handle_csv(payload_content: bytes) -> Iterator[str]:
        """Processes UTF-8 encoded CSV content and yields each line as a string.
        The header row is lowercased; all other rows are returned as-is.

        Args:
            payload_content (bytes): Raw CSV data in UTF-8 encoding.

        Yields:
            str: Each line of the CSV, with the header lowercased.
        """
        buffer = StringIO(payload_content.decode("utf-8"))
        first = True
        for line in buffer:
            if first:
                yield line.lower()
                first = False
            else:
                yield line

    @staticmethod
    def _decompress_payload_content(brotli_data: bytes) -> bytes:
        """Decompress the brotli compressed content.

        Args:
            brotli_data (bytes): The compressed content.
        """
        try:
            return brotli.decompress(brotli_data)
        except brotli.error as e:
            msg_decoded = brotli_data.decode("utf-8")
            logger.error(f"Brotli decompression failed: {e}, error: {msg_decoded}")
            raise

    @staticmethod
    def _handle_decompressed_json_payload_content(
        msgpack_content: bytes, resp_model: BaseColumns | None
    ) -> Iterator[dict | BaseColumns]:
        """Handle the decompressed msgpack content.

        Args:
            msgpack_content (bytes): The decompressed msgpack content.
            resp_model (BaseColumns | None): Pydantic model to wrap the results.
              The model attributes MUST BE LOWERCASE, as the JSON keys are
              lower cased before being passed to the model.
        """
        json_data = msgpack.unpackb(msgpack_content, strict_map_key=False)
        return (
            (resp_model(**{k.lower(): v for k, v in r.items()}) for r in json_data)
            if resp_model
            else ({k.lower(): v for k, v in r.items()} for r in json_data)
        )


# Singleton instance of the `_run_sql` function
run_sql = _RunSQL().run