from uuid import UUID

import pytest
from pydantic import UUID4

from product_staging.models import (
    bulk_session as bulk_session_model,
    bulk_session_ingestion as bulk_session_ingestion_model,
    bulk_session_ingestion_execution as bulk_session_ingestion_execution_model,
    bulk_session_ingestion_project as bulk_session_ingestion_project_model,
    bulk_session_ingestion_product as bulk_session_ingestion_product_model,
)

json_file_name = "test.json"


@pytest.mark.asyncio
async def create_bulk_ingestion(vendor_uuid: UUID, identity_uuid: UUID4, session):
    # Create a bulk session
    bulk_session = await bulk_session_model.create_bulk_session(
        vendor_uuid=vendor_uuid, identity_uuid=identity_uuid, session=session
    )
    assert bulk_session

    # Create a bulk session ingestion
    ingestion = await bulk_session_ingestion_model.create_bulk_session_ingestion(
        bulk_session_id=bulk_session.bulk_session_id,
        assets_required=True,
        submit_products=True,
        identity_uuid=identity_uuid,
        json_file_name=json_file_name,
        send_notifications=True,
        session=session,
    )
    assert ingestion

    return ingestion, bulk_session


@pytest.mark.asyncio
async def create_multiple_bulk_ingestion(
    vendor_uuid: UUID, identity_uuid: UUID4, session
):
    ingestion0, session0 = await create_bulk_ingestion(
        vendor_uuid, identity_uuid, session
    )
    ingestion1 = await bulk_session_ingestion_model.create_bulk_session_ingestion(
        bulk_session_id=session0.bulk_session_id,
        assets_required=True,
        submit_products=True,
        identity_uuid=identity_uuid,
        json_file_name=json_file_name,
        send_notifications=True,
        session=session,
    )
    assert ingestion1
    ingestion2, session1 = await create_bulk_ingestion(
        vendor_uuid, identity_uuid, session
    )

    return [ingestion0, ingestion1, ingestion2], [session0, session1]


@pytest.mark.asyncio
async def create_bulk_ingestion_project(
    vendor_uuid: UUID, identity_uuid: UUID4, session
):
    ingestion, bulk_session = await create_bulk_ingestion(
        vendor_uuid, identity_uuid, session
    )
    project = await bulk_session_ingestion_project_model.create_bulk_session_ingestion_project(
        bulk_session_ingestion_id=ingestion["bulk_session_ingestion_id"],
        project_code="TEST_PROJECT",
        vendor_uuid=vendor_uuid,
        ingestion_status="in_progress",
        session=session,
    )
    assert project

    return project, ingestion, bulk_session


@pytest.mark.asyncio
async def create_bulk_ingestion_product(
    vendor_uuid: UUID, identity_uuid: UUID4, session
):
    project, ingestion, bulk_session = await create_bulk_ingestion_project(
        vendor_uuid, identity_uuid, session
    )
    product = await bulk_session_ingestion_product_model.create_bulk_session_ingestion_product(
        bulk_session_ingestion_id=ingestion["bulk_session_ingestion_id"],
        bulk_session_ingestion_project_id=project["bulk_session_ingestion_project_id"],
        product_id=123,
        upc="456",
        product_code="ABC",
        session=session,
    )
    assert product

    return product, project, ingestion, bulk_session


@pytest.mark.asyncio
async def create_bulk_ingestion_execution(
    vendor_uuid: UUID, identity_uuid: UUID4, session
):
    ingestion, _ = await create_bulk_ingestion(vendor_uuid, identity_uuid, session)

    # Create an execution
    execution_arn = (
        "arn:aws:states:us-east-1:123456789012:execution:MyStateMachine:my-execution"
    )
    execution = await bulk_session_ingestion_execution_model.upsert_bulk_session_ingestion_execution(
        bulk_session_ingestion_id=ingestion["bulk_session_ingestion_id"],
        execution_arn=execution_arn,
        session=session,
    )
    assert execution

    return ingestion, execution, execution_arn