"""Lambda custom_reports function module."""

import calendar
from collections.abc import Callable
import io
import json
import math
from os import path
from typing import Any
from zipfile import ZIP_DEFLATED
from zipfile import ZipFile

from lambdacommon.common_config import logger
import pandas as pd
from pandas import DataFrame
import polars as pl
import sentry_sdk
from sentry_sdk import capture_exception
from sentry_sdk.integrations.aws_lambda import AwsLambdaIntegration

import config
from config import RevenueDisplayType
from src.connectors import s3
from src.connectors import snowflake
from src.connectors.ows_moneyhub import OwsMoneyhub
from src.exceptions import CustomReportException
from src.utils.constants import Encoding
from src.utils.constants import NumberFormat
from src.utils.features import is_feature_enabled

MAX_ROWS_PER_FILE = 1000000

if config.SENTRY_DSN:
    sentry_sdk.init(
        dsn=config.SENTRY_DSN,
        environment=config.ENVIRONMENT,
        integrations=[AwsLambdaIntegration(timeout_warning=True)],
    )


def _make_float_formatter(number_format: NumberFormat) -> Callable[[float], str]:
    """Create a float formatter function based on number format.

    Args:
        number_format (NumberFormat): The number format (US/EU).

    Returns:
        callable: Float formatter function.
    """
    if number_format == NumberFormat.EU:
        return lambda x: f'{x:.10f}'.replace('.', ',') if pd.notna(x) else ''
    else:
        return lambda x: f'{x:.10f}' if pd.notna(x) else ''


def _sanitize_for_tsv(dataframe: DataFrame) -> DataFrame:
    """Sanitize dataframe for TSV output by removing tabs and newlines.

    Args:
        dataframe (DataFrame): Input dataframe.

    Returns:
        DataFrame: Sanitized dataframe.
    """
    df_copy = dataframe.copy()
    for col in df_copy.select_dtypes(include=['object']).columns:
        df_copy[col] = (
            df_copy[col]
            .astype(str)
            .str.replace('\t', ' ')
            .str.replace('\n', ' ')
            .str.replace('\r', ' ')
        )
    return df_copy


def _write_file_to_buffer(
    dataframe: DataFrame, file_type: str, number_format: NumberFormat
) -> io.BytesIO:
    """Write dataframe to memory buffer (using Polars).

    Args:
        dataframe (DataFrame): Data frame containing the data to write.
        file_type (str): The type of the file.
        number_format (NumberFormat): The number format (US/EU).

    Returns:
        io.BytesIO: Buffer containing the file data.
    """
    internal_buffer = io.BytesIO()
    match file_type:
        case 'xls':
            separator = '\t'
            quoting = 'always'
            encoding = Encoding.UTF16
            include_bom = False
        case 'txt':
            separator = '\t'
            quoting = 'never'
            dataframe = _sanitize_for_tsv(dataframe)
            encoding = Encoding.UTF8
            include_bom = False
        case _:  # defaults to csv
            separator = ','
            quoting = 'non_numeric'
            encoding = Encoding.UTF8_SIG
            include_bom = True

    polars_dataframe = pl.from_pandas(dataframe, include_index=True)
    polars_dataframe.write_csv(  # type: ignore[call-overload]
        internal_buffer,
        separator=separator,
        include_header=True,
        line_terminator='\r\n',
        null_value='',
        decimal_comma=number_format == NumberFormat.EU,
        quote_style=quoting,
        include_bom=include_bom,
    )

    # reset the buffer
    internal_buffer.seek(0)

    # handle encoding
    if encoding != Encoding.UTF8:
        raw_bytes = internal_buffer.getvalue()
        internal_buffer.seek(0)
        internal_buffer.truncate()
        internal_buffer.write(raw_bytes.decode(Encoding.UTF8.value).encode(encoding.value))
        internal_buffer.seek(0)

    return internal_buffer


def _write_file_legacy(
    destination_path: str,
    dataframe: DataFrame | None,
    file_name: str,
    file_type: str,
    number_format: NumberFormat,
) -> None:
    """Write the zip file to disk (legacy single-file version).

    Args:
        destination_path (str): Path to store the file at.
        dataframe (DataFrame): Data frame containing the data to write.
        file_name (str): The name of the file.
        file_type (str): The type of the file.
        number_format (NumberFormat): The number format (US/EU).
    """
    if dataframe is None or len(dataframe) == 0:
        ZipFile(destination_path, 'w', ZIP_DEFLATED).close()
        return

    internal_buffer = _write_file_to_buffer(dataframe, file_type, number_format)

    with ZipFile(destination_path, 'w', ZIP_DEFLATED) as zip_file:
        zip_file.writestr(file_name, internal_buffer.getvalue())


def _write_file(
    destination_path: str,
    dataframe: DataFrame | None,
    file_name: str,
    file_type: str,
    number_format: NumberFormat,
) -> None:
    """Write the zip file to disk, splitting into multiple files if needed.

    Args:
        destination_path (str): Path to store the file at.
        dataframe (DataFrame): Data frame containing the data to write.
        file_name (str): The name of the file.
        file_type (str): The type of the file.
        number_format (NumberFormat): The number format (US/EU).
    """
    if dataframe is None or len(dataframe) == 0:
        ZipFile(destination_path, 'w', ZIP_DEFLATED).close()
        return

    total_rows = len(dataframe)
    num_files = math.ceil(total_rows / MAX_ROWS_PER_FILE)

    with ZipFile(destination_path, 'w', ZIP_DEFLATED) as zip_file:
        if num_files == 1:
            internal_buffer = _write_file_to_buffer(dataframe, file_type, number_format)
            zip_file.writestr(file_name, internal_buffer.getvalue())
        else:
            logger.info(f'Splitting {total_rows} rows into {num_files} files')
            base_name, extension = path.splitext(file_name)

            for file_num in range(1, num_files + 1):
                start_idx = (file_num - 1) * MAX_ROWS_PER_FILE
                end_idx = min(file_num * MAX_ROWS_PER_FILE, total_rows)

                df_chunk = dataframe.iloc[start_idx:end_idx]
                numbered_file_name = f'{base_name}_{file_num}{extension}'

                logger.info(
                    f'Creating file {file_num}/{num_files}: {numbered_file_name} '
                    f'(rows {start_idx}-{end_idx})'
                )

                internal_buffer = _write_file_to_buffer(df_chunk, file_type, number_format)
                zip_file.writestr(numbered_file_name, internal_buffer.getvalue())


def _build_statement_period_token(statement_periods_parsed: list[tuple[int, int]]) -> str:
    """Build time period token for file name.

    single: feb2025
    multi same year: feb-to-oct2025
    multi diff year: feb2025-to-mar2026
    """
    statement_periods_parsed.sort(key=lambda ym: (ym[0], ym[1]))
    (y1, m1) = statement_periods_parsed[0]
    (y2, m2) = statement_periods_parsed[-1]
    start_month = calendar.month_abbr[m1].lower()
    end_month = calendar.month_abbr[m2].lower()

    if len(statement_periods_parsed) == 1:
        return f'{start_month}{y1}'

    if y1 == y2:
        return f'{start_month}-to-{end_month}{y2}'

    return f'{start_month}{y1}-to-{end_month}{y2}'


def _build_file_name(report: dict, statement_periods_parsed: list[tuple[int, int]]) -> str:
    """Build the file name of a report.

    Args:
        report (dict): Report to make file name for.
        statement_periods_parsed (list): Parsed (start, end) statement period pairs.

    Returns:
        str: A file name appropriate for the report.
    """
    entity_id = str(report.get('subaccount_id') or report['account_id'])
    if report.get('contract_id'):
        entity_id += '_' + str(report['contract_id'])

    row_dimension = report['dimension_row']
    if row_dimension == 'service':
        row_dimension = 'store'

    # Build base filename
    periods_token = _build_statement_period_token(statement_periods_parsed)

    base_filename = 'report_{entity}_{periods}_{column}_{row}'.format(
        entity=entity_id,
        periods=periods_token,
        column=report['dimension_column'],
        row=row_dimension,
    )

    if report.get('subaccount_id'):
        subaccount_name = snowflake.get_subaccount_name(report['subaccount_id'])
        if subaccount_name:
            subaccount_clean = (
                subaccount_name.replace(' ', '_').replace('/', '_').replace('\\', '_')
            )
            base_filename += '_subaccount_{}'.format(subaccount_clean)

    # Apply character limit (max 255 chars for most filesystems)
    max_length = 251
    if len(base_filename) > max_length:
        base_filename = base_filename[:max_length]

    file_type = report.get('file_type', 'csv').lower()

    match file_type:
        case 'xls':
            extension = '.xls'
        case 'txt':
            extension = '.txt'
        case 'csv':
            extension = '.csv'
        case _:
            extension = '.csv'

    return base_filename + extension


def _generate_report(report_custom_id: int) -> None:
    """Generate the report.

    Args:
        report_custom_id (int): ID of the report to generate.
    """
    try:
        report = OwsMoneyhub.get_report_custom(report_custom_id)

        custom_filters = report.get('filters')
        if custom_filters:
            logger.info(f'Applying custom filters: {list(custom_filters.keys())}')

        logger.info('Requesting report data from Snowflake...')
        if report['dimension_column'] == config.DimensionType.FINANCIAL_DETAIL:
            data = snowflake.get_report_data_financial_detail(
                report['account_id'],
                report['contract_id'],
                report['statement_period_ids'],
                report['dimension_row'],
                report['revenue_type'],
            )
        else:
            data = snowflake.get_report_data(
                report['account_id'],
                report['contract_id'],
                report['statement_period_ids'],
                report['dimension_column'],
                report['dimension_row'],
                report['revenue_type'],
                report.get('revenue_display_type', RevenueDisplayType.NET),
                report.get('subaccount_id'),
                custom_filters,
            )

        logger.info('Creating zipped report file...')
        statement_periods_parsed = snowflake.get_statement_periods_parsed(
            report['statement_period_ids']
        )
        if statement_periods_parsed is None:
            raise CustomReportException('No statement periods found for report')
        filename = _build_file_name(report, statement_periods_parsed)
        number_format = NumberFormat[report['number_format'].upper()]

        local_path_parts = path.join(config.FILE_OUTPUT_PATH, filename).split('.')
        local_path_parts[-1] = 'zip'
        local_path = '.'.join(local_path_parts)

        # TODO remove legacy method after feature flag is rolled out
        is_enabled_multiple_files_ff = is_feature_enabled(
            'moneyhub_custom_reports_multiple_files', report['account_id']
        )
        if is_enabled_multiple_files_ff:
            logger.info('Using file splitting functionality (feature enabled)')
            _write_file(local_path, data, filename, report['file_type'], number_format)
        else:
            logger.info('Using legacy single-file functionality (feature disabled)')
            _write_file_legacy(local_path, data, filename, report['file_type'], number_format)

        logger.info(f'File created: {local_path}')

        if config.ENVIRONMENT == config.DEV_ENVIRONMENT:
            return

        # Upload to S3
        s3_object_path = path.join(
            str(report['account_id']), path.relpath(local_path, config.FILE_OUTPUT_PATH)
        )
        full_s3_path = s3.full_path(s3_object_path)
        logger.info(f'Uploading document to S3: {full_s3_path}')
        s3.upload_file(local_path, s3_object_path)

        # Update custom report db entry
        data = {
            'file_location': full_s3_path,
            'report_custom_status': config.CUSTOM_REPORT_STATUS_COMPLETE,
        }
        logger.info('Updating database entry ({}): {}'.format(report_custom_id, str(data)))

        OwsMoneyhub.update_report_custom(report_custom_id, **data)
    except Exception as e:
        logger.exception(str(e))
        data = {
            'file_location': None,
            'report_custom_status': config.CUSTOM_REPORT_STATUS_ERROR,
        }
        OwsMoneyhub.update_report_custom(report_custom_id, **data)
        raise e


def handler(event: dict[str, Any], _context: object) -> None:
    """Lambda entry point."""
    try:
        logger.info('event=' + str(event))

        if 'Records' in event:
            payloads = [json.loads(record['body']) for record in event['Records']]
        else:
            payloads = [event]

        for payload in payloads:
            if 'report_custom_id' not in payload:
                raise CustomReportException('Missing report_custom_id in event payload')

            _generate_report(int(payload['report_custom_id']))
    except CustomReportException as e:
        logger.exception(str(e))
        capture_exception(e)
        raise e