"""Lambda set_tracks function module."""

from typing import Any

from common.connectors.exceptions import NetworkException
from common.connectors.graphql.obo_graphql import upsert_bulk_session_ingestion_track
from common.schemas.ingestion import Event, ProductInfo

from src.connectors import graphql
from src.constants import KNOWN_RETRY_ERRORS


def handle(
    product_info: ProductInfo, execution_arn: str, correlation_id: str, identity_id: str
) -> ProductInfo:
    graphql.init(correlation_id, identity_id)
    product_id = product_info.product.product_id
    upc = product_info.product.upc
    vendor_id = product_info.product.vendor_id
    subaccount_id = product_info.product.subaccount_id
    tracks = product_info.tracks

    if not product_id or not upc:
        raise ValueError("Product ID and UPC are required")

    if not tracks:
        raise ValueError("Tracks are required")

    if not product_info.product.ingestion_product_id:
        raise ValueError("Ingestion Product ID is required")

    for track in tracks or []:
        upsert_bulk_session_ingestion_track(
            bulk_session_ingestion_product_id=product_info.product.ingestion_product_id,
            execution_arn=execution_arn,
            volume=track.volume,
            sequence_number=track.sequence_number,
        )

    updated_product_info = product_info.copy(
        update={
            "tracks": graphql.set_product_tracks(
                product_id=product_id,
                tracks=tracks,
            )
        }
    )

    try:
        # TODO Remove once setProductTracks does everything that saveTracks does
        for track in updated_product_info.tracks or []:
            graphql.update_track(upc, track)

        for track in updated_product_info.tracks or []:
            graphql.update_track_personnel(track, vendor_id, subaccount_id)

        for track in updated_product_info.tracks or []:
            upsert_bulk_session_ingestion_track(
                bulk_session_ingestion_product_id=product_info.product.ingestion_product_id,
                execution_arn=execution_arn,
                volume=track.volume,
                sequence_number=track.sequence_number,
                ingestion_status="success",
            )

        return updated_product_info

    except Exception as exc:
        handle_exception(
            exception=exc,
            execution_arn=execution_arn,
            product_info=updated_product_info,
        )
        return updated_product_info


def handle_exception(
    *, exception: Exception, execution_arn: str, product_info: ProductInfo
) -> None:
    """Handle exceptions."""
    for error_message in KNOWN_RETRY_ERRORS:
        if error_message in str(exception):
            raise NetworkException(
                f"Known retry error detected: {exception}"
            ) from exception

    for track in product_info.tracks or []:
        if not product_info.product.ingestion_product_id:
            break
        upsert_bulk_session_ingestion_track(
            bulk_session_ingestion_product_id=product_info.product.ingestion_product_id,
            execution_arn=execution_arn,
            volume=track.volume,
            sequence_number=track.sequence_number,
            ingestion_status="failure",
        )
    raise exception


def handler(event_data: dict[str, Any], context: Any) -> dict[str, Any]:
    """
    Lambda entry point.

    Args:
        event_data: Lambda event payload (should look like event.shadow.json)
        context: Lambda context.
    """
    event = Event(**event_data)
    correlation_id = event.correlation_id
    assert event.execution_arn
    event.product_info = handle(
        event.product_info, event.execution_arn, correlation_id, event.identity_uuid
    )
    return event.model_dump()