import asyncio

import pandas as pd
from box import Box
from fastapi import APIRouter
from fastapi.responses import JSONResponse, Response
from unidecode import unidecode

from ... import logger, models
from ...logic.reporting import ReportExporter
from ...models import AuditReport
from ...utils.xlsx import df_to_xlsx_bytes, format_output_xlsx

logger = logger.new_logger(__name__)


ROUTE: str = "/audits/export"


def create_router(state: Box, *args, **kwargs):
    """Create FastAPI router."""

    app = APIRouter()

    @app.post("/xlsx", response_class=JSONResponse)
    async def export_rows_to_xlsx(export: models.Export):
        """Export list of rows (dicts) to downloadable Excel byte stream.

        Returns:
            Response: Excel byte stream, ideal for downloading.
        """
        logger.debug(
            "Exporting {} rows to Excel (in {} sheets)...",
            sum(len(sheet.rows) for sheet in export.sheets),
            len(export.sheets),
        )

        def worker():
            sheets_as_df = []

            for i, sheet in enumerate(export.sheets):
                if columns := sheet.columns:
                    # A column order and display logic has been provided
                    # in the request, mapping column keys in the input
                    # to labels in the output.
                    col_keys, col_labels = zip(*columns)
                    col_keys = list(col_keys)
                    col_labels = list(col_labels)
                else:
                    col_keys = None
                    col_labels = None

                df = pd.DataFrame(sheet.rows, columns=col_keys)

                date_cols = filter(lambda x: "date_" in x.lower(), df.columns)

                for col in date_cols:
                    # Convert to datetime format if not already and
                    # format as a string in the desired format
                    df[col] = pd.to_datetime(df[col]).apply(
                        lambda x: (
                            x.strftime("%Y-%m-%d %H:%M:%S") if pd.notnull(x) else x
                        )
                    )

                if col_labels:
                    # Convert the column keys into the desired labels. This must be done
                    # right before exporting to Excel, as the column keys may be used for
                    # data manipulation and filtering.
                    df.columns = col_labels

                sheet_name = sheet.name or f"Sheet {i + 1}"
                sheets_as_df.append((sheet_name, df))

            output = df_to_xlsx_bytes(sheets_as_df)

            return format_output_xlsx(output)

        # Use non-blocking IO to generate the Excel file (useful for preventing
        # blocking the main thread in case of large exports; i.e. datasets of
        # thousands or tens of thousands of rows).
        formatted_excel_bytes = await asyncio.to_thread(worker)

        return Response(
            content=formatted_excel_bytes, media_type="application/octet-stream"
        )

    @app.post("/audit_report", response_class=JSONResponse)
    async def audit_report(data: AuditReport):
        """Generate an audit report.

        The returned response is either a PPTX or a PDF file, depending on the
        configuration provided in the request. It also includes the file name
        in the response headers.
        """
        exporter = ReportExporter(state.db, data)
        content, media_type, file_name = await exporter.export()
        response = Response(content, media_type=media_type)
        # File name in the headers must be ASCII safe,
        # so ensure non-ASCII characters are replaced
        response.headers["X-Custom-File-Name"] = unidecode(file_name)
        return response

    return app