"""Integration tests for lambda function."""

# flake8: noqa F401, F811
import msgpack
import pytest
import brotli
from ..conftest import sample_payload, SQL_QUERY_LIMIT
from ...src import app
from ...src.enums import PayloadKeys, OutputFormat
from common.src.enums import CompressionFormats, AWSPrefixes
from common.src.utils.compression import zstd_decompress
import zstandard
import json


@pytest.mark.parametrize(PayloadKeys.OUTPUT_FORMAT, [OutputFormat.JSON, None])
def test_handler_default_format_json(sample_payload, output_format) -> None:
    event = {**sample_payload}
    if output_format:
        event[PayloadKeys.OUTPUT_FORMAT] = output_format
    response = app.handler(event, None)
    assert isinstance(response, list)
    assert len(response) == SQL_QUERY_LIMIT


def test_handler_csv(sample_payload):
    event = {**sample_payload, PayloadKeys.OUTPUT_FORMAT: OutputFormat.CSV}
    response = app.handler(event, None)
    assert isinstance(response, str)
    assert len(response.splitlines()) == SQL_QUERY_LIMIT + 1


def test_handler_msgpack(sample_payload):
    event = {**sample_payload, PayloadKeys.OUTPUT_FORMAT: OutputFormat.MSGPACK}
    response = app.handler(event, None)
    unpacked = msgpack.unpackb(response, strict_map_key=False)
    assert isinstance(unpacked, list)
    assert len(unpacked) == SQL_QUERY_LIMIT


@pytest.mark.parametrize(
    "compression, decompress",
    [
        (CompressionFormats.BROTLI, brotli.decompress),
        (CompressionFormats.ZSTD, zstandard.decompress),
    ],
)
def test_handler_compressed(sample_payload, compression, decompress):
    event = {**sample_payload, PayloadKeys.COMPRESSION: compression}
    response = app.handler(event, None)
    assert isinstance(response, bytes)
    assert len(response) > 0
    decompressed = decompress(response)
    assert "TABLE" in decompressed.decode("utf-8")


def test_handler_staged(sample_payload):
    event = {**sample_payload, PayloadKeys.STAGED: True}
    response = app.handler(event, None)
    assert isinstance(response, bytes)
    assert len(response) > 0
    decompressed = zstd_decompress(response).read()
    assert "TABLE" in decompressed, "Expected 'TABLE' in decompressed data"
    first_csv_line = decompressed.splitlines()[0]
    assert "TABLE" in first_csv_line, "Expected 'TABLE' in first line of CSV"
    try:
        is_json = json.loads(decompressed)
    except json.JSONDecodeError:
        is_json = False
    assert not is_json, "Expected decompressed data to be CSV, not JSON"


def test_handler_staged_s3(sample_payload):
    event = {**sample_payload, PayloadKeys.STAGED: "S3"}
    response = app.handler(event, None)
    assert isinstance(response, list)
    assert len(response)
    assert all(
        s3_uri.startswith(AWSPrefixes.S3) for s3_uri in response
    ), f"Expected all S3 URIs to start with '{AWSPrefixes.S3}"
    assert all(
        s3_uri.endswith(".csv.zst") for s3_uri in response
    ), "Expected all S3 URIs to end with '.csv.zst'"