import os
import sys

import pyhectiqlab
from pyhectiqlab.metrics import MetricsManager
from pyhectiqlab.events_manager import EventsManager
from pyhectiqlab.stream import LogHandler
from pyhectiqlab.buffer import Buffer
from pyhectiqlab.watermark import Watermark
from pyhectiqlab.artifacts import SharedArtifactsManager
from pyhectiqlab.notebooks import get_notebook_file
from pyhectiqlab.utils import list_all_files_in_dir

from pyhectiqlab.mlmodels import download_mlmodel as ops_download_mlmodel
from pyhectiqlab.mlmodels import upload_mlmodel as ops_upload_mlmodel
from pyhectiqlab.datasets import download_dataset as ops_download_dataset
from pyhectiqlab.datasets import upload_dataset as ops_upload_dataset


from functools import partial
from sys import getsizeof
import packaging 
import tempfile
import logging
import shutil

class Run():
    
    def __init__(self, name:str, project:str=None, project_name:str=None, should_create_project = False, dry=False):
        """
        dry: bool: If true, the run does not send anything.
        should_create_project: bool. If true, a project is created if if does not exists
        """
        self._mock = dry
        self.events_manager = EventsManager(self._mock)
        # Check the python version
        self.events_manager.compare_python_version()    
        if project_name is not None:
            print('Argument `project_name` will be remove in a future release. Please use `project` instead.')
            project = project_name
        if dry:
            print("Dry mode activated.")
            self.run_view = {"id": "null"}
            self._mock = True
            self.events_manager = EventsManager(self._mock)
        elif self.events_manager.is_logged()==False:
            print("User not authentificated. Use `pyhectiqlab.login()` for a web prompt or `hectiqlab login` in command line.")
            self._failed_init()
        else:
            # try:
            if project is None:
                projects = self.events_manager.add_event("get_all_projects", args=(), auth=True, async_method=False)
                if len(projects["result"])==1:
                    project = projects["result"][0]["name"]
                    print(f"Connecting to project {project}.")
                    self.project_view = projects["result"][0]
                else:
                    print("User has access to multiple projects. Switching to dry mode. \n--------\nFor active mode, please use a project:")
                    for el in projects["result"]:
                        print(f"{el['name']}")
                    self._failed_init()
                    pass
            else:
                self.project_view = self.events_manager.add_event("create_project", args=(project, should_create_project), auth=True, async_method=False)
                if self.project_view is None:
                    print("Project does not exists. Dry mode activated. \n--------\nUse should_create_project=True to create a project.")
                    self._failed_init()
                    pass
                if self.project_view['status_code']==401:
                    self._failed_init()
                    pass
            if project:
                self.run_view = self.events_manager.add_event("create_run", args=(name, project), auth=True, async_method=False)
                try:
                    if self.run_view["is_new"]==False:
                        print(f"Attaching to an existing run (ID :{self.run_view['p_id']}).")
                    if self.run_view['readonly']:
                        print(f"Run is read only.")
                    self._read_only = self.run_view['readonly']
                except:
                    print("An error occured while creating the run. Switching to dry mode.")
                    self._failed_init()

                if "id" not in self.run_view:
                    print("An error occured while creating the run. Switching to dry mode.")
                    self._failed_init()

        self.name = name
        self.metrics_manager = MetricsManager(run_id=self.run_view["id"], push_method=self._push_metrics)
        self.metrics_manager.set_aggr('mean')
        self.metrics_manager.update_cache_settings(max_cache_timeout=20,
                                                    min_cache_flush_delay=5, 
                                                    max_cache_length=25)
        self.log_recorder = None
        self.is_recording_logs = False
        self.logs_buffer = Buffer(push_method=self._push_logs)
        return

    @property
    def read_only(self):
        return self._read_only

    def set_metrics_aggregate_method(self, aggr):
        """
        Set the aggregate method for the metrics manager. 

        aggr [str]: One of ['none', 'sum', 'max', 'mean']
        """
        print('`set_metrics_aggregate_method` will be remove in a future release. Please use `set_metrics_cache_settings` instead.')
        self.metrics_manager.set_aggr(aggr)

    def set_metrics_cache_settings(self, aggr: str = 'mean', max_cache_timeout=None, max_cache_length=None, min_cache_flush_delay=None):
        """
        Update the settings of the metrics manager.

        max_cache_timeout [int]: Time interval for flushing the cache (seconds)  (default is 20)
        max_cache_length [int]: Max number of elements before flushing the cache (default is 25)
        min_cache_flush_delay [float]: Minimum delay between to successive data send (for the same key). 
                If the number of matrix is added faster than this delay, the cache will increase to
                respect the delay. (default is 5 seconds)
        aggr [str]: The aggregate method. One of ['none', 'sum', 'max', 'mean']  (default is mean)
        """
        self.metrics_manager.set_aggr(aggr)
        self.metrics_manager.update_cache_settings(max_cache_timeout=max_cache_timeout, 
                                                    max_cache_length=max_cache_length,
                                                    min_cache_flush_delay=min_cache_flush_delay)

    def _failed_init(self):
        self.run_view = {"id": "null"}
        self._mock = True
        self.events_manager = EventsManager(self._mock)
        self._read_only = True

    def add_artifacts(self, filepath: str, step: int = None):
        """Log a file as an artifacts. If a file already exists with the name, it will be overwritten.

        filepath [str] : Path to the file.
        """
        assert self._read_only==False, 'Run is in read only.'
        filename = os.path.basename(filepath)
        num_bytes = os.path.getsize(filepath)
        content_bytes = open(filepath, "rb")
        args = (self._id, filename, content_bytes, num_bytes, step)
        self.events_manager.add_event("add_file", args)

    def add_mlmodel(self, source_path: str, name: str, version: str = None, short_description: str = None,
                push_dir: bool = False, show_filenames: bool = False):
        """Add a mlmodel
        """
        assert self._read_only==False, 'Run is in read only.'
        ops_upload_mlmodel(source_path=source_path, 
            mlmodel_name=name,
             run_id=self._id, 
             version=version,
             short_description=short_description,
             push_dir=push_dir)

    def download_mlmodel(self, mlmodel_name: str, version: str = None, save_path: str = './', overwrite:bool = False):

        dir_path = ops_download_mlmodel(mlmodel_name=mlmodel_name, 
                        project_id=self.project_view["id"], 
                        version=version, 
                        save_path=save_path, 
                        overwrite=overwrite)
        if dir_path is not None:
            print(dirpath) 

    def add_dataset(self, source_path: str, name: str, version: str = None, short_description: str = None,
                push_dir: bool = False, show_filenames: bool = False):
        """Add a mlmodel
        """
        assert self._read_only==False, 'Run is in read only.'
        ops_upload_dataset(source_path=source_path, 
            dataset_name=name,
             run_id=self._id, 
             version=version,
             short_description=short_description,
             push_dir=push_dir)

    def download_dataset(self, dataset_name: str, version: str = None, save_path: str = './', overwrite:bool = False):

        dir_path = ops_download_dataset(dataset_name=dataset_name, 
                        project_id=self.project_view["id"], 
                        version=version, 
                        save_path=save_path, 
                        overwrite=overwrite)
        if dir_path is not None:
            print(dirpath)        

    def add_config(self, config, prefix=""):
        assert self._read_only==False, 'Run is in read only.'
        config.push_as_meta(self, prefix=os.path.join("config/",prefix))

    def add_tf_model(self, model, filename, step: int = None):
        assert self._read_only==False, 'Run is in read only.'
        save_path = tempfile.mkdtemp()
        model.save_weights(f"{save_path}/{filename}")
        p = tempfile.mkdtemp()
        shutil.make_archive(f"{p}/{filename}", 'zip', save_path)
        self.add_artifacts(f"{p}/{filename}.zip", step=step)

    def add_directory_as_zip_artifact(self, dirpath, step: int = None):
        assert self._read_only==False, 'Run is in read only.'
        p = tempfile.mkdtemp()
        filename = os.path.basename(dirpath)
        shutil.make_archive(f"{p}/{filename}", 'zip', root_dir=dirpath)
        self.add_artifacts(f"{p}/{filename}.zip", step=step)

    def add_shared_artifact(self, filename):
        """Push an artifact to be shared accross the project.
        """
        assert self._read_only==False, 'Run is in read only.'
        m = SharedArtifactsManager()
        m.push_artifact(filepath=filename, project_id=self.project_view["id"])

    def list_shared_artifact(self):
        """Returns a list of the shared artifacts for this project.
        """
        m = SharedArtifactsManager()
        return m.list_artifacts(project_id=self.project_view["id"])

    def add_metrics(self, key, value, step):
        """Add a metrics. The metrics is appended even if there already exists a value
        for the same key and step. 

        key [str]: Key used for the metrics. The format with backlash such as `metrics/train`
            is recognized to organize the metrics.

        value [float or string]. If tensor, is casted to float.
        step [int]: Training step. Casted to integer.
        """
        assert self._read_only==False, 'Run is in read only.'
        self.metrics_manager.add(key, float(value), int(step))

    def add_tag(self, name: str, description: str = None, color: str = None):
        """Add tag to a run. If the tag does not already exists by this name,
        a new tag is created with the given description and color.
        """
        assert self._read_only==False, 'Run is in read only.'
        args = (self._id, name, description, color)
        self.events_manager.add_event("add_tag", args)

    def add_package_versions(self, packages=None, with_sys=True, with_python=True, with_git=True):
        """Save the version of the imported packages.

        Use:
        add_package_versions(globals())
        """
        assert self._read_only==False, 'Run is in read only.'
        water = Watermark()
        versions = {}
        if with_python:
            versions["python"] = water.get_pyversions()
        if with_sys:
            versions["system"] = water.get_sysinfo()
        if packages:
            versions["packages"] = water.get_all_import_versions(packages)
        if with_git:
            git_info = water.get_git_info()
            if git_info:
                versions["git"] = water.get_git_info()

        args = (self._id, versions)
        self.events_manager.add_event("push_package_versions", args)

    def add_package_repo_state(self, package: str = None):
        """Save the branch/commit/origin of the git repo in which
        a package is located.

        Example:
            self.add_package_repo_state('pyhectiqlab')
        """
        assert self._read_only==False, 'Run is in read only.'
        water = Watermark()
        result = water.get_repo_info(package)
        data = {package: result}
        args = (self._id, data)
        self.events_manager.add_event("push_git_package_state", args)


    def add_current_notebook(self, stamp: str = 'datetime'):
        """If the run is executed in a jupyter notebook, this
        will quickly add your current notebook as an artifact.

        stamp [str]: Add a stamp on the notebook name.
                    One of 'datetime', 'date', None
        """
        assert self._read_only==False, 'Run is in read only.'
        filename = get_notebook_file(stamp)
        self.add_artifacts(filename)

    def failed(self):
        """Set the run to the failed status.
        """
        assert self._read_only==False, 'Run is in read only.'
        args = (self._id, "failed")
        self.events_manager.add_event("set_run_status", args)

    def _stopped(self):
        """Set the run to the stopped status.
        """
        assert self._read_only==False, 'Run is in read only.'
        args = (self._id, "stopped")
        self.events_manager.add_event("set_run_status", args)

    def completed(self):
        """Set the run to the completed with success status.
        """
        assert self._read_only==False, 'Run is in read only.'
        args = (self._id, "completed")
        self.events_manager.add_event("set_run_status", args)

    def pending(self):
        """Set the run to the pending status.
        """
        assert self._read_only==False, 'Run is in read only.'
        args = (self._id, "pending")
        self.events_manager.add_event("set_run_status", args)

    def running(self):
        """Set the run to the running status.
        """
        assert self._read_only==False, 'Run is in read only.'
        args = (self._id, "running")
        self.events_manager.add_event("set_run_status", args)

    def training(self):
        """Set the run to the training status.
        """
        assert self._read_only==False, 'Run is in read only.'
        args = (self._id, "training")
        self.events_manager.add_event("set_run_status", args)

    def add_meta(self, key, value):
        """Add or update a meta information in the format of
        key/value.
        """ 
        assert self._read_only==False, 'Run is in read only.'
        args = (self._id, key, value)
        self.events_manager.add_event("push_meta", args)

    def set_note(self, text):
        """Add or update a custom note field for the run.
        """
        assert self._read_only==False, 'Run is in read only.'
        args = (self._id, text)
        self.events_manager.add_event("set_note", args)

    def start_recording_logs(self, logger_name=None):
        """Start tailing the logs on a logger name.
        If logger_name is None, the default __name__ will be taken and
        you may use the returned logger.
        """
        assert self._read_only==False, 'Run is in read only.'
        self.is_recording_logs = True
        self.logs_buffer.start(key="logs")
        
        print("Start recording logs.")
        if (self.log_recorder is None) or self.logger_name!=logger_name:
            func = partial(self.logs_buffer.add, key="logs")
            self.log_recorder = LogHandler(func)
            console  = logging.StreamHandler()  
            console.setLevel(logging.DEBUG)

            name = logger_name
            if name is None:
                name = __name__
            self.logger_name = name

            logger = logging.getLogger(name)
            logger.setLevel(logging.DEBUG)
            logger.addHandler(self.log_recorder)

            # Handle all exceptions
            def handle_exception(exc_type, exc_value, exc_traceback):
                if issubclass(exc_type, KeyboardInterrupt):
                    self._stopped()
                else:
                    self.add_tag(name="failed")
                    self.failed()
                    
                logger.error("Exception", exc_info=(exc_type, exc_value, exc_traceback))
                # Flush cache before leaving
                self.logs_buffer.flush_cache()
                self.metrics_manager.flush_cache()
                sys.exit(0)                
                
            sys.excepthook = handle_exception

            return logger
        else:
            return logging.getLogger(self.logger_name)

    def stop_recording_logs(self):
        """Stop recording the logs from the logger. 
        """
        assert self._read_only==False, 'Run is in read only.'
        self.is_recording_logs = False
        self.logs_buffer.stop(key="logs")

    def _push_logs(self, key, msg):
        assert self._read_only==False, 'Run is in read only.'
        if key=="logs":
            if self.is_recording_logs:
                args = (self._id, "".join(msg))
                self.events_manager.add_event("append_logs", args)

    def _push_metrics(self, key, values):
        assert self._read_only==False, 'Run is in read only.'
        args = (self._id, key, values)
        self.events_manager.add_event("push_metrics", args)

    @property
    def _id(self):
        return self.run_view["id"]
    