"""Helper utilities for integration tests."""

import gzip
import io
import json
import os

import boto3
import mysql.connector
from utils.dockerized_lambda_client import DockerizedLambdaAPIClient

QA_BASE_URL = os.environ.get('QA_BASE_URL', 'http://lambda:8080')


def dockerized_lambda_api_client(headers):
    """Create adjustments_json_validation dockerized lambda object."""
    return DockerizedLambdaAPIClient(QA_BASE_URL, headers)


def get_db_connection():
    """Return a MySQL connection to the royalty_accounting database."""
    return mysql.connector.connect(
        host='mysql',
        port='3306',
        database='royalty_accounting',
        user='royalties',
        password='1234',
    )


def get_invalid_file_location(statement_period_adjustment_file_id: int) -> str:
    """Return the invalid_file_location for a statement_period_adjustment_file row.

    Args:
        statement_period_adjustment_file_id: the primary key of the record.

    Returns:
        The S3 URI stored in invalid_file_location.
    """
    cnx = get_db_connection()
    cursor = cnx.cursor()
    cursor.execute(
        'SELECT invalid_file_location FROM statement_period_adjustment_file '
        'WHERE statement_period_adjustment_file_id = %s',
        (statement_period_adjustment_file_id,),
    )
    row = cursor.fetchone()
    cursor.close()
    cnx.close()
    assert row is not None, (
        f'No statement_period_adjustment_file found with id={statement_period_adjustment_file_id}'
    )
    assert row[0] is not None, (
        f'invalid_file_location is NULL for statement_period_adjustment_file id={statement_period_adjustment_file_id}'
    )
    return row[0]


def get_s3_file_contents(s3_uri: str) -> list[dict]:
    """Fetch a gzip-compressed JSON-lines file from S3 and return its records.

    Args:
        s3_uri: full S3 URI in the form ``s3://bucket/key``.

    Returns:
        A list of dicts, one per line in the decompressed file.
    """
    assert s3_uri.startswith('s3://'), f'Expected an s3:// URI, got: {s3_uri}'
    without_scheme = s3_uri[len('s3://') :]
    bucket, _, key = without_scheme.partition('/')

    s3_client = boto3.client('s3')
    response = s3_client.get_object(Bucket=bucket, Key=key)
    compressed_data = response['Body'].read()

    with gzip.GzipFile(fileobj=io.BytesIO(compressed_data)) as gz:
        text = gz.read().decode('utf-8')

    records = [json.loads(line) for line in text.splitlines() if line.strip()]
    return records