# Licensed to the Apache Software Foundation (ASF) under one
# or more contributor license agreements.  See the NOTICE file
# distributed with this work for additional information
# regarding copyright ownership.  The ASF licenses this file
# to you 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.
"""TuningRecord database"""
from typing import Any, Callable, List, Optional, Union

# isort: off
from typing_extensions import Literal

# isort: on

from tvm._ffi import register_object
from tvm.ir.module import IRModule
from tvm.runtime import Object
from tvm.target import Target
from tvm.tir.schedule import Schedule, Trace

from .. import _ffi_api
from ..arg_info import ArgInfo
from ..utils import _json_de_tvm


@register_object("meta_schedule.Workload")
class Workload(Object):
    """A workload, i.e. an IRModule and its structural hash.

    Parameters
    ----------
    mod : IRModule
        The workload's IRModule
    """

    mod: IRModule

    def __init__(self, mod: IRModule) -> None:
        self.__init_handle_by_constructor__(
            _ffi_api.Workload,  # type: ignore # pylint: disable=no-member
            mod,
        )

    def as_json(self) -> Any:
        """Export the workload to JSON as a python object.

        Returns
        -------
        json : Any
            The JSON serialized as a python object (e.g. a Dict or List).
            Use json.dumps() to get the associated json string.
        """
        return _json_de_tvm(_ffi_api.WorkloadAsJSON(self))  # type: ignore # pylint: disable=no-member

    @staticmethod
    def from_json(json_obj: Any) -> "Workload":
        """Create a workload from a json object.

        Parameters
        ----------
        json_obj : Any
            The json object to parse.

        Returns
        -------
        tuning_record : TuningRecord
            The parsed tuning record.
        """
        return _ffi_api.WorkloadFromJSON(json_obj)  # type: ignore # pylint: disable=no-member


@register_object("meta_schedule.TuningRecord")
class TuningRecord(Object):
    """The class of tuning records.

    Parameters
    ----------
    trace : tvm.ir.Trace
        The trace of the tuning record.
    workload : Workload
        The workload of the tuning record.
    run_secs : Optional[List[float]]
        The run time of the tuning record.
    target : Optional[Target]
        The target of the tuning record.
    args_info : Optional[List[ArgInfo]]
        The argument information of the tuning record.
    """

    trace: Trace
    workload: Workload
    run_secs: Optional[List[float]]
    target: Optional[Target]
    args_info: Optional[List[ArgInfo]]

    def __init__(  # type: ignore # pylint: disable=too-many-arguments
        self,
        trace: Trace,
        workload: Workload,
        run_secs: Optional[List[float]] = None,
        target: Optional[Target] = None,
        args_info: Optional[List[ArgInfo]] = None,
    ) -> None:
        self.__init_handle_by_constructor__(
            _ffi_api.TuningRecord,  # type: ignore # pylint: disable=no-member
            trace,
            workload,
            run_secs,
            target,
            args_info,
        )

    def as_measure_candidate(self) -> Any:
        """Generate a measure candidate given an initial IR module and a trace
        stored in the tuning record.

        Returns
        -------
        candidate : MeasureCandidate
            A generated candidate.
        """
        return _ffi_api.TuningRecordAsMeasureCandidate(self)  # type: ignore # pylint: disable=no-member

    def as_json(self) -> Any:
        """Export the tuning record to a JSON string.

        Returns
        -------
        json_str : str
            The JSON string exported.
        """
        return _json_de_tvm(_ffi_api.TuningRecordAsJSON(self))  # type: ignore # pylint: disable=no-member

    @staticmethod
    def from_json(json_obj: Any, workload: Workload) -> "TuningRecord":
        """Create a tuning record from a json object.

        Parameters
        ----------
        json_obj : Any
            The json object to parse.
        workload : Workload
            The workload.

        Returns
        -------
        tuning_record : TuningRecord
            The parsed tuning record.
        """
        return _ffi_api.TuningRecordFromJSON(json_obj, workload)  # type: ignore # pylint: disable=no-member


@register_object("meta_schedule.Database")
class Database(Object):
    """The abstract database interface."""

    DatabaseType = Union["Database", Literal["json", "memory"]]

    def has_workload(self, mod: IRModule) -> bool:
        """Check if the database has the given workload.
        Parameters
        ----------
        mod : IRModule
            The IRModule to be searched for.
        Returns
        -------
        result : bool
            Whether the database has the given workload.
        """
        return _ffi_api.DatabaseHasWorkload(self, mod)  # type: ignore # pylint: disable=no-member

    def commit_workload(self, mod: IRModule) -> Workload:
        """Commit a workload to the database if missing.

        Parameters
        ----------
        mod : IRModule
            The IRModule to be searched for or added.

        Returns
        -------
        workload : Workload
            The workload corresponding to the given IRModule.
        """
        return _ffi_api.DatabaseCommitWorkload(self, mod)  # type: ignore # pylint: disable=no-member

    def commit_tuning_record(self, record: TuningRecord) -> None:
        """Commit a tuning record to the database.

        Parameters
        ----------
        record : TuningRecord
            The tuning record to add.
        """
        _ffi_api.DatabaseCommitTuningRecord(self, record)  # type: ignore # pylint: disable=no-member

    def get_top_k(self, workload: Workload, top_k: int) -> List[TuningRecord]:
        """Get the top K valid tuning records of given workload from the database.

        Parameters
        ----------
        workload : Workload
            The workload to be searched for.
        top_k : int
            The number of top records to get.

        Returns
        -------
        top_k_records : List[TuningRecord]
            The top K records.
        """
        return _ffi_api.DatabaseGetTopK(self, workload, top_k)  # type: ignore # pylint: disable=no-member

    def get_all_tuning_records(self) -> List[TuningRecord]:
        """Get all the tuning records from the database.

        Returns
        -------
        tuning_records : List[TuningRecord]
            All tuning records from the database.
        """
        return _ffi_api.DatabaseGetAllTuningRecords(self)  # type: ignore # pylint: disable=no-member

    def __len__(self) -> int:
        """Get the number of records in the database.

        Returns
        -------
        num_records : int
            The number of records in the database
        """
        return _ffi_api.DatabaseSize(self)  # type: ignore # pylint: disable=no-member

    def query_tuning_record(
        self,
        mod: IRModule,
        target: Target,
        workload_name: str,
    ) -> Optional[TuningRecord]:
        """Query the best record of the given workload from the database.

        Parameters
        ----------
        mod : IRModule
            The IRModule to be searched for.
        target : Target
            The target to be searched for.
        workload_name : str
            The name of the workload to be searched for.

        Returns
        -------
        tuning_record : Optional[TuningRecord]
            The best record of the given workload; None if not found.
        """
        return _ffi_api.DatabaseQueryTuningRecord(self, mod, target, workload_name)  # type: ignore # pylint: disable=no-member

    def query_schedule(
        self,
        mod: IRModule,
        target: Target,
        workload_name: str,
    ) -> Optional[Schedule]:
        """Query the best schedule of the given workload from the database.

        Parameters
        ----------
        mod : IRModule
            The IRModule to be searched for.
        target : Target
            The target to be searched for.
        workload_name : str
            The name of the workload to be searched for.

        Returns
        -------
        schedule : Optional[tvm.tir.Schedule]
            The best schedule of the given workload; None if not found.
        """
        return _ffi_api.DatabaseQuerySchedule(self, mod, target, workload_name)  # type: ignore # pylint: disable=no-member

    def query_ir_module(
        self,
        mod: IRModule,
        target: Target,
        workload_name: str,
    ) -> Optional[IRModule]:
        """Query the best IRModule of the given workload from the database.

        Parameters
        ----------
        mod : IRModule
            The IRModule to be searched for.
        target : Target
            The target to be searched for.
        workload_name : str
            The name of the workload to be searched for.

        Returns
        -------
        ir_module : Optional[IRModule]
            The best IRModule of the given workload; None if not found.
        """
        return _ffi_api.DatabaseQueryIRModule(self, mod, target, workload_name)  # type: ignore # pylint: disable=no-member

    def dump_pruned(self, destination: "Database") -> None:
        """Dump the pruned database to files of JSONDatabase format.

        Parameters
        ----------
        destination : Database
            The destination database to be dumped to.
        """
        return _ffi_api.DatabaseDumpPruned(  # type: ignore # pylint: disable=no-member
            self, destination
        )

    def query(
        self,
        mod: IRModule,
        target: Target,
        *,
        workload_name: str = "main",
        kind: Union[
            Literal["schedule"],
            Literal["record"],
            Literal["ir_module"],
        ] = "schedule",
    ) -> Union[Schedule, IRModule, TuningRecord]:
        """Query the database to retrieve the best optimization outcome of the given workload.

        Parameters
        ----------
        mod : IRModule
            The IRModule to be searched for.
        target : Target
            The target to be searched for.
        kind : str = "schedule" | "record" | "ir_module"
            The kind of the optimization outcome to be returned.

        Returns
        -------
        result : Union[tvm.tir.Schedule, IRModule, TuningRecord]
            The best optimization outcome of the given workload.
        """
        if kind == "schedule":
            return self.query_schedule(mod, target, workload_name)
        if kind == "record":
            return self.query_tuning_record(mod, target, workload_name)
        if kind == "ir_module":
            return self.query_ir_module(mod, target, workload_name)
        raise ValueError(f'Unknown kind: {kind}. Candidates are: "schedule", "record", "ir_module"')

    def __enter__(self) -> "Database":
        """Entering the scope of the context manager"""
        _ffi_api.DatabaseEnterWithScope(self)  # type: ignore # pylint: disable=no-member
        return self

    def __exit__(self, ptype, value, trace) -> None:
        """Exiting the scope of the context manager"""
        _ffi_api.DatabaseExitWithScope(self)  # type: ignore # pylint: disable=no-member

    @staticmethod
    def current() -> Optional["Database"]:
        """Get the current database under scope."""
        return _ffi_api.DatabaseCurrent()  # type: ignore # pylint: disable=no-member

    @staticmethod
    def create(  # pylint: disable=keyword-arg-before-vararg
        kind: Union[
            Literal[
                "json",
                "memory",
                "union",
                "ordered_union",
            ],
            Callable[[Schedule], bool],
        ] = "json",
        *args,
        **kwargs,
    ) -> "Database":
        """Create a Database.

        Parameters
        ----------
        kind : str = "json" | "memory" | "union" | "ordered_union" | Callable[[tvm.tir.Schedule],
        bool]
            The kind of the database to be created. The following kinds are supported:
            "json", "memory", "union", "ordered_union", and a custom schedule function.

        Returns
        -------
        database : Database
            The created database.
        """
        from . import (  # pylint: disable=import-outside-toplevel
            JSONDatabase,
            MemoryDatabase,
            OrderedUnionDatabase,
            ScheduleFnDatabase,
            UnionDatabase,
        )

        if callable(kind):
            return ScheduleFnDatabase(kind, *args, **kwargs)  # type: ignore
        if kind == "json":
            return JSONDatabase(*args, **kwargs)
        if kind == "memory":
            return MemoryDatabase(*args, **kwargs)  # type: ignore
        if kind == "union":
            return UnionDatabase(*args, **kwargs)  # type: ignore
        if kind == "ordered_union":
            return OrderedUnionDatabase(*args, **kwargs)  # type: ignore
        raise ValueError(f"Unknown Database: {kind}")


create = Database.create  # pylint: disable=invalid-name


@register_object("meta_schedule.PyDatabase")
class _PyDatabase(Database):
    """
    A TVM object database to support customization on the python side.
    This is NOT the user facing class for function overloading inheritance.

    See also: PyDatabase
    """

    def __init__(
        self,
        f_has_workload: Callable = None,
        f_commit_workload: Callable = None,
        f_commit_tuning_record: Callable = None,
        f_get_top_k: Callable = None,
        f_get_all_tuning_records: Callable = None,
        f_query_tuning_record: Callable = None,
        f_query_schedule: Callable = None,
        f_query_ir_module: Callable = None,
        f_size: Callable = None,
        module_equality: str = "structural",
    ):
        """Constructor."""

        self.__init_handle_by_constructor__(
            _ffi_api.DatabasePyDatabase,  # type: ignore  # pylint: disable=no-member
            f_has_workload,
            f_commit_workload,
            f_commit_tuning_record,
            f_get_top_k,
            f_get_all_tuning_records,
            f_query_tuning_record,
            f_query_schedule,
            f_query_ir_module,
            f_size,
            module_equality,
        )


class PyDatabase:
    """
    An abstract database with customized methods on the python-side.
    This is the user facing class for function overloading inheritance.

    Note: @derived_object is required for proper usage of any inherited class.
    """

    _tvm_metadata = {
        "cls": _PyDatabase,
        "methods": [
            "has_workload",
            "commit_workload",
            "commit_tuning_record",
            "get_top_k",
            "get_all_tuning_records",
            "query_tuning_record",
            "query_schedule",
            "query_ir_module",
            "__len__",
        ],
    }

    def has_workload(self, mod: IRModule) -> bool:
        """Check if the database has the given workload.
        Parameters
        ----------
        mod : IRModule
            The IRModule to be searched for.
        Returns
        -------
        result : bool
            Whether the database has the given workload.
        """
        raise NotImplementedError

    def commit_workload(self, mod: IRModule) -> Workload:
        """Commit a workload to the database if missing.

        Parameters
        ----------
        mod : IRModule
            The IRModule to be searched for or added.

        Returns
        -------
        workload : Workload
            The workload corresponding to the given IRModule.
        """
        raise NotImplementedError

    def commit_tuning_record(self, record: TuningRecord) -> None:
        """Commit a tuning record to the database.

        Parameters
        ----------
        record : TuningRecord
            The tuning record to add.
        """
        raise NotImplementedError

    def get_top_k(self, workload: Workload, top_k: int) -> List[TuningRecord]:
        """Get the top K tuning records of given workload from the database.

        Parameters
        ----------
        workload : Workload
            The workload to be searched for.
        top_k : int
            The number of top records to get.

        Returns
        -------
        top_k_records : List[TuningRecord]
            The top K records.
        """
        raise NotImplementedError

    def get_all_tuning_records(self) -> List[TuningRecord]:
        """Get all the tuning records from the database.

        Returns
        -------
        tuning_records : List[TuningRecord]
            All tuning records from the database.
        """
        raise NotImplementedError

    def query_tuning_record(
        self, mod: IRModule, target: Target, workload_name: Optional[str] = None
    ) -> Optional[TuningRecord]:
        """Query a tuning record from the database.

        Parameters
        ----------
        mod : IRModule
            The IRModule to be searched for.
        target : Target
            The target to be searched for.
        workload_name : Optional[str]
            The workload name to be searched for.

        Returns
        -------
        record : Optional[TuningRecord]
            The tuning record corresponding to the given workload.
        """
        # Using self._outer to replace the self pointer
        return _ffi_api.DatabaseQueryTuningRecord(  # type: ignore # pylint: disable=no-member
            self._outer(), mod, target, workload_name  # type: ignore # pylint: disable=no-member
        )

    def query_schedule(
        self, mod: IRModule, target: Target, workload_name: Optional[str] = None
    ) -> Optional[Schedule]:
        """Query a schedule from the database.

        Parameters
        ----------
        mod : IRModule
            The IRModule to be searched for.
        target : Target
            The target to be searched for.
        workload_name : Optional[str]
            The workload name to be searched for.

        Returns
        -------
        schedule : Optional[Schedule]
            The schedule corresponding to the given workload.
        """
        # Using self._outer to replace the self pointer
        return _ffi_api.DatabaseQuerySchedule(  # type: ignore # pylint: disable=no-member
            self._outer(), mod, target, workload_name  # type: ignore # pylint: disable=no-member
        )

    def query_ir_module(
        self, mod: IRModule, target: Target, workload_name: Optional[str] = None
    ) -> Optional[IRModule]:
        """Query an IRModule from the database.

        Parameters
        ----------
        mod : IRModule
            The IRModule to be searched for.
        target : Target
            The target to be searched for.
        workload_name : Optional[str]
            The workload name to be searched for.

        Returns
        -------
        mod : Optional[IRModule]
            The IRModule corresponding to the given workload.
        """
        # Using self._outer to replace the self pointer
        return _ffi_api.DatabaseQueryIRModule(  # type: ignore # pylint: disable=no-member
            self._outer(), mod, target, workload_name  # type: ignore # pylint: disable=no-member
        )

    def __len__(self) -> int:
        """Get the number of records in the database.

        Returns
        -------
        num_records : int
            The number of records in the database
        """
        raise NotImplementedError
