"""Lambda function module."""
import uuid
from typing import Dict

from ddex_ingester_common.helpers.asset import get_product_track_assets
from ddex_ingester_common.lambda_exceptions import AudioException
from ddex_ingester_common.logging import utils as logging_utils
from ddex_ingester_common.models.state_machine import \
    product as StateMachineProduct
from ddex_ingester_common.schemas.state_machine_schema import (
    StateMachineSchema, TrackSchema)

import config
from config import graphql_gateway
from constants.asset import ENCODING_COMPLETED

logger = logging_utils.get_logger(config.app_logger)


def handler(event, context):
    """Lambda entrypoint."""
    context = StateMachineSchema().load(event.get('context'))
    track = TrackSchema().load(event.get('track'))

    correlation_id = context.correlation_id or str(uuid.uuid4())
    context.correlation_id = correlation_id
    logging_utils.update_logger_correlation_id(logger, correlation_id)
    logging_utils.update_logger_with_message_ids(
        logger,
        context.message_id,
        context.message_thread_id,
        context.execution_name
    )
    graphql_gateway.set_headers(
        {
            'Orchard-User-Id': config.OA_USER,
            'Correlation-Id': correlation_id,
        }
    )
    # Call ows assets for status of uploaded track asset
    status = get_asset_status(context.product, track.tuid)
    logger.info(f'Asset finished processing: {status}')

    return StateMachineSchema().dump(context)


def get_asset_status(product: StateMachineProduct, tuid: int) -> Dict:
    """Get track asset status."""
    # Check if it's enabled at a vendor or subaccount level
    return _v2_asset_status(product.product_id, tuid)


def _v2_asset_status(product_id: int, tuid: int) -> Dict:
    """V2 flow for asset status checking."""
    assets = get_product_track_assets(
        graphql_gateway, product_id, True
    )

    status = next(
        (asset for asset in assets
         if asset.get('trackUniqueId') == tuid),
        None
    )

    # Note: V2 endpoint does not have a correction field.
    # Replacement asset replaces the original in the GraphQL response.

    logger.info(f'Audio asset status: {status}')
    if not status or status.get('status') != ENCODING_COMPLETED:
        raise AudioException(f'Asset status is not {ENCODING_COMPLETED}')
    return status