import os

import requests

from fedn.common.telemetry import trace_all_methods

__all__ = ['APIClient']


@trace_all_methods
class APIClient:
    """ An API client for interacting with the statestore and controller.

    :param host: The host of the api server.
    :type host: str
    :param port: The port of the api server.
    :type port: int
    :param secure: Whether to use https.
    :type secure: bool
    :param verify: Whether to verify the server certificate.
    :type verify: bool
    """

    def __init__(self, host, port=None, secure=False, verify=False, token=None, auth_scheme=None):
        self.host = host
        self.port = port
        self.secure = secure
        self.verify = verify
        self.headers = {}
        # Auth scheme passed as argument overrides environment variable.
        # "Token" is the default auth scheme.
        if not auth_scheme:
            auth_scheme = os.environ.get("FEDN_AUTH_SCHEME", "Bearer")
        # Override potential env variable if token is passed as argument.
        if not token:
            token = os.environ.get("FEDN_AUTH_TOKEN", False)

        if token:
            self.headers = {"Authorization": f"{auth_scheme} {token}"}

    def _get_url(self, endpoint):
        if self.secure:
            protocol = 'https'
        else:
            protocol = 'http'
        if self.port:
            return f'{protocol}://{self.host}:{self.port}/{endpoint}'
        return f'{protocol}://{self.host}/{endpoint}'

    def _get_url_api_v1(self, endpoint):
        return self._get_url(f'api/v1/{endpoint}')

    # --- Clients --- #

    def get_client(self, id: str):
        """ Get a client from the statestore.

        :param id: The client id to get.
        :type id: str
        :return: Client.
        :rtype: dict
        """
        response = requests.get(self._get_url_api_v1(f'clients/{id}'), verify=self.verify, headers=self.headers)

        _json = response.json()

        return _json

    def get_clients(self, n_max: int = None):
        """ Get clients from the statestore.

        :param n_max: The maximum number of clients to get (If none all will be fetched).
        :type n_max: int
        return: Clients.
        rtype: dict
        """
        _headers = self.headers.copy()

        if n_max:
            _headers['X-Limit'] = str(n_max)

        response = requests.get(self._get_url_api_v1('clients'), verify=self.verify, headers=_headers)

        _json = response.json()

        return _json

    def get_clients_count(self):
        """ Get the number of clients in the statestore.

        :return: The number of clients.
        :rtype: dict
        """
        response = requests.get(self._get_url_api_v1('clients/count'), verify=self.verify, headers=self.headers)

        _json = response.json()

        return _json

    def get_client_config(self, checksum=True):
        """ Get client config from controller. Optionally include the checksum.
        The config is used for clients to connect to the controller and ask for combiner assignment.

        :param checksum: Whether to include the checksum of the package.
        :type checksum: bool
        :return: The client configuration.
        :rtype: dict
        """
        _params = {
            'checksum': "true" if checksum else "false"
        }

        response = requests.get(self._get_url('get_client_config'), params=_params, verify=self.verify, headers=self.headers)

        _json = response.json()

        return _json

    def get_active_clients(self, combiner_id: str = None, n_max: int = None):
        """ Get active clients from the statestore.

        :param combiner_id: The combiner id to get active clients for.
        :type combiner_id: str
        :param n_max: The maximum number of clients to get (If none all will be fetched).
        :type n_max: int
        :return: Active clients.
        :rtype: dict
        """
        _params = {"status": "online"}

        if combiner_id:
            _params['combiner'] = combiner_id

        _headers = self.headers.copy()

        if n_max:
            _headers['X-Limit'] = str(n_max)

        response = requests.get(self._get_url_api_v1('clients/'), params=_params, verify=self.verify, headers=_headers)

        _json = response.json()

        return _json

    # --- Combiners --- #

    def get_combiner(self, id: str):
        """ Get a combiner from the statestore.

        :param id: The combiner id to get.
        :type id: str
        :return: Combiner.
        :rtype: dict
        """
        response = requests.get(self._get_url_api_v1(f'combiners/{id}'), verify=self.verify, headers=self.headers)

        _json = response.json()

        return _json

    def get_combiners(self, n_max: int = None):
        """ Get combiners in the network.

        :param n_max: The maximum number of combiners to get (If none all will be fetched).
        :type n_max: int
        :return: Combiners.
        :rtype: dict
        """
        _headers = self.headers.copy()

        if n_max:
            _headers['X-Limit'] = str(n_max)

        response = requests.get(self._get_url_api_v1('combiners/'), verify=self.verify, headers=_headers)

        _json = response.json()

        return _json

    def get_combiners_count(self):
        """ Get the number of combiners in the statestore.

        :return: The number of combiners.
        :rtype: dict
        """
        response = requests.get(self._get_url_api_v1('combiners/count'), verify=self.verify, headers=self.headers)

        _json = response.json()

        return _json

    # --- Controllers --- #

    def get_controller_status(self):
        """ Get the status of the controller.

        :return: The status of the controller.
        :rtype: dict
        """
        response = requests.get(self._get_url('get_controller_status'), verify=self.verify, headers=self.headers)

        _json = response.json()

        return _json

    # --- Models --- #

    def get_model(self, id: str):
        """ Get a model from the statestore.

        :param id: The id (or model property) of the model to get.
        :type id: str
        :return: Model.
        :rtype: dict
        """
        response = requests.get(self._get_url_api_v1(f'models/{id}'), verify=self.verify, headers=self.headers)

        _json = response.json()

        return _json

    def get_models(self, session_id: str = None, n_max: int = None):
        """ Get models from the statestore.

        :param session_id: The session id to get models for. (optional)
        :type session_id: str
        :param n_max: The maximum number of models to get (If none all will be fetched).
        :type n_max: int
        :return: Models.
        :rtype: dict
        """
        _params = {}

        if session_id:
            _params['session_id'] = session_id

        _headers = self.headers.copy()

        if n_max:
            _headers['X-Limit'] = str(n_max)

        response = requests.get(self._get_url_api_v1('models/'), params=_params, verify=self.verify, headers=_headers)

        _json = response.json()

        return _json

    def get_models_count(self):
        """ Get the number of models in the statestore.

        :return: The number of models.
        :rtype: dict
        """
        response = requests.get(self._get_url_api_v1('models/count'), verify=self.verify, headers=self.headers)

        _json = response.json()

        return _json

    def get_active_model(self):
        """ Get the latest model from the statestore.

        :return: The latest model.
        :rtype: dict
        """
        _headers = self.headers.copy()
        _headers['X-Limit'] = "1"

        response = requests.get(self._get_url_api_v1('models/'), verify=self.verify, headers=_headers)
        _json = response.json()

        if "result" in _json and len(_json["result"]) > 0:
            return _json["result"][0]

        return _json

    def get_model_trail(self, id: str = None, include_self: bool = True, reverse: bool = True, n_max: int = None):
        """ Get the model trail.

        :param id: The id (or model property) of the model to start the trail from. (optional)
        :type id: str
        :param n_max: The maximum number of models to get (If none all will be fetched).
        :type n_max: int
        :return: Models.
        :rtype: dict
        """
        if not id:
            model = self.get_active_model()
            if "id" in model:
                id = model["id"]
            else:
                return model

        _headers = self.headers.copy()

        _count: int = n_max if n_max else self.get_models_count()
        _headers['X-Limit'] = str(_count)
        _headers['X-Reverse'] = "true" if reverse else "false"

        _include_self_str: str = "true" if include_self else "false"

        response = requests.get(self._get_url_api_v1(f'models/{id}/ancestors?include_self={_include_self_str}'), verify=self.verify, headers=_headers)
        _json = response.json()

        return _json

    def download_model(self, id: str, path: str):
        """ Download the model with id id.

        :param id: The id (or model property) of the model to download.
        :type id: str
        :param path: The path to download the model to.
        :type path: str
        :return: Message with success or failure.
        :rtype: dict
        """
        response = requests.get(self._get_url_api_v1(f'models/{id}/download'), verify=self.verify, headers=self.headers)

        if response.status_code == 200:

            with open(path, 'wb') as file:
                file.write(response.content)
            return {'success': True, 'message': 'Model downloaded successfully.'}
        else:
            return {'success': False, 'message': 'Failed to download model.'}

    def set_active_model(self, path):
        """ Set the initial model in the statestore and upload to model repository.

        :param path: The file path of the initial model to set.
        :type path: str
        :return: A dict with success or failure message.
        :rtype: dict
        """
        with open(path, 'rb') as file:
            response = requests.post(self._get_url('set_initial_model'), files={'file': file}, verify=self.verify, headers=self.headers)
        return response.json()

    # --- Packages --- #

    def get_package(self, id: str):
        """ Get a compute package from the statestore.

        :param id: The id of the compute package to get.
        :type id: str
        :return: Package.
        :rtype: dict
        """
        response = requests.get(self._get_url_api_v1(f'packages/{id}'), verify=self.verify, headers=self.headers)

        _json = response.json()

        return _json

    def get_packages(self, n_max: int = None):
        """ Get compute packages from the statestore.

        :param n_max: The maximum number of packages to get (If none all will be fetched).
        :type n_max: int
        :return: Packages.
        :rtype: dict
        """
        _headers = self.headers.copy()

        if n_max:
            _headers['X-Limit'] = str(n_max)

        response = requests.get(self._get_url_api_v1('packages/'), verify=self.verify, headers=_headers)

        _json = response.json()

        return _json

    def get_packages_count(self):
        """ Get the number of compute packages in the statestore.

        :return: The number of packages.
        :rtype: dict
        """
        response = requests.get(self._get_url_api_v1('packages/count'), verify=self.verify, headers=self.headers)

        _json = response.json()

        return _json

    def get_active_package(self):
        """ Get the (active) compute package from the statestore.

        :return: Package.
        :rtype: dict
        """
        response = requests.get(self._get_url_api_v1('packages/active'), verify=self.verify, headers=self.headers)

        _json = response.json()

        return _json

    def get_package_checksum(self):
        """ Get the checksum of the compute package.

        :return: The checksum.
        :rtype: dict
        """
        response = requests.get(self._get_url('get_package_checksum'), verify=self.verify, headers=self.headers)

        _json = response.json()

        return _json

    def download_package(self, path: str):
        """ Download the compute package.

        :param path: The path to download the compute package to.
        :type path: str
        :return: Message with success or failure.
        :rtype: dict
        """
        response = requests.get(self._get_url('download_package'), verify=self.verify, headers=self.headers)
        if response.status_code == 200:
            with open(path, 'wb') as file:
                file.write(response.content)
            return {'success': True, 'message': 'Package downloaded successfully.'}
        else:
            return {'success': False, 'message': 'Failed to download package.'}

    def set_active_package(self, path: str, helper: str, name: str = None, description: str = None):
        """ Set the compute package in the statestore.

        :param path: The file path of the compute package to set.
        :type path: str
        :param helper: The helper type to use.
        :type helper: str
        :return: A dict with success or failure message.
        :rtype: dict
        """
        with open(path, 'rb') as file:
            response = requests.post(self._get_url('set_package'), files={'file': file}, data={
                                     'helper': helper, 'name': name, 'description': description}, verify=self.verify, headers=self.headers)

        _json = response.json()

        return _json

    # --- Rounds --- #

    def get_round(self, id: str):
        """ Get a round from the statestore.

        :param round_id: The round id to get.
        :type round_id: str
        :return: Round (config and metrics).
        :rtype: dict
        """
        response = requests.get(self._get_url_api_v1(f'rounds/{id}'), verify=self.verify, headers=self.headers)

        _json = response.json()

        return _json

    def get_rounds(self, n_max: int = None):
        """ Get all rounds from the statestore.

        :param n_max: The maximum number of rounds to get (If none all will be fetched).
        :type n_max: int
        :return: Rounds.
        :rtype: dict
        """
        _headers = self.headers.copy()

        if n_max:
            _headers['X-Limit'] = str(n_max)

        response = requests.get(self._get_url_api_v1('rounds/'), verify=self.verify, headers=_headers)

        _json = response.json()

        return _json

    def get_rounds_count(self):
        """ Get the number of rounds in the statestore.

        :return: The number of rounds.
        :rtype: dict
        """
        response = requests.get(self._get_url_api_v1('rounds/count'), verify=self.verify, headers=self.headers)

        _json = response.json()

        return _json

    # --- Sessions --- #

    def get_session(self, id: str):
        """ Get a session from the statestore.

        :param id: The session id to get.
        :type id: str
        :return: Session.
        :rtype: dict
        """

        response = requests.get(self._get_url_api_v1(f'sessions/{id}'), self.verify, headers=self.headers)

        _json = response.json()

        return _json

    def get_sessions(self, n_max: int = None):
        """ Get sessions from the statestore.

        :param n_max: The maximum number of sessions to get (If none all will be fetched).
        :type n_max: int
        :return: Sessions.
        :rtype: dict
        """
        _headers = self.headers.copy()

        if n_max:
            _headers['X-Limit'] = str(n_max)

        response = requests.get(self._get_url_api_v1('sessions/'), verify=self.verify, headers=_headers)

        _json = response.json()

        return _json

    def get_sessions_count(self):
        """ Get the number of sessions in the statestore.

        :return: The number of sessions.
        :rtype: dict
        """
        response = requests.get(self._get_url_api_v1('sessions/count'), verify=self.verify, headers=self.headers)

        _json = response.json()

        return _json

    def get_session_status(self, id: str):
        """ Get the status of a session.

        :param id: The id of the session to get.
        :type id: str
        :return: The status of the session.
        :rtype: str
        """
        session = self.get_session(id)

        if session and "status" in session:
            return session["status"]

        return "Could not retrieve session status."

    def session_is_finished(self, id: str):
        """ Check if a session with id has finished.

        :param id: The id of the session to get.
        :type id: str
        :return: True if session is finished, otherwise false.
        :rtype: bool
        """
        status = self.get_session_status(id)
        return status and status.lower() == "finished"

    def start_session(
            self,
            id: str = None,
            aggregator: str = 'fedavg',
            model_id:  str = None,
            round_timeout: int = 180,
            rounds: int = 5,
            round_buffer_size: int = -1,
            delete_models: bool = True,
            validate: bool = True, helper: str = 'numpyhelper', min_clients: int = 1, requested_clients: int = 8
    ):
        """ Start a new session.

        :param id: The session id to start.
        :type id: str
        :param aggregator: The aggregator plugin to use.
        :type aggregator: str
        :param model_id: The id of the initial model.
        :type model_id: str
        :param round_timeout: The round timeout to use in seconds.
        :type round_timeout: int
        :param rounds: The number of rounds to perform.
        :type rounds: int
        :param round_buffer_size: The round buffer size to use.
        :type round_buffer_size: int
        :param delete_models: Whether to delete models after each round at combiner (save storage).
        :type delete_models: bool
        :param validate: Whether to validate the model after each round.
        :type validate: bool
        :param helper: The helper type to use.
        :type helper: str
        :param min_clients: The minimum number of clients required.
        :type min_clients: int
        :param requested_clients: The requested number of clients.
        :type requested_clients: int
        :return: A dict with success or failure message and session config.
        :rtype: dict
        """
        response = requests.post(self._get_url('start_session'), json={
            'session_id': id,
            'aggregator': aggregator,
            'model_id': model_id,
            'round_timeout': round_timeout,
            'rounds': rounds,
            'round_buffer_size': round_buffer_size,
            'delete_models': delete_models,
            'validate': validate,
            'helper': helper,
            'min_clients': min_clients,
            'requested_clients': requested_clients
        }, verify=self.verify, headers=self.headers)

        _json = response.json()

        return _json

    # --- Statuses --- #

    def get_status(self, id: str):
        """ Get a status object (event) from the statestore.

        :param id: The id of the status to get.
        :type id: str
        :return: Status.
        :rtype: dict
        """
        response = requests.get(self._get_url_api_v1(f'statuses/{id}'), verify=self.verify, headers=self.headers)

        _json = response.json()

        return _json

    def get_statuses(self, session_id: str = None, event_type: str = None, sender_name: str = None, sender_role: str = None, n_max: int = None):
        """ Get statuses from the statestore. Filter by input parameters

        :param session_id: The session id to get statuses for.
        :type session_id: str
        :param event_type: The event type to get.
        :type event_type: str
        :param sender_name: The sender name to get.
        :type sender_name: str
        :param sender_role: The sender role to get.
        :type sender_role: str
        :param n_max: The maximum number of statuses to get (If none all will be fetched).
        :type n_max: int
        :return: Statuses
        """
        _params = {}

        if session_id:
            _params["session_id"] = session_id

        if event_type:
            _params["type"] = event_type

        if sender_name:
            _params["sender.name"] = sender_name

        if sender_role:
            _params["sender.role"] = sender_role

        _headers = self.headers.copy()

        if n_max:
            _headers['X-Limit'] = str(n_max)

        response = requests.get(self._get_url_api_v1('statuses/'), params=_params, verify=self.verify, headers=_headers)

        _json = response.json()

        return _json

    def get_statuses_count(self):
        """ Get the number of statuses in the statestore.

        :return: The number of statuses.
        :rtype: dict
        """
        response = requests.get(self._get_url_api_v1('statuses/count'), verify=self.verify, headers=self.headers)

        _json = response.json()

        return _json

    # --- Validations --- #

    def get_validation(self, id: str):
        """ Get a validation from the statestore.

        :param id: The id of the validation to get.
        :type id: str
        :return: Validation.
        :rtype: dict
        """
        response = requests.get(self._get_url_api_v1(f'validations/{id}'), verify=self.verify, headers=self.headers)

        _json = response.json()

        return _json

    def get_validations(
        self,
        session_id: str = None,
        model_id: str = None,
        correlation_id: str = None,
        sender_name: str = None,
        sender_role: str = None,
        receiver_name: str = None,
        receiver_role: str = None,
        n_max: int = None
    ):
        """ Get validations from the statestore. Filter by input parameters.

        :param session_id: The session id to get validations for.
        :type session_id: str
        :param model_id: The model id to get validations for.
        :type model_id: str
        :param correlation_id: The correlation id to get validations for.
        :type correlation_id: str
        :param sender_name: The sender name to get validations for.
        :type sender_name: str
        :param sender_role: The sender role to get validations for.
        :type sender_role: str
        :param receiver_name: The receiver name to get validations for.
        :type receiver_name: str
        :param receiver_role: The receiver role to get validations for.
        :type receiver_role: str
        :param n_max: The maximum number of validations to get (If none all will be fetched).
        :type n_max: int
        :return: Validations.
        :rtype: dict
        """
        _params = {}

        if session_id:
            _params["sessionId"] = session_id

        if model_id:
            _params["modelId"] = model_id

        if correlation_id:
            _params["correlationId"] = correlation_id

        if sender_name:
            _params["sender.name"] = sender_name

        if sender_role:
            _params["sender.role"] = sender_role

        if receiver_name:
            _params["receiver.name"] = receiver_name

        if receiver_role:
            _params["receiver.role"] = receiver_role

        _headers = self.headers.copy()

        if n_max:
            _headers['X-Limit'] = str(n_max)

        response = requests.get(self._get_url_api_v1('validations/'), params=_params, verify=self.verify, headers=_headers)

        _json = response.json()

        return _json

    def get_validations_count(self):
        """ Get the number of validations in the statestore.

        :return: The number of validations.
        :rtype: dict
        """
        response = requests.get(self._get_url_api_v1('validations/count'), verify=self.verify, headers=self.headers)

        _json = response.json()

        return _json
