from typing import Any, Dict, Generator

import pytest

from tests import config
from tests.cross_lambda_integration.utils import (
    assert_payment_correction_data,
    assert_payment_vat_data,
    generate_file_data,
    get_abacus_event_id,
    get_payment_group_payment_id,
    resolve_statement_period_id,
    seed_data,
)
from tests.utils import db
from tests.utils.csv_helpers import create_and_upload_csv_to_s3, read_csv_from_buffer

LAMBDA_CALCULATE_PAYMENTS = config.LAMBDA_CALCULATE_PAYMENTS
S3_BUCKET = config.LOAD_FROM_S3_BUCKET

CLOSED_PERIOD_TEST_CASES = [
    {
        's3_folder': 'post_tax_corrections',
        'filename_prefix': 'TaxInsertionCorrection',
        'tax_table': 'worksheet_tax_correction',
        'statement_period_status': 'closed',
        'assert_fn': assert_payment_correction_data,
    },
    {
        's3_folder': 'post_tax_corrections_vat',
        'filename_prefix': 'TaxInsertionVAT',
        'tax_table': 'worksheet_tax_correction_vat',
        'statement_period_status': 'closed',
        'assert_fn': assert_payment_vat_data,
    },
]

CURRENT_PERIOD_TEST_CASES = [
    {
        's3_folder': 'post_tax_corrections',
        'filename_prefix': 'TaxInsertionCorrection',
        'tax_table': 'worksheet_tax_correction',
        'statement_period_status': 'current',
    },
    {
        's3_folder': 'post_tax_corrections_vat',
        'filename_prefix': 'TaxInsertionVAT',
        'tax_table': 'worksheet_tax_correction_vat',
        'statement_period_status': 'current',
    },
]


@pytest.mark.parametrize('test_parameters', CLOSED_PERIOD_TEST_CASES, indirect=True)
def test_tax_insertion_closed_statement_period(
    setup: Dict[str, Any],
    lambda_client: Any,
    db_session: Any,
    test_parameters: Dict[str, Any],
    teardown: None,
) -> None:
    """Test payment is generated for closed statement period."""

    account_id = seed_data()['account_id']

    lambda_event = {
        'event_name': 'calculate_payments',
        'abacus_event_id': setup['abacus_event_id'],
        'statement_period_id': setup['statement_period_id'],
        'target_type': 'payment_group_payment',
        'target_id': setup['payment_group_payment_id'],
    }

    response = lambda_client.invoke(LAMBDA_CALCULATE_PAYMENTS, lambda_event)
    assert response['StatusCode'] == 200, 'Lambda invocation failed'

    db.wait_for_generate_payments_to_complete(
        db_session, setup['payment_group_payment_id']
    )

    result = db.get_entity(
        db_session,
        'payment_group_payment_account',
        {'account_id': account_id},
    )

    assert (
        result is not None
    ), f"Expected data in 'payment_group_payment_account' for account_id={account_id}, but got None"

    # Custom assertion function from test parameters
    test_parameters['assert_fn'](setup['statement_period_id'], result)


@pytest.mark.parametrize('test_parameters', CURRENT_PERIOD_TEST_CASES, indirect=True)
def test_tax_insertion_current_statement_period(
    setup: Dict[str, Any],
    lambda_client: Any,
    db_session: Any,
    test_parameters: Dict[str, Any],
    teardown: None,
) -> None:
    """Ensure no payment is generated for current statement period."""

    account_id = seed_data()['account_id']

    lambda_event = {
        'event_name': 'calculate_payments',
        'abacus_event_id': setup['abacus_event_id'],
        'statement_period_id': setup['statement_period_id'],
        'target_type': 'payment_group_payment',
        'target_id': setup['payment_group_payment_id'],
    }

    response = lambda_client.invoke(LAMBDA_CALCULATE_PAYMENTS, lambda_event)
    assert response['StatusCode'] == 200, 'Lambda invocation failed'

    db.wait_for_generate_payments_to_complete(
        db_session, setup['payment_group_payment_id']
    )

    result = db.get_entity(
        db_session,
        'payment_group_payment_account',
        {
            'account_id': account_id,
            # Added to filter by current statement period as desired by the test case
            'current_statement_period_id': setup['statement_period_id'],
        },
    )

    assert result is None, f'Expected no data, but got: {result}'


@pytest.fixture
def test_parameters(request: Any) -> Any:
    """Return test parameters from indirect parametrization."""
    return request.param


@pytest.fixture
def setup(
    db_session: Any,
    s3_client: Any,
    test_parameters: Dict[str, Any],
) -> Dict[str, Any]:
    """Prepare test data and upload CSV."""

    statement_period_id = resolve_statement_period_id(
        db_session, test_parameters['statement_period_status']
    )
    payment_group_payment_id = get_payment_group_payment_id(
        db_session, statement_period_id, test_parameters['statement_period_status']
    )
    abacus_event_id = get_abacus_event_id(
        db_session, statement_period_id, payment_group_payment_id
    )

    if not payment_group_payment_id or not abacus_event_id:
        pytest.skip('Missing payment_group_payment_id or abacus_event_id.')

    file_data = generate_file_data(
        test_parameters['filename_prefix'], statement_period_id
    )
    csv_buffer, s3_key = create_and_upload_csv_to_s3(
        s3_client=s3_client,
        file_data=file_data,
        filename_prefix=test_parameters['filename_prefix'],
        s3_folder=test_parameters['s3_folder'],
        bucket=S3_BUCKET,
    )

    db.wait_for_data_in_db(
        db_session,
        test_parameters['tax_table'],
        {'account_id': seed_data()['account_id']},
    )

    return {
        'statement_period_id': statement_period_id,
        'payment_group_payment_id': payment_group_payment_id,
        'abacus_event_id': abacus_event_id,
        'csv_buffer': csv_buffer,
        's3_key': s3_key,
    }


@pytest.fixture
def teardown(
    s3_client: Any,
    db_session: Any,
    test_parameters: Dict[str, Any],
    setup: Dict[str, Any],
) -> Generator[None, None, None]:
    """Clean up DB and S3 after test."""
    yield

    account_id = seed_data()['account_id']
    pgp_id = setup['payment_group_payment_id']

    db.delete_entity_by_id(
        db_session, 'payment_group_payment_account', 'payment_group_payment_id', pgp_id
    )
    db.delete_entity_by_id(
        db_session,
        'worksheet_account_contract_payable_after_tax',
        'account_id',
        account_id,
    )

    for action in ['generate_payments', 'calculate_payments']:
        db.update_entity(
            db_session,
            'abacus_state',
            {
                'parent_table_id': pgp_id,
                'parent_table_name': 'payment_group_payment',
                'action_name': action,
            },
            {'action_status': 'init'},
        )

    for row in read_csv_from_buffer(setup['csv_buffer']):
        db.delete_entity_by_id(
            db_session,
            test_parameters['tax_table'],
            'account_id',
            int(row['account_id']),
        )

    s3_client.delete_if_object_present(S3_BUCKET, setup['s3_key'])