import os

# This script is used to upload logs during task bootstrapping, so
# it shouldn't have external dependencies besides Metaflow itself
# (e.g. no click for parsing CLI args).
from metaflow.datastore import DATASTORES
from metaflow.util import Path
from . import TASK_LOG_SOURCE

SMALL_FILE_LIMIT = 1024*1024

def save_logs():
    def _read_file(path):
        with open(path, 'rb') as f:
            return f.read()

    # these env vars are set by mflog.mflog_env
    pathspec = os.environ['MF_PATHSPEC']
    attempt = os.environ['MF_ATTEMPT']
    ds_type = os.environ['MF_DATASTORE']
    ds_root = os.environ.get('MF_DATASTORE_ROOT')
    paths = (os.environ['MFLOG_STDOUT'],\
             os.environ['MFLOG_STDERR'])

    flow_name, run_id, step_name, task_id = pathspec.split('/')
    Datastore = DATASTORES[ds_type]
    if ds_root is None:
        def print_clean(line, **kwargs):
            pass
        ds_root = Datastore.get_datastore_root_from_config(print_clean)
    Datastore.datastore_root = ds_root

    ds = Datastore(flow_name,
                   run_id=run_id,
                   step_name=step_name,
                   task_id=task_id,
                   attempt=int(attempt),
                   mode='w')

    try:
        streams = ('stdout', 'stderr')
        sizes = [(stream, path, os.path.getsize(path))
                 for stream, path in zip(streams, paths)
                 if os.path.exists(path)]

        if max(size for _, _, size in sizes) < SMALL_FILE_LIMIT:
            op = _read_file
        else:
            op = Path

        data = [(stream, op(path)) for stream, path, _ in sizes]
        ds.save_logs(TASK_LOG_SOURCE, data)
    except:
        # Upload failing is not considered a fatal error.
        # This script shouldn't return non-zero exit codes
        # for transient errors.
        pass

if __name__ == '__main__':
    save_logs()
    #to debug delays in logs, comment the line above and uncomment
    #this snippet:
    """
    import sys
    from metaflow.metaflow_profile import profile
    d = {}
    with profile('save_logs', stats_dict=d):
        save_logs()
    print('Save logs took %dms' % d['save_logs'], file=sys.stderr)
    """