# 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.
"""Configurations for measurements in the runner"""

import os
from threading import Thread
from typing import NamedTuple, Optional

from tvm import rpc


class EvaluatorConfig(NamedTuple):
    """Config Details of Evaluator

    Parameters
    ----------
    number: int
        The number of times to run this function for taking average.
        We call these runs as one `repeat` of measurement.
    repeat: int
        The number of times to repeat the measurement.
        In total, the function will be invoked (1 + number x repeat) times,
        where the first one is warm up and will be discarded.
        The returned result contains `repeat` costs,
        each of which is an average of `number` costs.
    min_repeat_ms: int
        Minimum repeat time in ms. if the execution latency is too short,
        increase the number of runs to the given time (in ms) to reduce the measurement error.
    enable_cpu_cache_flush: bool
        Whether to flush the cache on CPU.

    Note
    ----
    The total number of actual executions is 1+number*repeat because we would warm up 1 time before
    actual run. The number of runs would be increased if run time is below min_repeat_ms.
    """

    number: int = 3
    repeat: int = 1
    min_repeat_ms: int = 100
    enable_cpu_cache_flush: bool = False

    @staticmethod
    def _normalized(config: Optional["EvaluatorConfig"]) -> "EvaluatorConfig":
        if config is None:
            return EvaluatorConfig()
        config = EvaluatorConfig(
            number=config.number,
            repeat=config.repeat,
            min_repeat_ms=config.min_repeat_ms,
            enable_cpu_cache_flush=config.enable_cpu_cache_flush,
        )
        return config


class RPCConfig(NamedTuple):
    """RPC configuration

    Parameters
    ----------
    tracker_host: str
        Host of the RPC Tracker
    tracker_port: int
        Port of the RPC Tracker
    tracker_key: str
        Key of the Tracker
    session_timeout_sec: float
        Timeout of the RPC session
    session_priority: int
        Priority of the RPC session
    """

    tracker_host: str | None = None
    tracker_port: None | int | str = None
    tracker_key: str | None = None
    session_priority: int = 1
    session_timeout_sec: int = 10

    def _sanity_check(self) -> None:
        err_str = (
            "RPCConfig.{0} is not provided. Please provide it explicitly,"
            "or set environment variable {1}"
        )
        if self.tracker_host is None:
            raise ValueError(err_str.format("tracker_host", "TVM_TRACKER_HOST"))
        if self.tracker_port is None:
            raise ValueError(err_str.format("tracker_port", "TVM_TRACKER_PORT"))
        if self.tracker_key is None:
            raise ValueError(err_str.format("tracker_key", "TVM_TRACKER_KEY"))

    @staticmethod
    def _normalized(config: Optional["RPCConfig"]) -> "RPCConfig":
        if config is None:
            config = RPCConfig()
        tracker_host = config.tracker_host or os.environ.get("TVM_TRACKER_HOST", None)
        tracker_port = config.tracker_port or os.environ.get("TVM_TRACKER_PORT", None)
        tracker_key = config.tracker_key or os.environ.get("TVM_TRACKER_KEY", None)
        if isinstance(tracker_port, str):
            tracker_port = int(tracker_port)
        config = RPCConfig(
            tracker_host=tracker_host,
            tracker_port=tracker_port,
            tracker_key=tracker_key,
            session_priority=config.session_priority,
            session_timeout_sec=config.session_timeout_sec,
        )
        config._sanity_check()  # pylint: disable=protected-access
        return config

    def connect_tracker(self) -> rpc.TrackerSession:
        """Connect to the tracker

        Returns
        -------
        tracker : TrackerSession
            The connected tracker session
        """
        tracker: rpc.TrackerSession | None = None

        def _connect():
            nonlocal tracker
            tracker = rpc.connect_tracker(self.tracker_host, self.tracker_port)

        t = Thread(target=_connect)
        t.start()
        t.join(self.session_timeout_sec)
        if t.is_alive() or tracker is None:
            raise ValueError(
                "Unable to connect to the tracker using the following configuration:\n"
                f"    tracker host: {self.tracker_host}\n"
                f"    tracker port: {self.tracker_port}\n"
                f"    timeout (sec): {self.session_timeout_sec}\n"
                "Please check the tracker status via the following command:\n"
                "     python3 -m tvm.exec.query_rpc_tracker "
                f"--host {self.tracker_host} --port {self.tracker_port}"
            )
        return tracker

    def connect_server(self) -> rpc.RPCSession:
        """Connect to the server

        Returns
        -------
        session : RPCSession
            The connected rpc session
        """
        tracker = self.connect_tracker()
        session: rpc.RPCSession = tracker.request(
            key=self.tracker_key,
            priority=self.session_priority,
            session_timeout=self.session_timeout_sec,
        )
        return session

    def count_num_servers(self, allow_missing=True) -> int:
        """Count the number of servers available in the tracker

        Parameters
        ----------
        allow_missing : bool
            Whether to allow no server to be found.

        Returns
        -------
        num_servers : int
            The number of servers
        """
        tracker = self.connect_tracker()
        tracker_summary = tracker.summary()
        result: int = 0
        for item in tracker_summary["server_info"]:
            _, item_key = item["key"].split(":")
            if item_key == self.tracker_key:
                result += 1
        if result == 0 and not allow_missing:
            raise ValueError(
                "Unable to find servers with the specific key using the following configuration:\n"
                f"    tracker host: {self.tracker_host}\n"
                f"    tracker port: {self.tracker_port}\n"
                f"    tracker key: {self.tracker_key}\n"
                f"    timeout (sec): {self.session_timeout_sec}\n"
                "Please check the tracker status via the following command:\n"
                "     python3 -m tvm.exec.query_rpc_tracker "
                f"--host {self.tracker_host} --port {self.tracker_port}\n"
                f'and look for key: "{self.tracker_key}"'
            )
        return result
