# -*- coding: utf-8 -*-
# Copyright (c) 2022 Intel Corporation
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#   http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
"""Status update function factory."""

from typing import Callable

from neural_compressor.ux.components.db_manager.db_operations import (
    BenchmarkAPIInterface,
    OptimizationAPIInterface,
    ProfilingAPIInterface,
)

STATUS_UPDATE_FUNCTIONS = {
    "benchmark": BenchmarkAPIInterface.update_benchmark_status,
    "profiling": ProfilingAPIInterface.update_profiling_status,
    "optimization": OptimizationAPIInterface.update_optimization_status,
}


def get_status_update_function(job_type: str) -> Callable[[dict], dict]:
    """Get update status function for specified job type."""
    function = STATUS_UPDATE_FUNCTIONS.get(job_type)
    if function is not None:
        return function
    else:
        raise KeyError()
