"""
This module contains the logic for the consumer.
"""

import asyncio
import csv
from datetime import datetime
from functools import lru_cache
from functools import partial
from io import BytesIO
from typing import Iterator

import xlsxwriter

from common.src import logger
from common.src.aws import new_s3_client
from common.src.aws.utils import get_s3_bucket_key_from_uri, s3_object_body_stream
from common.src.connectors.snowflake import run_sql
from common.src.enums import FileExtensions
from common.src.typings import OrchLabelId
from common.src.typings import S3URI
from common.src.utils import compression
from . import sql_queries
from .enums import ColumnAlias, ReportTabNames
from .models import ReportColumn, ReportSheet
from ..threadpool import executor
from ..typings import PeriodId
from ... import config

logger = logger.new_logger(__name__)


def run(label_id: OrchLabelId, period_id: PeriodId) -> tuple[str, str]:
    """Wrapper to run the consumer in async mode.

    Args:
        label_id (OrchLabelId): The label ID.
        period_id (PeriodId): The period ID.

    Returns:
        tuple[str, str]: The S3 bucket and key where the Excel file has been
            uploaded.
    """
    return asyncio.run(_main(label_id, period_id))


async def _main(label_id: OrchLabelId, period_id: PeriodId) -> tuple[str, str]:
    """Main function to run the consumer in async mode.

    This function will run the SQL queries to get the data which will
    then be used to generate the report, which in turn will be built
    into an Excel file and its bytes will be stored in an S3 bucket.

    Returns:
        tuple[str, str]: The S3 bucket and key where the Excel file has been
            uploaded.
    """
    loop = asyncio.get_running_loop()
    run_exec = partial(loop.run_in_executor, executor)

    async def run_with_delay(fn, delay: int):
        await asyncio.sleep(delay)
        return await run_exec(fn)

    tasks = [
        run_exec(
            lambda: run_sql(sql_queries.full_detail(label_id, period_id), staged="s3")
        ),
        run_with_delay(
            lambda: run_sql(sql_queries.royalties_by_composer(label_id, period_id)), 1
        ),
    ]

    logger.debug(
        "Running SQL queries in parallel and doing ReportColumn stuff "
        "and S3 Client instantiation in the meantime..."
    )

    # Order of columns below will be the order of the columns in the Excel file

    composer_columns = [
        ReportColumn(
            alias=ColumnAlias.PUB_WRITER_ID,
            human_name="Composers Pub Writer ID",
        ),
        ReportColumn(alias=ColumnAlias.LEGAL_NAME, human_name="Composers Legal Name"),
        ReportColumn(alias=ColumnAlias.IPI, human_name="Composers IPI"),
        ReportColumn(alias=ColumnAlias.PRO, human_name="Composers PRO"),
        ReportColumn(alias=ColumnAlias.ADJUSTED_GROSS, human_name="Adjusted Gross"),
        ReportColumn(alias=ColumnAlias.NET_REVENUE, human_name="Net Revenue"),
        ReportColumn(
            alias=ColumnAlias.NET_REVENUE_PREFERRED_CURRENCY,
            human_name="Net Revenue in Preferred Currency",
        ),
    ]

    full_detail_columns = [
        ReportColumn(alias=ColumnAlias.SONG_NO, human_name="song_no"),
        ReportColumn(alias=ColumnAlias.SONG, human_name="song"),
        ReportColumn(alias=ColumnAlias.WRITER, human_name="writer"),
        ReportColumn(alias=ColumnAlias.SOURCE1_NAME, human_name="src_nm"),
        ReportColumn(alias=ColumnAlias.SOURCE1_COUNTRY, human_name="src_ctry"),
        ReportColumn(alias=ColumnAlias.SOURCE2_NAME, human_name="src2_nm"),
        ReportColumn(alias=ColumnAlias.SOURCE2_COUNTRY, human_name="src2_ctry"),
        ReportColumn(alias=ColumnAlias.SOURCE3_NAME, human_name="src3_nm"),
        ReportColumn(alias=ColumnAlias.SOURCE3_COUNTRY, human_name="src3_ctry"),
        ReportColumn(alias=ColumnAlias.SOURCE4_NAME, human_name="src4_nm"),
        ReportColumn(alias=ColumnAlias.SOURCE4_COUNTRY, human_name="src4_ctry"),
        ReportColumn(alias=ColumnAlias.INCOME_TYPE, human_name="inc_typ"),
        ReportColumn(alias=ColumnAlias.SH_ID, human_name="sh_id"),
        ReportColumn(alias=ColumnAlias.REPORTING_PERIOD, human_name="rptg_pd"),
        ReportColumn(alias=ColumnAlias.SALES_PERIOD, human_name="sales_pd"),
        ReportColumn(alias=ColumnAlias.PRODUCT_NUMBER, human_name="prod_no"),
        ReportColumn(alias=ColumnAlias.ARTIST_PRODUCT_NUMBER, human_name="art_prd_no"),
        ReportColumn(alias=ColumnAlias.UNITS, human_name="units"),
        ReportColumn(alias=ColumnAlias.SONG_SHARE_PCT, human_name="sng_shr_pct"),
        ReportColumn(alias=ColumnAlias.CONTROL_PCT, human_name="cntrl_pct"),
        ReportColumn(alias=ColumnAlias.AMOUNT, human_name="amount"),
        ReportColumn(alias=ColumnAlias.DF, human_name="df"),
        ReportColumn(alias=ColumnAlias.SOURCE_PRODUCT, human_name="src_prod"),
        ReportColumn(alias=ColumnAlias.ISWC_CD, human_name="iswc_cd"),
        ReportColumn(alias=ColumnAlias.PUB_SONG_ID, human_name="ext_song"),
        ReportColumn(alias=ColumnAlias.ARTIST, human_name="artist"),
        ReportColumn(alias=ColumnAlias.SOURCE_SONG, human_name="src_song"),
        ReportColumn(alias=ColumnAlias.ISRC, human_name="isrc"),
        ReportColumn(alias=ColumnAlias.VENDOR_ID, human_name="vendor_id"),
        ReportColumn(alias=ColumnAlias.ADJUSTED_GROSS, human_name="adjusted_gross"),
        ReportColumn(alias=ColumnAlias.FEE_PCT, human_name="fee_pct"),
        ReportColumn(alias=ColumnAlias.NET_REVENUE, human_name="net_revenue"),
        ReportColumn(
            alias=ColumnAlias.PREFERRED_CURRENCY,
            human_name="preferred_currency",
        ),
        ReportColumn(
            alias=ColumnAlias.CURRENCY_CONVERSION_RATE,
            human_name="currency_conversion_rate",
        ),
        ReportColumn(
            alias=ColumnAlias.NET_REVENUE_PREFERRED_CURRENCY,
            human_name="net_revenue_preferred_currency",
        ),
    ]

    # Use the CPU idleness to do some work
    s3_client = new_s3_client(config.AWS_REGION)
    xlsx_key = _generate_key(label_id, period_id)

    full_detail_s3_files, royalties_composer = await asyncio.gather(*tasks)

    logger.debug(
        "SQL queries finished successfully, proceeding with sheet construction..."
    )

    sheets = [
        ReportSheet(
            data=royalties_composer,
            name=ReportTabNames.COMPOSER,
            columns=composer_columns,
        ),
        ReportSheet(
            data=_iter_s3_staged_results(s3_client, full_detail_s3_files),
            name=ReportTabNames.FULL_DETAIL,
            columns=full_detail_columns,
        ),
    ]

    xlsx_buffer = _build_xlsx(sheets)
    logger.info(
        "Excel report built successfully for label ID: {} and period ID: {}",
        label_id,
        period_id,
    )

    s3_client.upload_fileobj(xlsx_buffer, config.S3_BUCKET, xlsx_key)
    logger.info(
        "Excel report uploaded to S3 bucket: {} with key: {}",
        config.S3_BUCKET,
        xlsx_key,
    )

    return config.S3_BUCKET, xlsx_key


def _generate_key(label_id: OrchLabelId, period_id: PeriodId) -> str:
    """Generate the S3 key for the Excel file.

    The key is generated based on the label ID and period ID, and outputs
    a string in the format:

        L<label_id>_<period_id>_Publishing_Statement_<Month Name>_<Year>.xlsx

    The key has a nested folder structure:
        1st level: label_id
        2nd level: period_id
        3rd level: Publishing_Statement_<Month Name>_<Year>.xlsx

    Args:
        label_id (OrchLabelId): The label ID.
        period_id (PeriodId): The period ID.

    Returns:
        str: The generated S3 key. E.g.:
            L12345_312_Publishing_Statement_April_2025.xlsx
    """
    now = datetime.now()
    current_month_name = now.strftime("%B")
    return (
        f"{label_id}/{period_id}/"  # Nested folder structure
        f"L{label_id}_{period_id}_"
        f"Publishing_Statement_{current_month_name}_{now.year}{FileExtensions.XLSX}"
    )


def _iter_s3_staged_results(s3_client, uris: list[S3URI]) -> Iterator[dict[str, str]]:
    """Iterate over the staged results in S3 in a memory-efficient way, by handling
    one file at a time and getting file data + decompressing ZSTD + reading CSV
    using streams.

    Args:
        s3_client: The S3 client to use for reading the staged files.
        uris (list[S3URI]): List of S3 URIs to read (the files staged by Snowflake).
    """
    headers = None
    for uri in uris:
        bucket_name, key = get_s3_bucket_key_from_uri(uri)
        logger.debug("Fetching staged file: {} ...", uri)

        with s3_object_body_stream(s3_client, bucket_name, key, delete=True) as content:
            csv_file_buffer = compression.zstd_decompress(content)

            # Only the first file has headers, so we persist them for the additional files
            reader = (
                csv.DictReader(csv_file_buffer)
                if headers is None
                else csv.DictReader(csv_file_buffer, fieldnames=headers)
            )
            if headers is None:
                headers = reader.fieldnames
                # Sanity check. All headers must be nonempty strings, containing only
                # alphanumeric characters and underscores.
                assert all(
                    isinstance(h, str) and h.isidentifier() for h in headers
                ), f"Invalid headers: {headers}"

            for row in reader:
                # Replace \N with empty string (those are empty values in Snowflake)
                yield {k: ("" if v == "\\N" else v) for k, v in row.items()}


def _build_xlsx(sheets: list[ReportSheet]) -> BytesIO:
    """Build the report in Excel format.

    The function uses the xlsxwriter library for performance reasons.

    Args:
        sheets (list[ReportSheet]): List of ReportSheet objects containing
            the data to be included in the report.

    Returns:
        BytesIO: A BytesIO object containing the Excel file.
    """
    buffer = BytesIO()

    # Use constant_memory mode to reduce memory usage by writing rows directly to file
    # Use in_memory to avoid writing to disk
    workbook = xlsxwriter.Workbook(
        buffer,
        {
            "in_memory": True,
            "constant_memory": True,
            "default_date_format": None,  # Avoid unnecessary formatting overhead
        },
    )

    for sheet in sheets:
        logger.debug("Building sheet: {} ...", sheet.name)

        # Precompute aliases and headers only once
        aliases = [col.alias.lower() for col in sheet.columns]
        headers = [col.human_name for col in sheet.columns]

        max_rows_per_sheet = 1_000_000
        tab_idx = 0
        row_in_sheet = 1

        worksheet = workbook.add_worksheet(sheet.name)
        worksheet.write_row(0, 0, headers)  # Write headers once

        col_widths = [len(h) for h in headers]

        # Preallocate a reusable list for each row to avoid re-creating lists,
        # this has a performance impact when writing large amounts of data
        # (millions of rows)
        row_values = [""] * len(aliases)

        for row in sheet.data:
            if row_in_sheet == max_rows_per_sheet:
                # Auto-fit column widths before starting a new tab
                for i, width in enumerate(col_widths):
                    worksheet.set_column(i, i, width + 1)

                tab_idx += 1
                worksheet = workbook.add_worksheet(f"{sheet.name} ({tab_idx + 1})")
                worksheet.write_row(0, 0, headers)
                row_in_sheet = 1
                col_widths = [len(h) for h in headers]

            row_lower = {k.lower(): v for k, v in row.items()}

            # Reuse preallocated row_values list to avoid per-row allocations
            for i, a in enumerate(aliases):
                val = row_lower.get(a)
                val_str = "" if val is None else str(val)
                casted_val = _try_cast_number(val_str)
                row_values[i] = casted_val
                col_widths[i] = max(col_widths[i], len(val_str))

            worksheet.write_row(row_in_sheet, 0, row_values)
            row_in_sheet += 1

        # Final column width adjustment after last sheet is filled
        for i, width in enumerate(col_widths):
            worksheet.set_column(i, i, width + 1)

    # Clear the cache, we don't need it anymore for now
    _try_cast_number.cache_clear()

    workbook.close()
    buffer.seek(0)
    return buffer


@lru_cache(maxsize=4096)
def _try_cast_number(value: str) -> str | int | float:
    """Try to cast a string to a number, so that the value appears as a number
    in Excel. Uses a cache for performance.
    """
    if not value or not value[0].isdigit() and value[0] not in "-.":
        return value  # fast skip for obvious non-numbers
    try:
        return float(value) if "." in value else int(value)
    except ValueError:
        return value