"""Job Creation Logic."""

import json
import uuid
from typing import Any

import jsonschema
from botocore.exceptions import ClientError

from spec.jsonschema import jobs as jobs_validation
from spec.jsonschema.job_outputs import JOB_TYPE_TO_CREATE_JOB_OUTPUTS_SCHEMA_MAP
from video import config
from video.connectors import mediaconvert, mysql, stepfunctions
from video.constants import job_io_fields, job_statuses, job_types
from video.exceptions import InvalidRequest, JobNotFound
from video.logic import access_control, context as context_logic
from video.models.s3 import cloud_transfer_progress, s3_token
from video.models.sql import queries
from video.models.sql.classes import product_video


class Return(Exception):
    """Helper functions raise this if the calling function should return."""

    pass


def _if_cancelling_workflow_apply_updates(
    manage_job_response: dict[str, Any],
) -> int | None:
    """If cancelling workflow apply updates.

    manage_job_response (dict): Manage job response.

    Returns:
        None or int: If cancelling workflow, the workflow_job_id is returned.
    Raises:
        Return: Indicates calling function should return early.
    """
    stepfunctions_client = stepfunctions.get_stepfunctions_client()
    job_id = manage_job_response.get("id")
    parent_id = manage_job_response.get("parent_id")
    is_root_job = not parent_id

    if not is_root_job:
        return None
    if manage_job_response["status"] != job_statuses.CANCELLED:
        raise Return
    if manage_job_response["type"] not in job_types.WORKFLOWS:
        raise Return

    workflow_job_id = job_id

    try:
        stepfunctions_client.stop_execution(
            executionArn=manage_job_response["inputs"][
                job_io_fields.WORKFLOW_STATE_MACHINE_EXECUTION_ARN
            ]
        )
    except ClientError:
        pass

    return workflow_job_id


def _if_erroring_workflow_child_apply_updates(
    data: dict[str, Any],
    manage_job_response: dict[str, Any],
    job_data_to_persist: list[dict[str, Any]],
    session: Any,
) -> int | None:
    """If erroring workflow child apply updates.

    data (dict): Job management request data.
    manage_job_response (dict): Manage job response.
    job_data_to_persist (dict): Data that needs to be persisted is added here.
    session (db_session): db session.

    Returns:
        None or int: If erroring a workflow child, the workflow_job_id is
        returned.
    Raises:
        Return: Indicates calling function should return early.
    """
    new_status = data.get("status")
    parent_id = manage_job_response.get("parent_id")
    is_root_job = not parent_id

    if is_root_job:
        return None
    if manage_job_response["status"] != job_statuses.ERROR:
        raise Return
    parent_job = queries.get_jobs({"id": parent_id}, session)[0]
    if parent_job["status"] in job_statuses.FINAL_JOB_STATUSES:
        raise Return
    if parent_job["type"] not in job_types.WORKFLOWS:
        raise Return

    workflow_job_id = parent_id

    job_data_to_persist.append({"id": workflow_job_id, "status": new_status})

    return workflow_job_id


def _cancel_all_workflow_children(
    workflow_job_id: int | None,
    job_data_to_persist: list[dict[str, Any]],
    session: Any,
) -> None:
    """Cancel all workflow children.

    workflow_job_id (int): Workflow id.
    job_data_to_persist (dict): Data that needs to be persisted is added here.
    session (db_session): db session.
    """
    mediaconvert_client = mediaconvert.get_mediaconvert_client()

    child_jobs = (
        queries.get_jobs({"parent_id": workflow_job_id}, session)
        if workflow_job_id
        else []
    )

    for child_job in child_jobs:
        if child_job["status"] not in job_statuses.FINAL_JOB_STATUSES:
            job_data_to_persist.append(
                {"id": child_job["id"], "status": job_statuses.CANCELLED}
            )

        if not child_job["inputs"]:
            continue

        for key, value in child_job["inputs"].items():
            if key not in job_io_fields.MEDIACONVERT_JOB_ID_FIELDS:
                continue
            try:
                mediaconvert_client.cancel_job(Id=value)
            except ClientError:
                pass


def if_in_workflow_apply_updates(
    data: dict[str, Any], manage_job_response: dict[str, Any]
) -> None:
    """Manage workflow.

    Cancelling a parent job that is a workflow type will cause all of its
    children that aren't in a final state (CANCELLED, COMPLETE, ERROR) to be
    cancelled. The state machine run will also be cancelled and any
    mediaconvert jobs associated with the workflow will be canceled (this
    cancellation only works if the jobs are in the mediaconvert queue,
    otherwise it doesn't have any effect, in progress mediaconvert jobs cannot
    be cancelled).

    Cancelling a child job whose parent is a workflow type has no cascading
    effects.

    Erroring a child job whose parent is a workflow type will cause all of its
    siblings that aren't in a final state (CANCELLED, COMPLETE, ERROR) to be
    cancelled. The state machine is not stopped in this case as it should
    already have been stopped by a worker calling send_task_failure() which
    would result in the state machine terminating itself. Any mediaconvert jobs
    associated with the workflow will be canceled (this cancellation only works
    if the jobs are in the mediaconvert queue, otherwise it doesn't have any
    effect).

    Erroring a parent job that is a workflow type has no cascading effects.

    data (dict): Job management request data.
    manage_job_response (dict): Manage job response.
    """
    with mysql.db_session() as session:
        job_data_to_persist: list[dict[str, Any]] = []

        try:
            workflow_job_id = _if_cancelling_workflow_apply_updates(
                manage_job_response
            ) or _if_erroring_workflow_child_apply_updates(
                data, manage_job_response, job_data_to_persist, session
            )
        except Return:
            return

        _cancel_all_workflow_children(workflow_job_id, job_data_to_persist, session)

        if job_data_to_persist:
            queries.persist_job_data(job_data_to_persist, session=session)


def manage_job(data: dict[str, Any]) -> dict[str, Any]:
    """Create a job and add inputs, outputs, and statuses."""
    job_id: int | None = data.get("id")
    existing_job: dict[str, Any] | None = None
    if job_id is not None:
        existing_jobs = queries.get_jobs({"id": job_id})
        existing_job = existing_jobs[0] if existing_jobs else None
        if existing_job is None:
            raise JobNotFound()

    check_has_access(data)
    validate_job_management_request_data(data, existing_job)
    manage_job_response = persist_new_job_data(data)

    response_schema = jobs_validation.build_response_schema_jobs(
        manage_job_response["type"]
    )
    jsonschema.validate(manage_job_response, response_schema)

    if existing_job:
        if_in_workflow_apply_updates(data, manage_job_response)

    return manage_job_response


def check_has_access(data: dict[str, Any]) -> None:
    """Check whether user has access or throws error."""
    product_id = data.get("context", {}).get("product_id")
    job_id: int | str | None = data.get("id")
    if product_id:
        context_logic.add_to_context({"product_id": product_id})
        access_control.check_access_to_product(product_id)
    if job_id is not None:
        access_control.check_access_to_job(int(job_id))


def get_jobs(job_data_filters: dict[str, Any]) -> list[dict[str, Any]]:
    """Get jobs.

    Args:
        job_data_filters (dict): Filters for jobs.
    """
    return queries.get_jobs(job_data_filters)


def persist_new_job_data(job: dict[str, Any]) -> dict[str, Any]:
    """Manage a job.

    This includes creating a job, adding inputs and outputs, or changing
    status.

    Args:
        job (dict): A job.

    Returns:
        response.Response: The managed job.
    """
    job_id = queries.persist_job_data(job)[0]

    return queries.get_jobs({"id": job_id})[0]


def validate_job_management_request_data(
    data: dict[str, Any], existing_job: dict[str, Any] | None
) -> None:
    """Validate data sent with a job management request.

    Args:
        data (dict): Job management request data.
        existing_job (dict): Existing job.
    """
    new_id = data.get("id")
    new_status = data.get("status")
    new_type = data.get("type")
    new_inputs = data.get("inputs")
    new_outputs = data.get("outputs")

    if new_id is None:
        if not new_type:
            raise InvalidRequest("Job type is required")
        if new_type not in job_types.JOB_TYPES:
            raise InvalidRequest(f"Unknown job type: {new_type}")
        jsonschema.validate(
            data,
            jobs_validation.build_request_schema_create_job(str(new_type)),
        )
        return

    assert existing_job is not None
    existing_status = existing_job["status"]
    existing_type = existing_job["type"]

    if new_status:
        jsonschema.validate(
            data,
            jobs_validation.build_request_schema_change_job_status(
                existing_type,
                existing_status,
            ),
        )
        return

    if new_inputs:
        jsonschema.validate(
            data,
            jobs_validation.build_request_schema_create_job_inputs(
                existing_type,
            ),
        )
        return

    if new_outputs:
        jsonschema.validate(
            data,
            jobs_validation.build_request_schema_create_job_outputs(
                existing_type,
            ),
        )
        return


def setup_workflow(
    setup_info: dict[str, Any], data: dict[str, Any]
) -> list[dict[str, Any]]:
    """Set up workflow and inputs.

    Args:
        setup_info (dict): Setup info for workflow.
        data (dict): Data to be applied to this workflow.
    """
    workflow_job_type = setup_info["workflow_job_type"]
    workflow_inputs = setup_info.get("workflow_inputs") or {}
    to_setup = setup_info["jobs_to_setup"]
    job_type_to_inputs_function = setup_info.get("job_type_to_inputs", {})

    check_has_access(data)

    with mysql.db_session() as session:
        workflow_job_id = queries.persist_job_data(
            {
                "type": workflow_job_type,
                "inputs": workflow_inputs,
            },
            session=session,
        )[0]
        job_data_to_persist = []

        for job_type in to_setup:
            job_inputs_func = job_type_to_inputs_function.get(job_type, lambda _: {})

            job_data = {
                "parent_id": workflow_job_id,
                "type": job_type,
            }
            job_inputs = job_inputs_func(workflow_job_id)
            if job_inputs:
                job_data["inputs"] = job_inputs

            job_data_to_persist.append(job_data)

        s3_token_expiration = next(
            (
                job_data["inputs"][job_io_fields.S3_TOKEN_EXPIRATION]
                for job_data in job_data_to_persist
                if job_io_fields.S3_TOKEN_EXPIRATION in job_data.get("inputs", {})
            ),
            None,
        )
        if s3_token_expiration:
            workflow_inputs[job_io_fields.S3_TOKEN_EXPIRATION] = s3_token_expiration

        stepfunctions_execution_name = "{}_{}".format(
            str(uuid.uuid1()), workflow_job_id
        )

        stepfunctions_input = {
            **workflow_inputs,
            "workflow_job_type": workflow_job_type,
            "workflow_job_id": workflow_job_id,
        }

        if job_data_to_persist:
            queries.persist_job_data(job_data_to_persist, session=session)

        session.commit()

        state_machine_execution = stepfunctions.start_execution(
            json.dumps(stepfunctions_input),
            stepfunctions_execution_name,
            workflow_job_type,
        )

        state_machine_execution_arn = state_machine_execution["executionArn"]

        if workflow_job_type == job_types.WORKFLOW_APPROVAL:
            product_video.upsert(
                product_video={
                    "release_id": data["context"]["product_id"],
                    "latest_approval_job_id": workflow_job_id,
                }
            )
        elif workflow_job_type == job_types.WORKFLOW_INGESTION_REENCODE:
            # Point the product at the re-encode run so the streamable preview
            # URL and per-run thumbnails (resolved through latest_pipeline_run_id)
            # read the corrected render. The chosen thumbnail is denormalized
            # with the run id, so re-point it for a crop fix or clear it for a
            # re-pick when a clip (trim) fix shifts the timeline.
            product_id = data["context"]["product_id"]
            reencode_update: dict[str, Any] = {
                "release_id": product_id,
                "latest_pipeline_run_id": workflow_job_id,
            }
            reencode_update.update(
                _reencode_thumbnail_updates(
                    product_id=product_id,
                    reencode_job_id=workflow_job_id,
                    overrides=data["overrides"],
                )
            )
            product_video.upsert(product_video=reencode_update)

        queries.persist_job_data(
            {
                "id": workflow_job_id,
                "inputs": {
                    job_io_fields.WORKFLOW_STATE_MACHINE_EXECUTION_ARN: state_machine_execution_arn
                },
            },
            session=session,
        )

        return queries.get_jobs({"parent_id": workflow_job_id})


def get_file_attributes(file_data: dict[str, Any]) -> dict[str, Any]:
    """Add the file attributes.

    Args:
        file_data (dict): Has the filesize and filename.

    Returns:
        (dict): The filesize and the filename.
    """
    file_info = {"filename": file_data["name"], "filesize": file_data["size"]}
    response_schema = jobs_validation.build_response_schema_file_upload_info()
    jsonschema.validate(file_info, response_schema)
    return file_info


def get_transfer_from_browser_to_s3_inputs(
    workflow_job_id: int, data: dict[str, Any]
) -> dict[str, Any]:
    """Create s3 token and upload path.

    Args:
        workflow_job_id_(dict): Workflow id.
        data (dict): Data to be applied to this workflow.

    Returns:
        (dict): s3 credentials returned via video.models.s3_token._to_dict.
    """
    inputs = get_file_attributes(data)
    s3_dict = s3_token.get_s3_token(workflow_job_id, data["type"], inputs["filename"])
    inputs.update(s3_dict)
    return inputs


def setup_ingest_from_browser_workflow(data: dict[str, Any]) -> list[dict[str, Any]]:
    """Create ingest from browser workflow and inputs.

    Args:
        data (dict): Data to be applied to this workflow.
    """
    jobs_to_setup = [
        job_types.TRANSFER_FROM_BROWSER_TO_S3,
        *job_types.COMMON_INGESTION_WORKFLOW_JOB_TYPES,
    ]

    return setup_workflow(
        {
            "workflow_job_type": job_types.WORKFLOW_INGEST_FROM_BROWSER,
            "jobs_to_setup": jobs_to_setup,
            "job_type_to_inputs": {
                job_types.TRANSFER_FROM_BROWSER_TO_S3: (
                    lambda workflow_job_id: get_transfer_from_browser_to_s3_inputs(
                        workflow_job_id, data
                    )
                )
            },
        },
        data,
    )


def setup_approval_workflow(data: dict[str, Any]) -> list[dict[str, Any]]:
    """Create ingest approval workflow.

    Args:
        data (dict): Job management request data.
    """
    jobs_to_setup = [
        job_types.GET_CREATE_MEZZANINES_INPUTS,
        job_types.GET_PRODUCT_METADATA,
        job_types.EXTRACT_PRORES_MEZZANINE_METADATA,
        job_types.CREATE_MEZZANINES,
        job_types.WRITE_MEZZ_VIDEO_LOCATIONS_TO_VIDEO_ASSET_TABLE,
        job_types.CONVERT_THUMBNAILS_TO_TIFFS,
        job_types.WRITE_THUMBNAIL_LOCATIONS_TO_VIDEO_ASSET_TABLE,
        job_types.MARK_VIDEO_PRODUCT_AS_APPROVED,
    ]
    workflow_ingest_job_id = data.get("workflow_ingest_job_id")

    return setup_workflow(
        {
            "workflow_job_type": job_types.WORKFLOW_APPROVAL,
            "workflow_inputs": {
                job_io_fields.WORKFLOW_INGEST_JOB_ID: workflow_ingest_job_id
            },
            "jobs_to_setup": jobs_to_setup,
            "job_type_to_inputs": {},
        },
        data,
    )


# validate_analysis derives part of its output set (output resolution, audio
# adjustment, clip timecodes) from its analysis inputs, so the re-encode drops
# those derived fields from the seed and lets the re-run recompute them from the
# (possibly overridden) inputs. Deriving the set from validate_analysis's output
# schema keeps it in sync as those outputs change.
_REENCODE_RECOMPUTED_FIELDS = frozenset(
    JOB_TYPE_TO_CREATE_JOB_OUTPUTS_SCHEMA_MAP[job_types.VALIDATE_ANALYSIS]["properties"]
)


def _collect_reencode_seed_inputs(workflow_ingest_job_id: int) -> dict[str, Any]:
    """Gather the render inputs of a prior ingest run to re-encode from.

    The seed is the union of the inputs that validate_analysis and
    create_streamable_preview_and_thumbnails received in the referenced run,
    minus the values validate_analysis derives — those are recomputed when it
    re-runs with the (possibly overridden) inputs. Both jobs must have completed,
    so the seed reflects a fully-rendered run rather than a partial or failed one.
    """
    prior_jobs = queries.get_jobs({"parent_id": workflow_ingest_job_id})
    job_by_type = {job["type"]: job for job in prior_jobs}

    seed: dict[str, Any] = {}
    for job_type in (
        job_types.VALIDATE_ANALYSIS,
        job_types.CREATE_STREAMABLE_PREVIEW_AND_THUMBNAILS,
    ):
        prior_job = job_by_type.get(job_type)
        if not prior_job or prior_job.get("status") != job_statuses.COMPLETE:
            raise InvalidRequest(
                f"Ingest job {workflow_ingest_job_id} has no completed "
                f"{job_type} job to re-encode from"
            )
        seed.update(prior_job.get("inputs") or {})

    return {
        field: value
        for field, value in seed.items()
        if field not in _REENCODE_RECOMPUTED_FIELDS
    }


def _reencode_thumbnail_updates(
    *,
    product_id: int,
    reencode_job_id: int,
    overrides: dict[str, Any],
) -> dict[str, Any]:
    """Compute the thumbnail product_video fields to change for a re-encode.

    The chosen thumbnail (thumbnail_path) is stored denormalized as
    "{run_id}/{frame}", so it does not follow the latest_pipeline_run_id
    repoint. A crop correction is spatial — the same frame exists under the
    re-encode run, so the path is re-pointed to it. A clip (trim) correction
    shifts the output timeline, so the chosen frame is no longer valid:
    thumbnail_path and thumbnail_at_milliseconds are both cleared so a fresh
    thumbnail is picked in Workstation. A custom uploaded thumbnail (a static
    image, not a frame) and an unset thumbnail are left as-is.
    """
    current_thumbnail_path = product_video.get(product_id).get("thumbnail_path")
    if not current_thumbnail_path:
        return {}

    # A custom uploaded thumbnail is a static image (stored as
    # "{run_id}/custom.jpg"), not a frame of the video, so the timeline cannot
    # invalidate it — leave it untouched.
    thumbnail_frame_file = current_thumbnail_path.split("/", 1)[-1]
    if thumbnail_frame_file == "custom.jpg":
        return {}

    # A clip (trim) correction shifts the output timeline — a start change moves
    # every frame's index, and an end change can cut the chosen frame — so the
    # choice no longer reliably matches; clear it so a fresh thumbnail is picked.
    if overrides.keys() & {
        job_io_fields.INPUT_CLIP_START_TIMECODE,
        job_io_fields.INPUT_CLIP_END_TIMECODE,
    }:
        return {"thumbnail_path": None, "thumbnail_at_milliseconds": None}

    return {"thumbnail_path": f"{reencode_job_id}/{thumbnail_frame_file}"}


def setup_ingestion_reencode_workflow(
    data: dict[str, Any],
) -> list[dict[str, Any]]:
    """Re-render a product from a prior ingest run with corrected decisions.

    The caller references a completed ingest run and overrides the decisions
    that can misfire — the crop (cropdetect) and the clip boundaries (the trim,
    which validate_analysis derives by combining blackdetect and silencedetect).
    The workflow re-runs from validate_analysis, which re-derives the dependent
    outputs (output resolution, audio adjustment) and, when the clip is not
    overridden, the trim, then the streamable preview and thumbnails re-render
    under the new run's id.
    """
    jsonschema.validate(
        data,
        jobs_validation.build_request_schema_setup_ingestion_reencode_workflow(),
    )

    workflow_ingest_job_id = data[job_io_fields.WORKFLOW_INGEST_JOB_ID]
    seed = _collect_reencode_seed_inputs(workflow_ingest_job_id)
    seed.update(data["overrides"])
    # Replace the prior run's identity fields with this run's: drop its
    # workflow_job_id (setup_workflow assigns the new run's id) and take
    # product_id from the request rather than the copied run.
    seed.pop(job_io_fields.WORKFLOW_JOB_ID, None)
    seed[job_io_fields.PRODUCT_ID] = data["context"]["product_id"]
    seed[job_io_fields.WORKFLOW_INGEST_JOB_ID] = workflow_ingest_job_id

    jobs_to_setup = [
        job_types.VALIDATE_ANALYSIS,
        job_types.CREATE_STREAMABLE_PREVIEW_AND_THUMBNAILS,
    ]

    return setup_workflow(
        {
            "workflow_job_type": job_types.WORKFLOW_INGESTION_REENCODE,
            "workflow_inputs": seed,
            "jobs_to_setup": jobs_to_setup,
        },
        data,
    )


def setup_ingest_from_google_drive_workflow(
    data: dict[str, Any],
) -> list[dict[str, Any]]:
    """Create ingest from Google Drive workflow.

    Args:
        data (dict): Job management request data.
    """
    jobs_to_setup = [
        job_types.TRANSFER_FROM_GOOGLE_DRIVE_TO_S3,
        *job_types.COMMON_INGESTION_WORKFLOW_JOB_TYPES,
    ]

    return setup_workflow(
        {
            "workflow_job_type": job_types.WORKFLOW_INGEST_FROM_GOOGLE_DRIVE,
            "jobs_to_setup": jobs_to_setup,
            "workflow_inputs": {
                job_io_fields.GOOGLE_DRIVE_FILES: data["inputs"][
                    job_io_fields.GOOGLE_DRIVE_FILES
                ],
                job_io_fields.GOOGLE_DRIVE_AUTHORIZATION: data["inputs"][
                    job_io_fields.GOOGLE_DRIVE_AUTHORIZATION
                ],
            },
        },
        data,
    )


def setup_ingest_from_dropbox_workflow(data: dict[str, Any]) -> list[dict[str, Any]]:
    """Create ingest from Dropbox workflow.

    Args:
        data (dict): Job management request data.
    """
    jobs_to_setup = [
        job_types.TRANSFER_FROM_DROPBOX_TO_S3,
        *job_types.COMMON_INGESTION_WORKFLOW_JOB_TYPES,
    ]

    return setup_workflow(
        {
            "workflow_job_type": job_types.WORKFLOW_INGEST_FROM_DROPBOX,
            "jobs_to_setup": jobs_to_setup,
            "workflow_inputs": {
                job_io_fields.DROPBOX_FILES: data["inputs"][
                    job_io_fields.DROPBOX_FILES
                ],
            },
        },
        data,
    )


def setup_ingest_from_s3_workflow(data: dict[str, Any]) -> list[dict[str, Any]]:
    """Create ingest from S3 workflow and inputs.

    Args:
        data (dict): Data to be applied to this workflow.
    """
    jobs_to_setup = [
        job_types.TRANSFER_FROM_S3_TO_S3,
        *job_types.COMMON_INGESTION_WORKFLOW_JOB_TYPES,
    ]

    return setup_workflow(
        {
            "workflow_job_type": job_types.WORKFLOW_INGEST_FROM_S3,
            "jobs_to_setup": jobs_to_setup,
            "job_type_to_inputs": {
                job_types.TRANSFER_FROM_S3_TO_S3: (
                    lambda workflow_job_id: {
                        job_io_fields.INPUT_VIDEO_S3_BUCKET: config.VIDEO_BUCKET_NAME,
                        job_io_fields.INPUT_VIDEO_S3_KEY: data["path"],
                        job_io_fields.WORKFLOW_JOB_ID: workflow_job_id,
                    }
                )
            },
        },
        data,
    )


def get_cloud_transfer_progress(workflow_job_id: int, filename: str) -> Any:
    """Get cloud transfer progress."""
    return cloud_transfer_progress.get_cloud_transfer_progress(
        workflow_job_id, filename
    )