# -*- 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.
# pylint: disable=no-member
"""INC Bench Benchmark API interface."""
import os
import shutil
from typing import List, Optional, Union

from sqlalchemy.orm import sessionmaker

from neural_compressor.ux.components.benchmark import Benchmarks
from neural_compressor.ux.components.configuration_wizard.configuration_parser import (
    ConfigurationParser,
)
from neural_compressor.ux.components.db_manager.db_manager import DBManager
from neural_compressor.ux.components.db_manager.db_models.benchmark import Benchmark
from neural_compressor.ux.components.db_manager.db_models.benchmark_result import BenchmarkResult
from neural_compressor.ux.components.db_manager.db_operations.project_api_interface import (
    ProjectAPIInterface,
)
from neural_compressor.ux.components.db_manager.params_interfaces import (
    BenchmarkAddParamsInterface,
    BenchmarkEditParamsInterface,
)
from neural_compressor.ux.components.jobs_management import jobs_control_queue, parse_job_id
from neural_compressor.ux.utils.consts import WORKSPACE_LOCATION, ExecutionStatus
from neural_compressor.ux.utils.exceptions import ClientErrorException
from neural_compressor.ux.utils.logger import log
from neural_compressor.ux.utils.utils import normalize_string

db_manager = DBManager()
Session = sessionmaker(bind=db_manager.engine)


class BenchmarkAPIInterface:
    """Interface for queries connected with benchmark."""

    @staticmethod
    def delete_benchmark(data: dict) -> dict:
        """Delete benchmark from database and clean workspace."""
        try:
            benchmark_id: int = int(data.get("id", None))
            benchmark_name: str = str(data.get("name", None))
            job_id = parse_job_id("benchmark", benchmark_id)
            jobs_control_queue.abort_job(job_id)
        except ValueError:
            raise ClientErrorException("Could not parse value.")
        except TypeError:
            raise ClientErrorException("Missing project id or project name.")
        with Session.begin() as db_session:
            benchmark_details = Benchmark.details(db_session, benchmark_id)
            project_id = benchmark_details["project_id"]
            project_details = ProjectAPIInterface.get_project_details({"id": project_id})
            removed_benchmark_id = Benchmark.delete_benchmark(
                db_session=db_session,
                benchmark_id=benchmark_id,
                benchmark_name=benchmark_name,
            )

        if removed_benchmark_id is not None:
            try:
                model_id = benchmark_details["model"]["id"]
                normalized_project_name = normalize_string(project_details["name"])
                normalized_benchmark_name = normalize_string(benchmark_name)
                normalized_model_name = normalize_string(benchmark_details["model"]["name"])
                benchmark_location = os.path.join(
                    WORKSPACE_LOCATION,
                    f"{normalized_project_name}_{project_id}",
                    "models",
                    f"{normalized_model_name}_{model_id}",
                    "benchmarks",
                    f"{normalized_benchmark_name}_{benchmark_id}",
                )
                shutil.rmtree(benchmark_location, ignore_errors=True)
            except Exception:
                log.debug("Could not find benchmark directory.")

        return {"id": removed_benchmark_id}

    @staticmethod
    def get_benchmark_details(data: dict) -> dict:
        """Parse input data and get benchmark details."""
        try:
            benchmark_id: int = int(data.get("id", None))
        except ValueError:
            raise ClientErrorException("Incorrect benchmark id.")
        except TypeError:
            raise ClientErrorException("Could not find benchmark id.")

        with Session.begin() as db_session:
            benchmark_details = Benchmark.details(
                db_session,
                benchmark_id,
            )
        return benchmark_details

    @staticmethod
    def list_benchmarks(data: dict) -> dict:
        """List benchmarks assigned to project."""
        try:
            project_id: int = int(data.get("project_id", None))
        except ValueError:
            raise ClientErrorException("Incorrect project id.")
        except TypeError:
            raise ClientErrorException("Could not find project id.")

        with Session.begin() as db_session:
            benchmarks_list = Benchmark.list(
                db_session,
                project_id,
            )
        return benchmarks_list

    @staticmethod
    def update_benchmark_accuracy(data: dict) -> dict:
        """Update benchmark accuracy."""
        try:
            benchmark_id: int = int(data.get("id", None))
        except ValueError:
            raise ClientErrorException("Incorrect benchmark id.")
        except TypeError:
            raise ClientErrorException("Could not find benchmark id.")

        try:
            accuracy: float = float(data.get("status", None))
        except ValueError as err:
            raise ClientErrorException(err)

        with Session.begin() as db_session:
            response_data = BenchmarkResult.update_accuracy(
                db_session,
                benchmark_id,
                accuracy,
            )
        return response_data

    @staticmethod
    def update_benchmark_performance(data: dict) -> dict:
        """Update benchmark performance."""
        try:
            benchmark_id: int = int(data.get("id", None))
        except ValueError:
            raise ClientErrorException("Incorrect benchmark id.")
        except TypeError:
            raise ClientErrorException("Could not find benchmark id.")

        try:
            performance: float = float(data.get("status", None))
        except ValueError as err:
            raise ClientErrorException(err)

        with Session.begin() as db_session:
            response_data = BenchmarkResult.update_performance(
                db_session,
                benchmark_id,
                performance,
            )
        return response_data

    @staticmethod
    def update_benchmark_status(data: dict) -> dict:
        """Update benchmark status."""
        try:
            benchmark_id: int = int(data.get("id", None))
        except ValueError:
            raise ClientErrorException("Incorrect benchmark id.")
        except TypeError:
            raise ClientErrorException("Could not find benchmark id.")

        try:
            status: ExecutionStatus = ExecutionStatus(data.get("status", None))
        except ValueError as err:
            raise ClientErrorException(err)

        with Session.begin() as db_session:
            response_data = Benchmark.update_status(
                db_session,
                benchmark_id,
                status,
            )
        return response_data

    @staticmethod
    def update_benchmark_duration(data: dict) -> dict:
        """Update duration of benchmark."""
        try:
            benchmark_id: int = int(data.get("id", None))
        except ValueError:
            raise ClientErrorException("Incorrect benchmark id.")
        except TypeError:
            raise ClientErrorException("Could not find benchmark id.")

        try:
            duration = int(data.get("duration", None))
        except ValueError as err:
            raise ClientErrorException(err)

        with Session.begin() as db_session:
            response_data = Benchmark.update_duration(
                db_session,
                benchmark_id,
                duration,
            )
        return response_data

    @staticmethod
    def update_paths(data: dict) -> dict:
        """Update config path and output log path."""
        response = {}
        try:
            benchmark_id: int = int(data.get("id", None))
        except ValueError:
            raise ClientErrorException("Incorrect benchmark id.")
        except TypeError:
            raise ClientErrorException("Could not find benchmark id.")

        config_path: Optional[str] = data.get("config_path")
        log_path: Optional[str] = data.get("log_path")

        with Session.begin() as db_session:
            config_path_response = Benchmark.update_config_path(
                db_session=db_session,
                benchmark_id=benchmark_id,
                path=config_path,
            )
            response.update(config_path_response)
            log_path_response = Benchmark.update_log_path(
                db_session=db_session,
                benchmark_id=benchmark_id,
                path=log_path,
            )
            response.update(log_path_response)
        return response

    @staticmethod
    def update_execution_command(data: dict) -> dict:
        """Update benchmark execution command."""
        try:
            benchmark_id: int = int(data.get("id", None))
        except ValueError:
            raise ClientErrorException("Incorrect benchmark id.")
        except TypeError:
            raise ClientErrorException("Could not find benchmark id.")

        execution_command: Optional[Union[str, List[str]]] = data.get("execution_command")
        if isinstance(execution_command, list):
            execution_command = " ".join(map(str, execution_command))
        with Session.begin() as db_session:
            response_data = Benchmark.update_execution_command(
                db_session=db_session,
                benchmark_id=benchmark_id,
                execution_command=execution_command,
            )
        return response_data

    @staticmethod
    def add_benchmark(data: dict) -> dict:
        """Add benchmark to database."""
        parser = ConfigurationParser()
        parsed_data = parser.parse(data)

        benchmark_params: BenchmarkAddParamsInterface = BenchmarkAPIInterface.parse_benchmark_data(
            parsed_data,
        )

        with Session.begin() as db_session:
            benchmark_id = Benchmark.add(
                db_session=db_session,
                project_id=benchmark_params.project_id,
                name=benchmark_params.name,
                model_id=benchmark_params.model_id,
                dataset_id=benchmark_params.dataset_id,
                mode=benchmark_params.mode,
                batch_size=benchmark_params.batch_size,
                iterations=benchmark_params.iterations,
                number_of_instance=benchmark_params.number_of_instance,
                cores_per_instance=benchmark_params.cores_per_instance,
                warmup_iterations=benchmark_params.warmup_iterations,
                execution_command=benchmark_params.command_line,
            )
        return {
            "benchmark_id": benchmark_id,
        }

    @staticmethod
    def add_result(data: dict) -> None:
        """Add benchmark result to database."""
        try:
            benchmark_id: int = int(data.get("benchmark_id", None))
        except ValueError:
            raise ClientErrorException("Incorrect benchmark id.")
        except TypeError:
            raise ClientErrorException("Could not find benchmark id.")

        accuracy: Optional[float] = None
        performance: Optional[float] = None

        try:
            accuracy = float(data["accuracy"])
        except Exception:
            pass

        try:
            performance = float(data["performance"])
        except Exception:
            pass
        with Session.begin() as db_session:
            BenchmarkResult.add(
                db_session=db_session,
                benchmark_id=benchmark_id,
                accuracy=accuracy,
                performance=performance,
            )

    @staticmethod
    def edit_benchmark(data: dict) -> dict:
        """Edit existing benchmark."""
        parser = ConfigurationParser()
        parsed_input_data = parser.parse(data)
        parsed_benchmark_data: BenchmarkEditParamsInterface = (
            BenchmarkAPIInterface.parse_benchmark_edit_data(
                parsed_input_data,
            )
        )
        response: dict = {"id": parsed_benchmark_data.id}
        with Session.begin() as db_session:
            benchmark = Benchmark.details(
                db_session,
                parsed_benchmark_data.id,
            )

            if benchmark.get("status", None) is not None:
                raise ClientErrorException("Can not modify benchmark  that has been run.")

            if parsed_benchmark_data.dataset_id is not None:
                response.update(
                    Benchmark.update_dataset(
                        db_session=db_session,
                        benchmark_id=parsed_benchmark_data.id,
                        dataset_id=parsed_benchmark_data.dataset_id,
                    ),
                )

            if parsed_benchmark_data.batch_size is not None:
                response.update(
                    Benchmark.update_batch_size(
                        db_session=db_session,
                        benchmark_id=parsed_benchmark_data.id,
                        batch_size=parsed_benchmark_data.batch_size,
                    ),
                )

            if parsed_benchmark_data.mode is not None:
                response.update(
                    Benchmark.update_mode(
                        db_session=db_session,
                        benchmark_id=parsed_benchmark_data.id,
                        mode=parsed_benchmark_data.mode,
                    ),
                )

            if parsed_benchmark_data.cores_per_instance is not None:
                response.update(
                    Benchmark.update_cores_per_instance(
                        db_session=db_session,
                        benchmark_id=parsed_benchmark_data.id,
                        cores_per_instance=parsed_benchmark_data.cores_per_instance,
                    ),
                )

            if parsed_benchmark_data.number_of_instance is not None:
                response.update(
                    Benchmark.update_number_of_instance(
                        db_session=db_session,
                        benchmark_id=parsed_benchmark_data.id,
                        number_of_instance=parsed_benchmark_data.number_of_instance,
                    ),
                )

        return response

    @staticmethod
    def parse_benchmark_data(data: dict) -> BenchmarkAddParamsInterface:
        """Parse input data for benchmark."""
        benchmark_data = BenchmarkAddParamsInterface()
        try:
            benchmark_data.project_id = int(data.get("project_id", None))
            benchmark_data.name = str(data.get("name", None))
            benchmark_data.model_id = int(data.get("model_id", None))
            benchmark_data.dataset_id = int(data.get("dataset_id", None))
            benchmark_data.mode = str(data.get("mode", Benchmarks.PERF))
            benchmark_data.batch_size = int(data.get("batch_size", 100))
            benchmark_data.iterations = int(data.get("iterations", -1))
            benchmark_data.cores_per_instance = int(data.get("cores_per_instance", 4))
            benchmark_data.warmup_iterations = int(data.get("warmup_iterations", 10))
            benchmark_data.command_line = str(data.get("command_line"))
        except ValueError:
            raise ClientErrorException("Could not parse value")
        except TypeError:
            raise ClientErrorException("Could not find required parameter.")

        try:
            benchmark_data.number_of_instance = int(data.get("number_of_instance", None))
        except TypeError:
            from neural_compressor.ux.utils.hw_info import HWInfo

            hw_info = HWInfo()
            benchmark_data.number_of_instance = (
                hw_info.cores_per_socket // benchmark_data.cores_per_instance
            )

        return benchmark_data

    @staticmethod
    def parse_benchmark_edit_data(data: dict) -> BenchmarkEditParamsInterface:
        """Parse data for editing benchmark."""
        benchmark_data = BenchmarkEditParamsInterface()
        try:
            benchmark_data.id = int(data.get("id", None))
            if "dataset_id" in data:
                benchmark_data.dataset_id = int(data.get("dataset_id", None))
            if "batch_size" in data:
                benchmark_data.batch_size = int(data.get("batch_size", None))
            if "mode" in data:
                benchmark_data.mode = str(data.get("mode", None))
            if "cores_per_instance" in data:
                benchmark_data.cores_per_instance = int(data.get("cores_per_instance", None))
            if "number_of_instance" in data:
                benchmark_data.number_of_instance = int(data.get("number_of_instance", None))
        except ValueError:
            raise ClientErrorException("Could not parse value")
        except TypeError:
            raise ClientErrorException("Could not find required parameter.")

        return benchmark_data

    @staticmethod
    def clean_status(status_to_clean: ExecutionStatus) -> dict:
        """Clean specified optimization status."""
        with Session.begin() as db_session:
            response = Benchmark.clean_status(
                db_session=db_session,
                status_to_clean=status_to_clean,
            )
        return response
