import os
import uuid
import csv
from ftplib import FTP
import yt.wrapper as yt_wrapper
from datacloud.dev_utils.logging.logger import get_basic_logger
from datacloud.dev_utils.s3.lib import s3_utils
from datacloud.dev_utils.yt import yt_utils

logger = get_basic_logger(__name__)


class BaseStorageException(Exception):
    def __init__(self, *args, **kwargs):
        super(BaseStorageException, self).__init__(*args, **kwargs)


class StorageItemNotExists(BaseStorageException):
    def __init__(self, *args, **kwargs):
        super(StorageItemNotExists, self).__init__(*args, **kwargs)


class GeneralStorage(object):
    def retrieve(self, storage_path, local_path):
        raise NotImplementedError

    def store(self, storage_path, local_path):
        raise NotImplementedError

    def list(self, folder):
        raise NotImplementedError

    def exists(self, path):
        raise NotImplementedError

    def mkdir(self, storage_path):
        raise NotImplementedError


class FtpStorage(GeneralStorage):
    class Context(object):
        def __init__(self, host, user, passwd):
            self.ftp = FTP(host, user, passwd)

        def __enter__(self):
            return self.ftp

        def __exit__(self, *exc):
            self.ftp.close()

    def context(self):
        return self.Context(self._host, self._user, self._passwd)

    def __init__(self, host, user, passwd):
        super(GeneralStorage, self).__init__()
        self._host = host
        self._user = user
        self._passwd = passwd

    def retrieve(self, storage_path, local_path):
        with self.context() as ftp, open(local_path, 'wb+') as fh:
            ftp.retrbinary('RETR {}'.format(storage_path), fh.write)

    def store(self, storage_path, local_path):
        with self.context() as ftp, open(local_path, 'rb') as fh:
            ftp.storbinary('STOR {}'.format(storage_path), fh)

    def list(self, folder):
        with self.context() as ftp:
            for item in ftp.nlst(folder):
                yield item

    def exists(self, path):
        parent = '/'.join(path.split('/')[:-1])
        item = parent.split('/')[-1]
        return item in self.list(parent)

    def mkdir(self, storage_path):
        with self.context() as ftp:
            ftp.mkd(storage_path)


class YtTableStorage(GeneralStorage):
    def __init__(self, yt_client, schema=None):
        super(GeneralStorage, self).__init__()
        self.yt_client = yt_client
        self.schema = schema

    def retrieve(self, storage_path, local_path):
        data = []
        for row in self.yt_client.read_table(storage_path):
            data.append(row)
        assert data, 'data is empty, check {}'.format(storage_path)
        columns = data[0].keys()

        with open(local_path, 'w+') as fh:
            writer = csv.DictWriter(fh, fieldnames=columns)
            writer.writeheader()
            for row in data:
                writer.writerow(row)

    def store(self, storage_path, local_path):
        with open(local_path, 'r') as fh:
            data = [row for row in csv.DictReader(fh)]
        table = yt_wrapper.TablePath(storage_path, schema=self.schema)
        if yt_utils.check_table_exists(storage_path, self.yt_client):
            assert 'Table {} alreay exists'.format(storage_path)
        self.yt_client.write_table(table, data)

    def list(self, folder):
        return self.yt_client.list(folder)

    def exists(self, path):
        return self.yt_client.exists(path)

    def mkdir(self, storage_path):
        yt_utils.create_folders([storage_path], self.yt_client)


class S3Storage(object):
    def __init__(self, client, bucket):
        self._client = client
        self._bucket = bucket

    def retrieve(self, storage_path, local_path):
        s3_utils.download(self._client, self._bucket, storage_path, local_path)

    def store(self, storage_path, local_path):
        s3_utils.upload(self._client, self._bucket, local_path, storage_path)

    def list(self, folder):
        if not self.exists(folder):
            raise StorageItemNotExists(folder)
        for item in s3_utils.list(self._client, self._bucket,
                                  folder, absolute=False):
            yield item

    def exists(self, path):
        return s3_utils.exists(self._client, self._bucket, path)

    def mkdir(self, storage_path):
        s3_utils.mkdir(self._client, self._bucket, storage_path)


def retrieve_tmp_file(storage, storage_path, local_path=None):
    if local_path is None:
        local_path = os.path.join('/tmp', uuid.uuid4().hex)
    if os.path.exists(local_path):
        logger.info('Old file with exists, remove {}'.format(local_path))
        os.remove(local_path)
    storage.retrieve(storage_path, local_path)
    return local_path


def move(from_storage, from_storage_path, to_storage, to_storage_path):
    tmp_file = retrieve_tmp_file(from_storage, from_storage_path)
    to_storage.store(to_storage_path, tmp_file)
    if os.path.exists(tmp_file):
        os.remove(tmp_file)
