from collections.abc import Generator
from contextlib import contextmanager
from dataclasses import dataclass
from typing import Any

from common.connectors.graphql import obo_graphql as graphql_common


@dataclass
class IngestionContext:
    """Typed context yielded by project_ingestion.

    Attributes:
        project_id: optional project id created during ingestion
        bulk_session_ingestion_project_id: optional id returned from upsert
            operations for the bulk_session_ingestion_project row
    """

    project_id: int | None = None
    bulk_session_ingestion_project_id: str | None = None


@contextmanager
def project_ingestion(
    *,
    bulk_session_ingestion_id: str,
    vendor_uuid: str,
    project_code: str,
) -> Generator[IngestionContext]:
    """
    Generator function that yields to a block within project ingestion
    status mutations.

    If the block completes successfully, the ingestion status is set to
    'success'. If an exception is raised, the ingestion status is set to
    'failure' and it is re-raised.

    Args:
        bulk_session_ingestion_id: the bulk session ingestion ID
        vendor_uuid: the vendor UUID for the project
        project_code: the project code for the project

    Raises:
        Exception: re-raises any exception from the block
    """
    try:
        ingestion_project = graphql_common.upsert_bulk_session_ingestion_project(
            bulk_session_ingestion_id=bulk_session_ingestion_id,
            project_code=project_code,
            vendor_uuid=vendor_uuid,
            ingestion_status="in_progress",
        )

        ctx = IngestionContext()
        if ingestion_project is not None:
            ctx.bulk_session_ingestion_project_id = ingestion_project.id

        yield ctx

        # On successful completion, include optional ctx data in the success update
        success_kwargs: dict[str, Any] = {
            "bulk_session_ingestion_id": bulk_session_ingestion_id,
            "project_code": project_code,
            "vendor_uuid": vendor_uuid,
            "ingestion_status": "success",
        }
        # If caller provided a project_id, include it for the success mutation
        if ctx.project_id is not None:
            success_kwargs["project_id"] = ctx.project_id

        success_result = graphql_common.upsert_bulk_session_ingestion_project(
            **success_kwargs
        )
        if success_result is not None:
            ctx.bulk_session_ingestion_project_id = success_result.id

    except Exception as ex:
        graphql_common.upsert_bulk_session_ingestion_project(
            bulk_session_ingestion_id=bulk_session_ingestion_id,
            project_code=project_code,
            vendor_uuid=vendor_uuid,
            ingestion_status="failure",
        )
        raise ex