from io import StringIO
import logging
from typing import Any, Dict, Generator, Tuple

import pytest
from tests import config
from tests.conftest import db_session
from tests.utils import db
from tests.utils.csv_helpers import create_and_upload_csv_to_s3, read_csv_from_buffer

from utils import (
    assert_tax_correction_data,
    assert_vat_tax_correction_data,
    correction_file_data,
    vat_file_data,
)

# Logger setup
log = logging.getLogger(__name__)
# Configuration settings
LAMBDA_LOAD_FROM_S3 = config.LAMBDA_LOAD_FROM_S3
S3_BUCKET = config.LOAD_FROM_S3_BUCKET


@pytest.mark.parametrize(
    'test_parameters',
    [
        {
            's3_folder': 'post_tax_corrections',
            'filename_prefix': 'Correction',
            'table_name': 'worksheet_tax_correction',
        }
    ],
)
def test_tax_corrections(
    setup_file: Tuple[StringIO, str],
    s3_client: Any,
    db_session: Any,
    test_parameters: Dict[str, Any],
    teardown_file: Any,
) -> None:
    """Test tax corrections data loaded into the database when a CSV file is uploaded to S3_BUCKET ."""

    csv_buffer, s3_key = setup_file

    # Lambda function should be invoked automatically after the file is uploaded
    s3_client.assert_object_exists(S3_BUCKET, s3_key)

    # Verify the file data is loaded into the database
    table_name = test_parameters['table_name']
    csv_rows = read_csv_from_buffer(csv_buffer)

    for csv_row in csv_rows:
        account_id = str(csv_row['account_id'])
        note = str(csv_row['note'])

        conditions = {'account_id': account_id, 'note': note}
        db_result = db.wait_for_data_in_db(db_session, table_name, conditions)

        assert (
            db_result is not None
        ), f'Entity not found in {table_name} table for account_id-{account_id}'
        assert_tax_correction_data(csv_row, db_result, account_id)


@pytest.mark.parametrize(
    'test_parameters',
    [
        {
            's3_folder': 'post_tax_corrections_vat',
            'filename_prefix': 'VAT',
            'table_name': 'worksheet_tax_correction_vat',
        }
    ],
)
def test_vat_tax_corrections(
    setup_file: Tuple[StringIO, str],
    s3_client: Any,
    db_session: Any,
    test_parameters: Dict[str, Any],
    teardown_file: Any,
) -> None:
    """Test VAT tax corrections data loaded into the database when a CSV file is uploaded to S3_BUCKET ."""

    csv_buffer, s3_key = setup_file

    # The lambda function should be triggered automatically after the file is uploaded
    s3_client.assert_object_exists(S3_BUCKET, s3_key)

    # Verify that the file data is loaded into the database
    table_name = test_parameters['table_name']
    csv_rows = read_csv_from_buffer(csv_buffer)

    for csv_row in csv_rows:
        contract_id = str(csv_row['contract_id'])
        note = str(csv_row['note'])

        conditions = {'contract_id': contract_id, 'note': note}
        db_result = db.wait_for_data_in_db(db_session, table_name, conditions)

        assert (
            db_result is not None
        ), f'Entity not found in {table_name} table for contract_id-{contract_id}'

        assert_vat_tax_correction_data(csv_row, db_result, contract_id)


@pytest.fixture
def setup_file(s3_client: Any, request: Any) -> Tuple[StringIO, str]:
    """
    Prepares a CSV file in memory, uploads it to S3, and returns the buffer and S3 key.
    """
    # Retrieve test parameters from the fixture
    test_parameters = request.getfixturevalue('test_parameters')
    filename_prefix = test_parameters['filename_prefix']
    s3_folder = test_parameters['s3_folder']

    # Select the appropriate file data
    if filename_prefix == 'Correction':
        file_data = correction_file_data()
    elif filename_prefix == 'VAT':
        file_data = vat_file_data()
    else:
        raise ValueError(f'Unsupported filename prefix: {filename_prefix}')

    # Generate filename and S3 key, create CSV buffer, upload to S3, and return buffer and key
    csv_buffer, s3_key = create_and_upload_csv_to_s3(
        s3_client=s3_client,
        file_data=file_data,
        filename_prefix=filename_prefix,
        s3_folder=s3_folder,
        bucket=S3_BUCKET,
    )

    return csv_buffer, s3_key


@pytest.fixture
def teardown_file(
    s3_client: Any,
    db_session: Any,
    request: Any,
    setup_file: Tuple[StringIO, str],
) -> Generator[None, None, None]:
    """Fixture for cleaning up the test data."""
    test_parameters = request.getfixturevalue('test_parameters')
    table_name = test_parameters['table_name']
    csv_buffer, s3_key = setup_file

    yield  # Allow the test to run

    log.info(f'Cleaning up test data in table {table_name}')

    # Remove the test data from the database
    csv_rows = read_csv_from_buffer(csv_buffer)
    for csv_row in csv_rows:
        account_id = str(csv_row['account_id'])
        db.delete_entity_by_id(db_session, table_name, 'account_id', int(account_id))

    # Remove the file from S3
    s3_client.delete_if_object_present(S3_BUCKET, s3_key)