"""Lambda update_session function module."""

from datetime import datetime
from typing import Any, Literal

from common.connectors.graphql import obo_graphql as graphql
from common.schemas.ingestion import BaseEvent

import config
from src.exceptions import IngestionException


def handle(
    event: BaseEvent,
) -> BaseEvent:
    final_ingestion_status: Literal["in_progress", "failure", "success"] = "success"
    # If there is already a failure outside the product loop, set ingestion status to failure
    if event.ingestion_status == "failure":
        final_ingestion_status = "failure"

    graphql.init_graphql_client(
        environment=config.ENVIRONMENT,
        service_name=config.M2M_APPLICATION_NAME,
        graphql_service_name=config.GRAPHQL_SERVICE_NAME,
        identity_id=event.identity_uuid,
        profile_id=config.PROFILE_ID,
        profile_type=config.PROFILE_TYPE,
        headers={
            "Correlation-Id": event.correlation_id,
            "Orchard-Roles": config.ROLE,
        },
    )

    # Check for failures within the products
    if final_ingestion_status != "failure":
        failed_products_response = graphql.get_bulk_session_ingestion_products(
            bulk_session_ingestion_id=event.bulk_session_ingestion_id,
            product_ingestion_status="failure",
            limit=1,
        )
        if failed_products_response and failed_products_response.products:
            # If any product has a failure status, set the ingestion status to failure
            final_ingestion_status = "failure"

    # If no products ingested successfully we failed while writing those failures
    if final_ingestion_status != "failure":
        success_products_response = graphql.get_bulk_session_ingestion_products(
            bulk_session_ingestion_id=event.bulk_session_ingestion_id,
            product_ingestion_status="success",
            limit=1,
        )
        if not success_products_response or not success_products_response.products:
            final_ingestion_status = "failure"

    # Update the bulk session ingestion status
    update_response = graphql.update_bulk_session_ingestion(
        bulk_session_ingestion_id=event.bulk_session_ingestion_id,
        ingestion_status=final_ingestion_status,
        completed_on=datetime.now(),
    )
    if update_response.ingestion_status != final_ingestion_status:
        raise ValueError("Failed to update bulk session ingestion status")

    # Raise an exception so that the entire step function will fail if there is any failure
    # This is necessary because failures in the map step are not propagated to the parent step
    if final_ingestion_status == "failure":
        raise IngestionException(
            "The ingestion process has failed for at least one product or step."
        )

    return BaseEvent(
        **{
            **event.model_dump(),
            **update_response.model_dump(exclude_unset=True),
        }
    )


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.
    """
    return handle(BaseEvent(**event_data)).model_dump()