"""Utilities associated with document generartion."""

import csv
from decimal import Decimal
import io
import math
import os
from os import path
import re
from typing import Any
from typing import Callable
from typing import Generator
from typing import Iterator
from typing import Literal
from zipfile import ZIP_DEFLATED
from zipfile import ZipFile

import jinja2
from lambdacommon.common_config import logger
import pandas
from pandas import DataFrame
import pdfkit
from playwright.sync_api import sync_playwright
import polars as pl
import pyarrow as pa

import config
from src.utils import as_decimal
from src.utils.constants import Encoding
from src.utils.constants import NumberFormat
from src.utils.constants import VAT_CATEGORY_DESCRIPTIONS
from src.utils.error_handling import LambdaException

MAX_ROWS_PER_FILE = 1000000
POLARS_ARROW_CONVERSION_THRESHOLD = 10000000


def format_row_numbers(row: list, number_format: NumberFormat) -> list:
    """Format Decimal numbers in a row using US or EU number format.

    Args:
        row (list): List of values, including Decimals.
        number_format (str): Format for numbers ('us' or 'eu')

    Returns:
        list: Row with formatted numbers as strings.
    """
    if number_format == NumberFormat.EU:
        thousands_sep = '.'
        decimal_sep = ','
    else:
        thousands_sep = ','
        decimal_sep = '.'

    def _format_number(val: Decimal) -> str:
        s = f'{val:.20f}'.rstrip('0').rstrip('.')
        integer_part, _, fractional_part = s.partition('.')
        if integer_part != '-0':
            integer_part = f'{int(integer_part):_}'.replace('_', thousands_sep)
        return f'{integer_part}{decimal_sep}{fractional_part}' if fractional_part else integer_part

    return [_format_number(item) if isinstance(item, Decimal) else item for item in row]


def _write_file(
    headers: list,
    iterator: Iterator[dict],
    max_rows: int,
    destination_path: str,
    row_processor: Callable[[dict], list] | None = None,
    file_type: str = 'csv',
) -> None:
    """Write the contents to a file (CSV, TXT, XLS) with appropriate formatting.

    Args:
        headers (list): Column headers.
        iterator (Iterator): Data rows.
        max_rows (int): Max number of rows to write.
        destination_path (str): File path.
        row_processor (Callable): Optional row processor.
        file_type (str): File format ('csv', 'txt', 'xls').
    """
    if file_type == 'xls':
        dialect = 'excel-tab'
        encoding = 'utf-16'
        quoting = csv.QUOTE_ALL
    elif file_type == 'txt':
        dialect = 'excel-tab'
        encoding = 'utf-8'
        quoting = csv.QUOTE_MINIMAL
    else:  # defaults to csv
        dialect = 'excel'
        encoding = 'utf-8-sig'
        quoting = csv.QUOTE_NONNUMERIC

    row_count = 0
    with open(destination_path, 'w', newline='', encoding=encoding) as file:
        writer = csv.writer(file, dialect=dialect, quoting=quoting)
        writer.writerow(headers)

        for row in iterator:
            content = row_processor(row) if row_processor else row
            writer.writerow(content)

            row_count += 1
            if row_count == max_rows:
                break


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

    Args:
        number_format (NumberFormat): Number format that informs the formatting.

    Returns:
        Callbable: Function to format float values.
    """

    def formatter(val: float) -> str:
        """Format the float value for output.

        Args:
            val (float): Value to format.

        Returns:
            str: Formatted text.
        """
        if number_format == NumberFormat.EU:
            thousands_sep = '.'
            decimal_sep = ','
        else:
            thousands_sep = ','
            decimal_sep = '.'

        s = f'{val:.12f}'.rstrip('0').rstrip('.')
        integer_part, _, fractional_part = s.partition('.')
        if integer_part != '-0':
            integer_part = f'{int(integer_part):_}'.replace('_', thousands_sep)

        return f'{integer_part}{decimal_sep}{fractional_part}' if fractional_part else integer_part

    return formatter


CONTROL_CHARS = dict.fromkeys(
    # strip CR, LF, TAB; also remove other ASCII control chars
    [*range(0x00, 0x09), 0x0A, 0x0B, 0x0C, 0x0D, *range(0x0E, 0x20)],
    None,
)


def _sanitize_for_tsv(df: DataFrame) -> DataFrame:
    """Sanitize a DataFrame for TSV output by stripping control characters.

    Args:
        df (DataFrame): DataFrame containing the data to sanitize.

    Returns:
        DataFrame: DataFrame with control characters removed from string values.
    """
    return df.applymap(lambda v: (v.translate(CONTROL_CHARS) if isinstance(v, str) else v))


def _write_file_pandas_polars(
    dataframe: DataFrame | None, file_type: str, number_format: NumberFormat, buffer: io.BytesIO
) -> None:
    """Write the file to the memory buffer (using Polars).

    Args:
        dataframe (DataFrame | None): Data frame containing the data to write.
        file_type (str): The type of file.
        number_format (NumberFormat): The number format (US/EU).
        buffer (io.BytesIO): Memory buffer to write the file to.
    """
    include_bom = False
    quoting: Literal['necessary', 'always', 'non_numeric', 'never']
    if file_type == 'xls':
        separator = '\t'
        quoting = 'always'
        encoding = Encoding.UTF16
    elif file_type == 'txt':
        separator = '\t'
        quoting = 'never'
        dataframe = _sanitize_for_tsv(dataframe)
        encoding = Encoding.UTF8
    else:  # defaults to csv
        separator = ','
        quoting = 'non_numeric'
        encoding = Encoding.UTF8_SIG
        include_bom = True

    if dataframe is not None:
        if len(dataframe) < POLARS_ARROW_CONVERSION_THRESHOLD:
            polars_dataframe: pl.DataFrame = pl.from_arrow(  # type: ignore[assignment]
                pa.Table.from_pandas(dataframe)
            )
        else:
            polars_dataframe = pl.from_pandas(dataframe)  # type: ignore[assignment]
        polars_dataframe.write_csv(
            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,
            float_scientific=False,
        )

        # handle encoding
        if encoding != Encoding.UTF8:
            # reset the buffer
            buffer.seek(0)
            utf8_content = buffer.read()  # polars write to utf8 by default

            encoded_content = utf8_content.decode(Encoding.UTF8.value).encode(
                encoding.value, errors='replace'
            )  # noqa: E501

            # wipe the buffer, then write with new encoded content
            buffer.seek(0)
            buffer.truncate()
            buffer.write(encoded_content)

        # reset buffer to start position
        buffer.seek(0)


def _write_file_pandas(
    dataframe: DataFrame | None, destination_path: str, file_type: str, number_format: NumberFormat
) -> None:
    """Write the actual file to disk (using Pandas).

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

    if file_type == 'xls':
        separator = '\t'
        encoding = 'utf-16'
        quoting = csv.QUOTE_ALL
    elif file_type == 'txt':
        separator = '\t'
        encoding = 'utf-8'
        quoting = csv.QUOTE_NONE
        dataframe = _sanitize_for_tsv(dataframe)
    else:  # defaults to csv
        separator = ','
        encoding = 'utf-8-sig'
        quoting = csv.QUOTE_NONNUMERIC

    with open(destination_path, 'w', encoding=encoding, newline='') as f:
        dataframe.to_csv(
            f,
            index=False,
            escapechar='\\',
            sep=separator,
            quoting=quoting,
            lineterminator='\r\n',
            na_rep='',
            float_format=_make_float_formatter(number_format),
        )


def create_report_file_pandas(
    dataframe_generator: Generator[DataFrame, None, None],
    total_rows: int,
    process_dataframe: Callable,
    destination_path: str,
    file_type: str = 'csv',
    number_format: NumberFormat | None = None,
) -> str:
    """Create the report file (using Pandas).

    Args:
        dataframe_generator (Generator): Generator that returns DataFrame objects.
        total_rows (int): Total rows of all the entries.
        process_dataframe (Callable): Function used to process the dataframe.
        destination_path (str): Path to store the file at.
        file_type (str): The type of file.
        number_format (NumberFormat): The number format (US/EU).

    Returns:
        str: Final path the file was stored at.
    """
    file_path, extension = path.splitext(destination_path)
    zip_path = f'{file_path}.zip'
    logger.info(f'Creating zipped report file: {zip_path}')

    def add_data_file(dataframe: DataFrame, file_num: int, zip_file: ZipFile) -> None:
        """Add data from the data frame to the ZIP file.

        Args:
            dataframe (DataFrame): Data frame with data to add.
            file_num (int): The number of the file in the sequence.
            zip_file (ZipFile): ZIP file to add the data file to.
        """
        memory_buffer = io.BytesIO()
        file_name = file_path.split('/')[-1]
        temp_file = f'{file_name}_{file_num}{extension}'
        _write_file_pandas_polars(
            dataframe, file_type, number_format or NumberFormat.US, memory_buffer
        )
        zip_file.writestr(temp_file, memory_buffer.getvalue())

    dataframe = None
    file_num = 1
    with ZipFile(zip_path, 'w', ZIP_DEFLATED) as zip_file:
        for df in dataframe_generator:
            process_dataframe(df)
            dataframe = pandas.concat([dataframe, df], ignore_index=True)

            if len(dataframe) > MAX_ROWS_PER_FILE:
                add_data_file(dataframe[0:MAX_ROWS_PER_FILE], file_num, zip_file)
                dataframe = dataframe[MAX_ROWS_PER_FILE:]
                file_num += 1

        if dataframe is not None and len(dataframe) >= 0:
            add_data_file(dataframe, file_num, zip_file)

    return zip_path


def create_report_file(
    headers: list,
    iterator: Iterator[dict],
    total_rows: int,
    destination_path: str,
    row_processor: Callable[[dict], list] | None = None,
    file_type: str = 'csv',
) -> str:
    """Create the report file based on headers and an iterator.

    Args:
        headers (list): Headers to include in the file.
        iterator (Iterator): Iterator for the rows of data.
        total_rows (int): Total rows in the iterator.
        destination_path (str): Path to write to.
        row_processor (Callable): Optional callable to process the row data.
        file_type (str): Type of file to write ('csv' or 'txt').

    Returns:
        str: Final destination of the file.
    """
    file_path, extension = path.splitext(destination_path)
    zip_path = f'{file_path}.zip'
    logger.info(f'Creating zipped report file: {zip_path}')

    with ZipFile(zip_path, 'w', ZIP_DEFLATED) as zip_file:
        for i in range(1, math.ceil(total_rows / MAX_ROWS_PER_FILE) + 1):
            target_file = f'{file_path}_{i}{extension}'
            logger.info(f'Creating temporary file: {target_file}')
            _write_file(headers, iterator, MAX_ROWS_PER_FILE, target_file, row_processor, file_type)
            zip_file.write(target_file, path.basename(target_file))
            os.remove(target_file)

    return zip_path


def process_transactions(
    ledger_vat_data: list, statement_period: dict, transactions: list, flip_signs: bool = False
) -> None:
    """Process transactions and groups.

    Args:
        ledger_vat_data (list): list of ledger vat records
        statement_period (dict): statement period dict.
        transactions (list): list of transaction to process.
        flip_signs (bool): bool flag to check if to flip the signs.

    Returns:
        None
    """
    for item in ledger_vat_data:
        base_amount = as_decimal(item['base_amount_payee_currency'])
        if flip_signs:
            base_amount *= -1

        if 'description' in item and item['description']:
            description = item['description']
        else:
            description = '{vat_category} {statement_period}'.format(
                vat_category=VAT_CATEGORY_DESCRIPTIONS[item['vat_category']],
                statement_period=statement_period['statement_period_name'],
            )
        transaction = {
            'description': description,
            'quantity': 1,
            'payee_currency_code': item['payee_currency_code'],
            'amount': base_amount,
            'amount_ex_tax': base_amount,
            'vat_rate': as_decimal(item['vat_rate']),
            'wht_rate': as_decimal(item['wht_rate']),
        }
        transactions.append(transaction)


def extract_vat_info(ledger_vat_data: list, flip_signs: bool) -> dict:
    """Extract VAT information from VAT ledger data.

    Args:
        ledger_vat_data (list): The ledger VAT data
        flip_signs (bool): Whether to flip the signs for distro invoices
    Returns:
        dict: VAT info used when generating documents
    """
    base_amount_total = as_decimal(0)
    payee_total_vat = as_decimal(0)
    vat_total_vat = as_decimal(0)
    has_vat = False

    vat_items = {}
    for item in ledger_vat_data:
        base_amount_payee_currency = as_decimal(item['base_amount_payee_currency'])
        vat_amount_payee_currency = as_decimal(item['vat_amount_payee_currency'])
        vat_amount_vat_currency = as_decimal(item['vat_amount_vat_currency'])
        vat_rate = as_decimal(item['vat_rate'])

        has_vat_amount_payee = item.get('vat_amount_payee_currency') is not None
        has_vat_amount_vat = item.get('vat_amount_vat_currency') is not None
        has_vat_rate = item.get('vat_rate') is not None
        if has_vat_amount_payee or has_vat_amount_vat or has_vat_rate:
            has_vat = True

        if flip_signs:
            base_amount_payee_currency = -1 * base_amount_payee_currency
            vat_amount_payee_currency = -1 * vat_amount_payee_currency
            vat_amount_vat_currency = -1 * vat_amount_vat_currency

        item_key = f'{vat_rate}-{item["vat_currency_code"]}-{item["payee_currency_code"]}'
        if item_key not in vat_items:
            vat_items[item_key] = {
                'base_amount_payee_currency': Decimal(0),
                'vat_rate': vat_rate,
                'vat_amount_payee_currency': Decimal(0),
                'vat_amount_vat_currency': Decimal(0),
                'vat_currency_code': item['vat_currency_code'],
                'payee_currency_code': item['payee_currency_code'],
            }

        vat_items[item_key]['base_amount_payee_currency'] += base_amount_payee_currency
        vat_items[item_key]['vat_amount_payee_currency'] += vat_amount_payee_currency
        vat_items[item_key]['vat_amount_vat_currency'] += vat_amount_vat_currency

        base_amount_total += base_amount_payee_currency
        payee_total_vat += vat_amount_payee_currency
        vat_total_vat += vat_amount_vat_currency

    return {
        'vat_items': list(vat_items.values()),
        'base_amount_total': base_amount_total,
        'payee_total_vat': payee_total_vat,
        'vat_total_vat': vat_total_vat,
        'has_vat': has_vat,
    }


def extract_wht_info(ledger_vat_data: list, flip_signs: bool) -> dict:
    """Extract WHT information from VAT ledger data.

    Args:
        ledger_vat_data (list): The ledger VAT data
        flip_signs (bool): Whether to flip the signs for distro invoices
    Returns:
        dict: WHT info used when generating documents
    """
    total_wht_amount_payee_currency = as_decimal(0)
    total_wht_amount_vat_currency = as_decimal(0)
    has_wht = False

    wht_items = {}
    for item in ledger_vat_data:
        base_amount_payee_currency = as_decimal(item['base_amount_payee_currency'])
        wht_amount_payee_currency = as_decimal(item['wht_amount_payee_currency'])
        wht_amount_vat_currency = as_decimal(item['wht_amount_vat_currency'])
        wht_rate = as_decimal(item['wht_rate'])

        has_wht_amount_payee = item.get('wht_amount_payee_currency') is not None
        has_wht_amount_vat = item.get('wht_amount_vat_currency') is not None
        has_wht_rate = item.get('wht_rate') is not None
        if has_wht_amount_payee or has_wht_amount_vat or has_wht_rate:
            has_wht = True

        if flip_signs:
            base_amount_payee_currency = -1 * base_amount_payee_currency
            wht_amount_payee_currency = -1 * wht_amount_payee_currency
            wht_amount_vat_currency = -1 * wht_amount_vat_currency

        item_key = f'{wht_rate}-{item["vat_currency_code"]}-{item["payee_currency_code"]}'
        if item_key not in wht_items:
            wht_items[item_key] = {
                'base_amount_payee_currency': Decimal(0),
                'wht_rate': wht_rate,
                'wht_amount_payee_currency': Decimal(0),
                'wht_amount_vat_currency': Decimal(0),
                'vat_currency_code': item['vat_currency_code'],
                'payee_currency_code': item['payee_currency_code'],
            }

        wht_items[item_key]['base_amount_payee_currency'] += base_amount_payee_currency
        wht_items[item_key]['wht_amount_payee_currency'] += wht_amount_payee_currency
        wht_items[item_key]['wht_amount_vat_currency'] += wht_amount_vat_currency

        total_wht_amount_payee_currency += wht_amount_payee_currency
        total_wht_amount_vat_currency += wht_amount_vat_currency

    return {
        'wht_items': list(wht_items.values()),
        'total_wht_amount_payee_currency': total_wht_amount_payee_currency,
        'total_wht_amount_vat_currency': total_wht_amount_vat_currency,
        'has_wht': has_wht,
    }


def format_address(address: dict) -> str:
    """Convert an address dict to a string.

    Args:
        address (dict): Address coming from the tax info
    Returns:
        str: A string representing the address
    """
    desired_fields = [
        'address_1',
        'address_2',
        'city',
        'zip',
        'province',
        'country_code',
    ]

    result = []
    for field in desired_fields:
        if field in address and address[field] is not None and address[field].strip() != '':
            result.append(address[field].strip())

    return ' '.join(result)


def generate_pdf_file_playwright(
    destination_path: str,
    template_file: str,
    footer_template_file: str,
    data: dict,
    footer_margin: str = '0.5in',
) -> None:
    """Generate a PDF file using Playwright.

    Args:
        destination_path (str): Path to store the generated file at.
        template_file (str): Path to the template file.
        footer_template_file (str): Path to the template file for the footer.
        data (dict): Data to send into the template.
        footer_margin (str): Margin for the footer.

    Returns:
        None
    """
    html_content = render_template(template_file, data)
    footer_content = render_template(footer_template_file, data)

    with sync_playwright() as p:
        browser = p.chromium.launch()
        page = browser.new_page()
        page.set_content(html_content, wait_until='networkidle')
        page.add_style_tag(path='src/documents/assets/styles_playwright.css')
        page.evaluate('async () => { await document.fonts.ready; }')  # force fonts to be loaded
        page.pdf(
            display_header_footer=True,
            path=destination_path,
            format='A4',
            print_background=True,
            header_template='<div></div>',
            footer_template=footer_content,
            margin={
                'top': '0.5in',
                'right': '0.75in',
                'bottom': footer_margin,
                'left': '0.75in',
            },
        )
        browser.close()


def generate_pdf_file(
    destination_path: str,
    template_file: str,
    data: dict,
    page_options: dict = {},
) -> None:
    """Generate a PDF file at the specified destination given a template, data, and page options.

    Args:
        destination_path (str): Path to store the generated file at.
        template_file (str): Path to the template file.
        data (dict): Data to send into the template.
        page_options (dict): Optional page options for generating the file.

    Returns:
        None
    """
    default_options = {
        'page-size': 'A4',
        'encoding': 'UTF-8',
        'margin-top': '0.5in',
        'margin-right': '0.75in',
        'margin-bottom': '0.5in',
        'margin-left': '0.75in',
        'enable-local-file-access': True,
    }

    pdfkit.from_string(
        render_template(template_file, data),
        output_path=destination_path,
        configuration=pdfkit.configuration(wkhtmltopdf=config.WKHTMLTOPDF_BIN_PATH),
        options={**default_options, **page_options},
        css='src/documents/assets/styles.css',
    )


def get_logo(sap_id: str) -> str:
    """Return the appropriate logo for a SAP ID.

    Args:
        sap_id (str): ID to get the logo for.

    Returns:
        str: Path to the logo image.

    Raises:
        LambdaException: If an appropriate logo doesn't exist.
    """
    if sap_id in config.ALTAFONTE_SAP_IDS:
        image_url = 'https://cdn.theorchard.io/assets/altafonte/icons/brand-text.png'
    elif sap_id in config.AWAL_SAP_IDS:
        image_url = 'https://cdn.theorchard.io/assets/awal/icons/brand-text.png'
    elif sap_id in config.KNR_SAP_IDS:
        image_url = 'https://cdn.theorchard.io/assets/knr/icons/brand-text.png'
    elif sap_id in config.ORCHARD_SAP_IDS:
        image_url = 'https://cdn.theorchard.io/assets/orchard/icons/brand-text.png'
    else:
        raise LambdaException(f'No logo for SAP ID: {sap_id}')

    return image_url


def render_template(template: str, data: dict) -> str:
    """Render a Jinja template.

    Args:
        template (str): Path to the template file
        data (dict): Data to send to the template
    Returns:
        str: Content of the rendered template
    """
    env = jinja2.Environment(
        loader=jinja2.FileSystemLoader('src/documents/templates'),
        autoescape=jinja2.select_autoescape(),
    )
    tmpl = env.get_template(template)
    return tmpl.render(**data)


def sanitize_filename(text: str) -> str:
    """Sanitize text meant to be used as a filename.

    Args:
        text (str): Text to sanitize
    Returns:
        str: Sanitized text
    """
    return re.sub(r'[^\w_-]', '_', text)


def sanitize_value(value: Any, apply_rounding: bool = False) -> Any:
    """Process a document value, converting it if necessary.

    Args:
        value (Any): Value to (potentially) convert
        apply_rounding (bool): Whether to apply rounding to the value
    Returns:
        Any: Processed value
    """
    if value == 0:
        return 0  # converts Decimal/float/whatever to an int

    if isinstance(value, float) and apply_rounding:
        return round(value, config.ROUNDING_DECIMAL_PLACES)

    if isinstance(value, Decimal):
        if apply_rounding:
            value = round(value, config.ROUNDING_DECIMAL_PLACES)
        return value.normalize()

    return value