"""Logic for generating a re-upload spreadsheet with failed products."""

import io
import logging
import uuid as uuid_mod
from uuid import UUID
from zipfile import ZIP_DEFLATED, ZipFile

from common.connectors import s3
from common.connectors.ows_product_staging_obo import (
    get_bulk_session,
    get_bulk_session_ingestion_products,
    get_bulk_session_metadata_file,
    update_bulk_session_ingestion_failure_file,
)
from common.constants.spreadsheet import columns
from common.constants.spreadsheet.sheets import (
    PRODUCTS,
    PROJECTS,
    WORKBOOK_FORMAT,
)
from common.schemas.s3_reference import S3Reference
from common.utils.openpyxl import load_workbook_resilient, maximum_row
from openpyxl import Workbook
from openpyxl.writer.excel import ExcelWriter

import config

logger = logging.getLogger(__name__)

MIN_ROW = 3

_PRODUCTS_COLS = WORKBOOK_FORMAT[PRODUCTS]
_PROJECTS_COLS = WORKBOOK_FORMAT[PROJECTS]

PRODUCT_CODE_INDEX_IN_PRODUCTS = _PRODUCTS_COLS.index(columns.PRODUCT_CODE)
PROJECT_CODE_INDEX_IN_PRODUCTS = _PRODUCTS_COLS.index(columns.PROJECT_CODE)
PROJECT_CODE_INDEX_IN_PROJECTS = _PROJECTS_COLS.index(columns.PROJECT_CODE)
PRODUCT_CODE_INDEX_DEFAULT = 0  # All other sheets have Product Code in column A


def _get_failed_product_codes(
    bulk_session_ingestion_id: str,
    identity_uuid: str,
) -> list[str]:
    """Fetch failed product codes via the REST API."""
    products = get_bulk_session_ingestion_products(
        bulk_session_ingestion_id=UUID(bulk_session_ingestion_id),
        impersonated_identity_uuid=identity_uuid,
        ingestion_status="failure",
    )
    return [p["product_code"] for p in products if p.get("product_code")]


def _get_filter_params(
    sheet_name: str,
    failed_product_codes: set[str],
    failed_project_codes: set[str],
) -> tuple[int, set[str]]:
    """Return (filter_column_index, match_codes) for a given sheet."""
    if sheet_name == PROJECTS:
        return PROJECT_CODE_INDEX_IN_PROJECTS, failed_project_codes
    if sheet_name == PRODUCTS:
        return PRODUCT_CODE_INDEX_IN_PRODUCTS, failed_product_codes
    return PRODUCT_CODE_INDEX_DEFAULT, failed_product_codes


def _collect_failed_project_codes(
    wb: Workbook,
    failed_product_codes: set[str],
) -> set[str]:
    """Scan the Products sheet and return project codes used by failed products."""
    result: set[str] = set()
    if PRODUCTS not in wb.sheetnames:
        return result

    products_ws = wb[PRODUCTS]
    for row in products_ws.iter_rows(
        min_row=MIN_ROW,
        max_row=maximum_row(products_ws),
    ):
        product_code = row[PRODUCT_CODE_INDEX_IN_PRODUCTS].value
        project_code = row[PROJECT_CODE_INDEX_IN_PRODUCTS].value
        if (
            product_code is not None
            and str(product_code) in failed_product_codes
            and project_code is not None
        ):
            result.add(str(project_code))
    return result


def _filter_workbook(
    original_stream: io.BytesIO,
    failed_product_codes: set[str],
) -> Workbook:
    """Filter a workbook in-place, deleting rows that don't match failed products.

    Loads the original workbook and removes non-matching data rows from each sheet,
    preserving all formatting, styles, data validations, column widths, etc.

    Args:
        original_stream: BytesIO stream of the original XLSX file.
        failed_product_codes: Set of product codes that failed ingestion.

    Returns:
        The modified Workbook containing only the rows for failed products.
    """
    wb = load_workbook_resilient(original_stream, data_only=True)
    failed_project_codes = _collect_failed_project_codes(wb, failed_product_codes)

    for sheet_name in WORKBOOK_FORMAT:
        if sheet_name not in wb.sheetnames:
            continue

        ws = wb[sheet_name]
        filter_col_index, match_codes = _get_filter_params(
            sheet_name, failed_product_codes, failed_project_codes
        )

        max_row = maximum_row(ws)
        for row_idx in range(max_row, MIN_ROW - 1, -1):
            cell_value = ws.cell(row=row_idx, column=filter_col_index + 1).value
            if cell_value is None or str(cell_value) not in match_codes:
                ws.delete_rows(row_idx)

    return wb


def _workbook_to_stream(workbook: Workbook) -> io.BytesIO:
    """Serialize a workbook to a BytesIO stream."""
    output = io.BytesIO()
    ExcelWriter(workbook, ZipFile(output, "w", ZIP_DEFLATED, allowZip64=True)).save()
    return output


def generate_reupload_spreadsheet(
    *,
    bulk_session_id: str,
    bulk_session_ingestion_id: str,
    identity_uuid: str,
) -> None:
    """Generate a re-upload spreadsheet containing only the failed products."""
    failed_product_codes = _get_failed_product_codes(
        bulk_session_ingestion_id, identity_uuid
    )
    if not failed_product_codes:
        logger.info(
            "No failed products found, skipping re-upload spreadsheet generation"
        )
        return

    logger.info(
        "Found %d failed products, generating re-upload spreadsheet",
        len(failed_product_codes),
    )

    session = get_bulk_session(
        bulk_session_id=UUID(bulk_session_id),
        impersonated_identity_uuid=identity_uuid,
    )

    if not session.latest_upload_file_id:
        logger.warning(
            "No latest_upload_file_id in bulk session, skipping re-upload generation"
        )
        return

    metadata_file = get_bulk_session_metadata_file(
        bulk_session_id=UUID(bulk_session_id),
        bulk_session_metadata_file_id=session.latest_upload_file_id,
        impersonated_identity_uuid=identity_uuid,
    )

    original_stream = s3.get_object_stream(
        S3Reference(bucket=metadata_file.s3_bucket, key=metadata_file.s3_filename)
    )

    filtered_wb = _filter_workbook(original_stream, set(failed_product_codes))
    filtered_stream = _workbook_to_stream(filtered_wb)

    failure_s3_key = f"ingestion_reupload_spreadsheets/{uuid_mod.uuid4()}"
    s3.put_object(
        S3Reference(bucket=config.S3_BUCKET, key=failure_s3_key),
        filtered_stream,
    )

    update_bulk_session_ingestion_failure_file(
        bulk_session_ingestion_id=UUID(bulk_session_ingestion_id),
        ingestion_failure_file=failure_s3_key,
        impersonated_identity_uuid=identity_uuid,
    )

    logger.info(
        "Re-upload spreadsheet generated and uploaded to S3: %s",
        failure_s3_key,
    )