# MAC address is used to generate IDs
import argparse
import logging
import json
import time
import pickle
from typing import Dict, List, Any, Tuple
import pathlib
import socket
import os
from hashlib import sha256

from getmac import get_mac_address as gma
import numpy as np

from stadle.lib.util.states import ClientState, IDPrefix
from stadle.lib.constants.environment import ConfigFileType
from stadle.lib.env.handler import EnvironmentVar, EnvironmentHandler


def set_env_config(config_type: EnvironmentVar, value):
    assert config_type != None, "config_type cannot be None"

    from stadle import EnvironmentHandler
    env_handler = EnvironmentHandler()

    if value:
        # user provides a custom value for the config file
        env_handler.update_variable(name=config_type.value, value=value)
    else:
        # user doesn't provide a custom value so the default is used.
        if config_type == EnvironmentVar.STADLE_DB_CONFIG_PATH:
            env_handler.update_variable(name=EnvironmentVar.STADLE_DB_CONFIG_PATH.value,
                                        value=os.path.join(env_handler.config_path,
                                                           ConfigFileType.STADLE_CORE_DB.value))
        elif config_type == EnvironmentVar.STADLE_AGGREGATOR_CONFIG_PATH:
            env_handler.update_variable(name=EnvironmentVar.STADLE_AGGREGATOR_CONFIG_PATH.value,
                                        value=os.path.join(env_handler.config_path,
                                                           ConfigFileType.STADLE_AGGREGATOR.value))
        elif config_type == EnvironmentVar.STADLE_AGENT_CONFIG_PATH:
            env_handler.update_variable(name=EnvironmentVar.STADLE_AGENT_CONFIG_PATH.value,
                                        value=os.path.join(env_handler.config_path, ConfigFileType.STADLE_AGENT.value))
        elif config_type == EnvironmentVar.STADLE_CONFIG_PATH:
            env_handler.update_variable(name=EnvironmentVar.STADLE_CONFIG_PATH.value,
                                        value=os.path.join(env_handler.config_path, ConfigFileType.STADLE_AGENT.value))
        else:
            raise ValueError(f"unsupported config type {config_type}")


def read_config(config_path: str) -> Dict[str, Any]:
    """
    Read a JSON configuration file to set up an aggregator
    :param config_path:
    :return: Dict[str,Any] - configs
    """
    with open(config_path) as jf:
        config = json.load(jf)
    return config


def generate_id() -> str:
    """
    Generate a system-wide unique ID based on
    MAC address and Instantiation time with a hash func (SHA256)
    :return: str - ID
    """
    macaddr = gma()
    in_time = time.time()

    raw = f'{macaddr}{in_time}'
    hash_id = sha256(raw.encode('utf-8'))
    return hash_id.hexdigest()


def get_id(data_path, id_file_name) -> str:
    id = None
    # if there is no directory to save agent id
    if not os.path.exists(data_path):
        os.makedirs(data_path)

    id_file_path = f'{data_path}/{id_file_name}'
    if os.path.exists(id_file_path):
        with open(id_file_path, 'r') as f:
            id = f.read()
    else:
        with open(id_file_path, 'w') as f:
            id = generate_id()
            f.write(id)

    return id


def generate_model_id(component_type: str, component_id: str, generation_time: float) -> str:
    """
    Generate a system-wide unique ID for a set of models based on
    (1) a component ID: the id of an entity that created the models and
    (2) generation time: the time the models were created.
    The ID is generated by a hash function (SHA256)
    :param component_type: str - a prefix indicating component type from IDPrefix
    :param component_id: str - ID of an entity that created the models
    :param generation_time: float - the time the models were created
    :return: str - Model ID
    """
    raw = f'{component_type}{component_id}{generation_time}'
    hash_id = sha256(raw.encode('utf-8'))
    return hash_id.hexdigest()


def generate_performance_dict(init_flag: bool = False,
                              performance: float = 0.0,
                              accuracy: float = 0.0,
                              loss_training: float = 0.0,
                              loss_valid: float = 0.0,
                              loss_test: float = 0.0,
                              f_score: float = 0.0,
                              reward: float = 0.0,
                              model_type: int = 0,
                              prev_perf_dict: Dict[str, np.array] = dict()):
    """
    Generate performance dictionary.
    """
    performance_dict = {}
    if init_flag:
        # TODO Initialization of performance data
        return dict()

    if 'performance_id' in prev_perf_dict:
        performance_dict['performance_id'] = prev_perf_dict['performance_id']

    performance_dict['accuracy'] = accuracy
    performance_dict['performance'] = performance
    performance_dict['loss_training'] = loss_training
    performance_dict['loss_valid'] = loss_valid
    performance_dict['loss_test'] = loss_test
    performance_dict['reward'] = reward
    performance_dict['f_score'] = f_score
    performance_dict['model_type'] = model_type

    return performance_dict


def compatible_data_dict_read(data_dict: Dict[str, Any]) -> List[Any]:
    # ID init
    # for compatibility with older versions
    if 'my_id' in data_dict.keys():
        id = data_dict['my_id']
    else:
        id = generate_id()

    if 'gene_time' in data_dict.keys():
        gene_time = data_dict['gene_time']
    else:
        gene_time = time.time()

    if 'models' in data_dict.keys():
        models = data_dict['models']
        version = 0.6
    else:
        models = data_dict
        version = 0.5

    if 'model_id' in data_dict.keys():
        model_id = data_dict['model_id']
    else:
        model_id = generate_model_id(IDPrefix.agent, id, gene_time)

    if 'base_model' in data_dict.keys():
        base_model = data_dict['base_model']
    else:
        base_model = None

    return id, gene_time, models, model_id, version, base_model


def compatible_performance_dict_read(performance_dict: Dict[str, Any]) -> List[Any]:
    # ID init

    gene_time = time.time()

    if 'performance_id' in performance_dict.keys():
        performance_id = performance_dict['performance_id']
    else:
        id = generate_id()
        performance_id = generate_model_id(IDPrefix.agent, id, gene_time)
        performance_dict['performance_id'] = performance_id

    return performance_id, performance_dict


def save_model_file(data_dict: Dict[str, Any],
                    path: str,
                    name: str,
                    performance_dict: Dict[str, float] = dict()):
    """
    Save a given set of models into a local file
    :param data_dict: Dict[str,np.array] - model_id, models
    :param path: str - path to the directory
    :param name: str - model file name
    :param performance_dict: Dict[str,float] - each entry shows a pair of a model id and its performance
    :return:
    """
    data_dict['performance'] = performance_dict

    fname = f'{path}/{name}'
    with open(fname, 'wb') as f:
        pickle.dump(data_dict, f)


def load_model_file(path: str, name: str) -> Tuple[Dict[str, Any], Dict[str, float]]:
    """
    Read a local model file and return the unpickled models
    :param path: str - path to the directory
    :param name: str - model file name
    :return: Dict[str,np.array] - models
    """
    fname = f'{path}/{name}'
    with open(fname, 'rb') as f:
        data_dict = pickle.load(f)

    performance_dict = data_dict.pop('performance')

    # data_dict only includes models
    return data_dict, performance_dict


def is_sgm_received(path: str, name: str) -> bool:
    fname = f'{path}/{name}'

    if os.path.isfile(fname):
        return True
    else:
        return False


def read_state(path: str, name: str) -> ClientState:
    """
    Read a local state file and return a Client state
    :param path: str - path to the directory
    :param name: str - model file name
    :return: ClientState - A state indicated in the file
    """
    fname = f'{path}/{name}'
    with open(fname, 'r') as f:
        st = f.read()

    # In case, the file is being written at a time,
    # Retry reading after 0.01 s
    if st == '':
        time.sleep(0.01)
        return read_state(path, name)

    return int(st)


def write_state(path: str, name: str, state: ClientState):
    """
    Change the state on the state file
    :param path: str - path to the directory
    :param name: str - model file name
    :param state: ClientState - A new Client state
    :return:
    """
    fname = f'{path}/{name}'
    with open(fname, 'w') as f:
        f.write(str(int(state)))


def get_ip():
    s = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
    try:
        # doesn't even have to be reachable
        s.connect(('1.1.1.1', 1))
        ip = s.getsockname()[0]
    except:
        ip = '127.0.0.1'
    finally:
        s.close()
    return ip


def client_arg_parser():
    parser = argparse.ArgumentParser(description='Evaluate Client arguments')
    parser.add_argument('--config_path', metavar='cfp', type=str,
                        help='admin agent configuration path')
    parser.add_argument('--simulation', action='store_true', required=False,
                        help='simulation or not (if flag added simulation mode)')
    parser.add_argument('--agent_name', metavar='ag_n', type=str, help='name to use for agent')
    parser.add_argument('--comm_protocol', metavar='com', type=str, help="Protocol to use for sending/receiving messages")
    parser.add_argument('--aggregator_ip', metavar='ip', type=str, help="Ip address of the aggregator server")
    parser.add_argument('--reg_port', metavar='rp', type=str, help="registration port")
    parser.add_argument('--exch_port', metavar='ep', type=str, help="exchange port")
    parser.add_argument('--token', metavar='t', type=str, help="token for validation")
    parser.add_argument('--model_path', metavar='m_path', type=str, help="model path")
    parser.add_argument('--agent_running', action='store_true', required=False, default=False,
                        help='agent running or not (if flag added simulation mode)')

    args = parser.parse_args()
    return args
