import os
import sys
import json
import re
import pandas as pd
from s3fs import S3FileSystem

from datetime import datetime, timedelta

class Dataset():
    """
    Class for specifying a dataset. This is typically
    an input to the transforms.

    Usage::

        # A dataset that has one directory for each day. Within that
        # there are atleast two sub-datasets
        Dataset(params={
            "name": 'athena-availability',
            "paths": [
                {
                    "name": "test",
                    "nature": "local",
                    "path": "%(data_root)s/shared/datasets/athena/v2/availability",

                },
                {
                    "name": "local",
                    "nature": "local",
                    "path": "%(enrich_data_dir)s/acme/Marketing/shared/datasets/athena/v2/availability",
                },
                {
                    "name": "test",
                    "nature": "local",
                    "path": "%(backup_root)s/%(node)s/data/acme/Marketing/shared/datasets/athena/v2/availability",
                },
            ],
            "match": {
                "generate": "generate_datetime_daily",
                "compare": "compare_datetime_pattern_match_range",
                "pattern": "%Y-%m-%d",
            },
            "subsets": [
                {
                    "name": "CatalogScore",
                    "pattern": ".*catalogscore.csv$",
                    "description": "Score of all products in the catalog"
                },
                {
                    "name": "ProductWeight",
                    "pattern": ".*productweight.csv$",
                    "description": "Assortment weight for each product"
                }
            ]
        })

    """

    def __init__(self, params, *args, **kwargs):

        super().__init__(*args, **kwargs)

        # Backward compatability
        if 'paths' not in params:
            params['paths'] = []
            for name in ['local', 'test', 'backup']:
                if name in params:
                    params['paths'].append({
                        'name': name,
                        'nature': 's3' if name == 'backup' else 'local',
                        'path': params[name]
                    })

        self.validate(params)

        self.isfile = params.get('isfile', False)
        """
        Specifies whether the dataset is simple (a file for each
        run) or more complex directory hierarchy.

        Example::

            # A complex dataset
            Dataset(params={
                "name": 'athena-availability',
                "isfile": False,
                ...

        """
        self.issortable = params.get('issortable', True)

        self.subsets = params.get('subsets',[])
        """
        A dataset oftentimes has multiple components. Define
        them here. We can give a name to each.

        Example::

          {
             ...
             "subsets": [
                 {
                     "name": "v1culltable",
                     "pattern": ".*v1culltable.csv$"
                 },...
              ]
          }
        """


        self.name = params['name']
        """
        Unique name of the dataset, e.g., Events
        """

        self.description = params.get('description','')
        """
        str: Description
        Example: "Daily events for Alpha Sensor"
        """

        self.resolve = kwargs.get('resolve',
                                  params.get('resolve',{}))
        """
        Parameters to resolve paths
        """

        self.paths = params.get('paths',{})
        """
        dict: Paths to be resolved
        Example::

            [
              {
                "name": "test",
                "nature": "local",
                "path": "%(data_root)s/shared/datasets/athena/v2/availability"
              }
           ]
        """

        self.match = params.get('match',{})
        """
        dict: Generating and matching rules

        Example::

            {
                "generate": "generate_datetime_daily",
                "match": "match_datetime_pattern_range",
                "pattern": "plpevents-%Y%m%d-%H%M%S",
            }
        """

    def sample(self, filename, safe=True, fd=None):
        """
        Sample a file belonging to this dataset.Subclass and
        overload this function if the data is sensitive.

        Args:
            filename (str): A file that belongs to this dataset
            safe (bool): Whether file is trusted
            fd (object): File descriptor for s3/gcs/other non-filesystems

        """

        if fd is None:
            # Text formats..
            if filename.split(".")[-1].lower() in ['csv', 'tsv']:
                fd = open(filename)
            else:
                fd = open(filename, 'rb')

        # Read the file
        if filename.lower().endswith('.csv'):
            df = pd.read_csv(fd, nrows=10)
        elif filename.lower().endswith('.tsv'):
            df = pd.read_csv(fd, nrows=10, sep='\t')
        else:
            df = pd.read_parquet(fd)
            df = df.head(10)

        return df

    @property
    def local(self):
        """
        Read the path with the name 'local'. Available for
        backward compatability
        """
        return self.get_path('local')

    @property
    def test(self):
        """
        Read the path with the name 'test'. Available for
        backward compatability
        """
        return self.get_path('test')

    @property
    def backup(self):
        """
        Read the path with the name 'backup'. Available for
        backward compatability
        """
        return self.get_path('backup')

    def get_paths(self):
        """
        Return the available paths (full list of dictionaries,
        not just the names
        """
        return self.paths

    def get_path_by_name(self,name, full=False, resolve=None):
        """
        Return the path definition for a path with a given name
        """

        if resolve is None:
            resolve = self.resolve

        selected = None
        for p in self.paths:
            if p['name'] == name:
                selected = p
                break
        if selected is None:
            raise Exception("Path with name {} not present".format(name))

        if full:
            return selected

        path = selected['path']

        try:
            if ((resolve is not None) and
                (isinstance(resolve, dict)) and
                (len(resolve) > 0)):
                path = path % resolve
        except:
            raise Exception("Path could not be resolved")

        return path

    def get_path_by_nature(self, nature, full=True, resolve=None):

        if resolve is None:
            resolve = self.resolve

        selected = None
        for p in self.paths:
            if p['nature'] == nature:
                selected = p
                break

        if selected is None:
            raise Exception("Path with nature {} not present".format(nature))

        if full:
            return selected

        path = selected['path']

        try:
            if ((resolve is not None) and
                (isinstance(resolve, dict)) and
                (len(resolve) > 0)):
                path = path % resolve
        except:
            raise Exception("Path could not be resolved")

        return path

    def has_subsets(self):
        """
        Does this dataset have components?
        """
        return len(self.subsets) > 0

    def get_subsets(self):
        """
        Return the names of subsets available
        """
        return [s['name'] for s in self.subsets]

    def get_subset_description(self,name):
        """
        Get the description for a subset
        """
        for s in self.subsets:
            if s['name'] == name:
                return s.get('description', self.description)

        raise Exception("Unknown subset: {}".format(name))

    def in_subset(self, name, spec):
        """
        Return the names of subsets available
        """

        if ((not isinstance(spec, dict)) or
            (len(spec) == 0) or
            ('filename' not in spec)):
            raise Exception("invalid subset check input. Expecting dict with filename")

        # Does the subset name exist?
        subset = None
        for s in self.subsets:
            if name == s['name']:
                subset = s
                break
        if subset is None:
            return False

        #
        if 'pattern' not in subset:
            return False

        filename = spec['filename']
        pattern = subset['pattern']
        return re.search(pattern, filename) is not None

    def get_path(self, name, full=False, resolve=None):
        return self.get_path_by_name(name, full, resolve)

    def set_resolve(self, resolve):
        assert isinstance(resolve, dict)
        assert len(resolve) > 0
        self.resolve.update(resolve)

    def validate(self, params):
        """
        Validate dataset arguments. It should be a dictionary
        with atleast three elements: name, match, and paths.

        Args:
          params (dict): Parameters for dataset

        Params dict should have the following:

          * name: string
          * match: dictionary with generate function (name or lambda), are function (name or lambda), and pattern (string).

        """
        if not isinstance(params, dict):
            raise Exception("Invalid parameters for dataset")

        for col in ['name', 'match', 'paths']:
            if col not in params:
                raise Exception("Data specification should have {}".format(col))

        if (not isinstance(params['paths'], list)):
            raise Exception("Paths should be a list of dictionaries")

        for elem in params['paths']:
            if not isinstance(elem, dict) or len(elem) == 0:
                raise Exception("paths element should be non-empty dict")
            if (('nature' not in elem) or (elem['nature'] not in ['local', 'remote', 's3', 'gcs'])):
                raise Exception("paths element have nature (local|remote|s3|gcs)")
            if 'path' not in elem:
                raise Exception("paths element should have a 'path'")

        if not isinstance(params['match'], dict):
            raise Exception("Match should be a dictionary")

        if 'pattern' not in params['match']:
            raise Exception("Matching requirement match.pattern missing")

        pattern = params['match']['pattern']
        if not isinstance(pattern, str) and not callable(pattern):
            raise Exception("Invalid pattern specification")

        generate = params['match'].get('generate', 'generate_datetime_daily')

        if (((isinstance(generate, str)) and (not hasattr(self, generate))) and
            (not callable(generate))):
            raise Exception("Matching requirement {} missing/invalid".format(generate))

        compare = params['match'].get('compare', 'compare_datetime_pattern_match_range')
        if (((isinstance(compare, str)) and (not hasattr(self, compare))) and
            (not callable(compare))):
            raise Exception("Matching requirement {} missing/invalid".format(compare))


    def __str__(self):
        return self.name


    def generate(self, start, end=None):
        """
        Given a start and an end, generate the datetime
        objects corresponding to each run

        Args:
           start (datetime): Start of the time range
           end (datetime): End of range. If None, default = start

        Returns:
           list (dict): List of dictionaries (name, timestamp)

        """
        if ((not isinstance(start, datetime))  or
            ((end is not None) and (not isinstance(end, datetime)))):
            raise Exception("Start or end is not a datetime")

        # One day
        if end is None:
            end = start

        # Handle ordering...
        if end < start:
            start, end = end, start

        func = self.match['generate']
        if isinstance(func,str):
            func = getattr(self, func)

        # Generate all possible names
        names = func(start, end)

        # Check if the name is in the names..
        return names

    def generate_datetime_daily(self, start, end):
        """
        Given a start and an end, generate the datetime
        objects corresponding to each run for each day.

        Args:
           start (datetime): Start of the time range
           end (datetime): End of range. If None, default = start
        """

        pattern = self.match['pattern']

        names = []
        current = start
        while current <= end:
            if isinstance(pattern, str):
                name = current.strftime(pattern)
            elif callable(pattern):
                name = pattern(current)
            else:
                raise Exception("Unsupported pattern type")

            names.append({
                'timestamp': current.isoformat(),
                'name': name
            })
            current += timedelta(days=1)

        return names

    def compare_datetime_pattern_match_range(self, name, start, end):
        """
        Given a start and an end, generate the datetime
        objects corresponding to each run for each day, and
        check whether a named file/directory exists in that list.

        Args:
           start (datetime): Start of the time range
           end (datetime): End of range. If None, default = start
        """

        if not isinstance(name, str):
            return False

        # Generate all possible <timestamp>: <name> combinations
        names = self.generate(start, end)

        # Check if the name is in the names or in the range. The
        # latter is required if there are multiple runs in a given day
        #
        return ((name in names.values()) or
                ((name >= min(names) and name <= max(names))))

    def match_content(self, fs, localdir, backupdir, dirname):
        """
        Check whether local path is replicated in the  s3/blob store

        Args:
           fs (object): s3fs handle
           localdir (str): Root local dir to check
           backupdir (str): Root remote dir in block store to check
           dirname (str): path within the localdir
        """

        localpath = os.path.join(localdir, dirname)
        if os.path.isfile(localpath):
            filesize = getsize(localpath)
            backupfile = os.path.join(backupdir, dirname)
            try:
                detail = fs.info(backupfile)
            except FileNotFoundError:
                return False, "Missing file in backup ({})".format(dirname)
            except Exception as e:
                return False, str(e)

            if filesize != detail['size']:
                return False, "Filesize mismatch: {}".format(dirname)

            return True, "Matched {} (size: {}M)".format(dirname, round(filesize/1000000, 1))

        matched = 0
        totalsize = 0
        for root, dirs, files in os.walk(localpath):
            for name in files:
                filename = os.path.join(root, name)
                relpath = os.path.relpath(filename, start=localpath)
                backupfile = os.path.join(backupdir, dirname, relpath)
                filesize = getsize(filename)
                try:
                    detail = self.fs.info(backupfile)
                except FileNotFoundError:
                    return False, "Missing file in backup ({})".format(relpath)
                except Exception as e:
                    return False, str(e)

                if filesize != detail['size']:
                    return False, "Filesize mismatch: {}".format(relpath)

                matched += 1
                totalsize += filesize

        return True, "Matched {} (size: {}M)".format(matched, round(totalsize/1000000, 1))


    def compare(self, name, start, end=None):
        """
        Given a start and an end, generate the datetime
        objects corresponding to each run

        Args:
           name (str): Directory name
           start (datetime): Start of the time range
           end (datetime): End of range. If None, default = start

        """
        if ((not isinstance(start, datetime))  or
            ((end is not None) and (not isinstance(end, datetime)))):
            raise Exception("Generate names in the range")

        # One day
        if end is None:
            end = start

        # Handle ordering...
        if end < start:
            start, end = end, start

        func = params['match'].get('compare', 'compare_datetime_pattern_match_range')
        if isinstance(func,str):
            func = getattr(self, func)

        # Generate all possible names
        return func(name, start, end)


    def listdir(self, name,
                fshandle=None,
                resolve=None,
                detail=False):
        """
        List the dataset directory
        """

        selected = None
        for p in self.paths:
            if p['name'] == name:
                selected = p
                break

        if selected is None:
            raise Exception("Unknown name of path: {}".format(name))

        if ((selected['nature'] in ['s3', 'gcs']) and
            ((fshandle is None) or (not isinstance(fshandle, S3FileSystem)))):
            raise Exception("Filesystem handle should be provided when nature is s3|gcs")

        if not isinstance(detail, bool):
            raise Exception("detail is a boolean field")

        path = self.paths[name]

        try:
            if ((resolve is not None) and
                (isinstance(resolve, dict)) and
                (len(resolve) > 0)):
                path = path % resolve
        except:
            raise Exception("Path could not be resolved")

        if fshandle is None:
            return os.listdir(path)

        results = fshandle.listdir(path, detail=detail)
        if not detail:
            results = [os.path.basename(r) for r in results]

        return results

def get_commands():

    return [
        {
            "valid": ['s3', 'local'],
            "name": "s3-to-local",
            "command": "aws s3 sync s3://%(source)s/%(subdir)s%(slash)s %(destination)s/%(subdir)s%(slash)s",
        },
        {
            "valid": ['local', 's3'],
            "name": "local-to-s3",
            "command": "aws s3 sync %(source)s/%(subdir)s%(slash)s s3://%(destination)s/%(subdir)s%(slash)s",
        },
        {
            "valid": ['s3', 's3'],
            "name": "s3-to-s3",
            "command": "aws s3 sync s3://%(source)s/%(subdir)s%(slash)s s3://%(destination)s/%(subdir)s%(slash)s",
        },
        {
            "valid": ['local', 'remote'],
            "name": "local-to-remote",
            "command" : "scp -C -r ubuntu@%(node)s:%(source)s %(destination)s/%(subdir)s"

        },
        {
            "valid": ['remote', 'local'],
            "name": "local-to-remote",
            "command" : "scp -C -r %(destination)s/%(subdir)s ubuntu@%(node)s:%(source)s"
        }
    ]



class DatasetRegistry(object):
    """
    Registry for datasets.

    This provides a search and resolution interface.

    There is a notion of a command - a scp/aws command line template.
    The registry allows enumeration of the commands and enables
    scripting.

    """
    def __init__(self, *args, **kwargs):
        """

        Args:
            commands (list): List of command templates. Optional. If not specified,
                             the system will use defaults.
            resolve (dict): Path resolution dictionary

        """
        self.commands = kwargs.get('commands', get_commands())
        self.resolve = kwargs.get('resolve',{})
        self.datasets = []
        self.validate()

    def validate(self):
        for c in self.commands:
            assert isinstance(c, dict)
            assert len(c) > 0
            assert 'valid' in c
            assert isinstance(c['valid'], list)
            assert len(c['valid']) == 2
            assert 'name' in c
            assert 'command' in c
            assert isinstance(c['command'], str)
            assert len(c['command']) > 0

    def get_command(self, name):
        """
        Get the command specification by specifying its name

        Args:
            name (str): Name of the command template
        """
        for c in self.commands:
            if c['name'] == name:
                return c
        raise Exception("Missing command: {}".format(name))

    def get_commands(self):
        """
        Get the command specifications as a list of dicts

        """
        for c in self.commands:
            if c['name'] == name:
                return c
        raise Exception("Missing command: {}".format(name))

    def add_datasets(self, items):
        """
        Add to the registry. Each element in the item list
        should be a Dataset (or subclass).

        Args:
           items (list): A list of datasets
        """
        if isinstance(items, list):
            for item in items:
                assert isinstance(item, Dataset)
                item.set_resolve(self.resolve)
            self.datasets.extend(items)
        else:
            assert isinstance(items, Dataset)
            items.set_resolve(self.resolve)
            self.datasets.append(items)

    def set_resolve(self, resolve):
        """
        Set the resolution parameters

        Args:
           resolve (dict): Resolution parameters

        """
        assert isinstance(resolve, dict)
        assert len(resolve) > 0
        self.resolve = resolve

    @property
    def params(self):
        return self.resolve

    def set_params(self, params):
        """
        Deprecated. Use set_resolve
        """
        return self.set_resolve(params)

    def list(self):
        """
        List all datasets in the registry
        """
        return [d.name for d in self.datasets]

    def find(self, name):
        """
        Find a dataset in the registry.

        Args:
            name (str): Name of the dataset
        """
        for d in self.datasets:
            if d.name == name:
                return d
        return None
