import struct
import os

from io import BytesIO

from . import NetworkError, FileServerException
from .zmq import ZMQReq
from pymoku.progress import progress


FS_CHUNK_SIZE = 0x40000

P_TYPE = 0x49

ACT_GET = 0x01
ACT_SET = 0x02
ACT_CSUM = 0x03
ACT_SIZE = 0x04
ACT_LIST = 0x05
ACT_FREE = 0x06
ACT_REMOVE = 0x07
ACT_RENAME = 0x08
ACT_RENAME_STATUS = 0x09

ERR_OK = 0x00
ERR_INVAL = 0x21
ERR_NOT_FOUND = 0x22
ERR_SPACE = 0x23
ERR_MOUNT = 0x24
ERR_OPERATION = 0x25
ERR_BUSY = 0x26
ERR_RONLY = 0x27
ERR_UNKNOWN = 0xFF

CSUM_FLAGS = {
    'none': 0x00,
    'crc32': 0x01,
    'sha256': 0x02,
    'sha256_short': 0x03,
}

msgs = {
    ERR_INVAL: "Invalid fileserver request parameters.",
    ERR_NOT_FOUND: "Could not find directory or file.",
    ERR_SPACE: "Insufficient space to perform action.",
    ERR_MOUNT: "Mount point has not been mounted.",
    ERR_OPERATION: "Unknown fileserver action requested.",
    ERR_BUSY: "Fileserver busy",
    ERR_RONLY: "Requested mount point was Read-Only.",
    ERR_UNKNOWN: "Unknown fileserver action requested."
}

PORT = 27184


class FileTransfer(object):

    def __init__(self, skt):
        self.skt = skt

    def _transfer(self, data, retries=1):

        for i in range(retries):
            try:
                rep = self.skt.transfer(data)
                ptype, status, act = struct.unpack("BBB", rep[:3])
                if data[0] != rep[0] or (status <= 5 and status > 0):
                    raise NetworkError("Status: 0x{:02X}, {:s}".format(
                        status, rep[2:].decode('utf-8')))
                break
            except NetworkError:
                if i == retries - 1:
                    raise

        if not status:
            return rep

        try:
            ex = FileServerException(msgs[status])
        except KeyError:
            ex = FileServerException(
                "Received invalid status ID: {:d}".format(status))
        raise ex

    def get_chunk(self, mp, fname, into=None, length=0, offset=0):
        path = mp + ":" + fname

        pkt = struct.pack('<BBB', P_TYPE, ACT_GET, len(path))
        pkt += bytes(path, 'utf-8')
        pkt += struct.pack("<QQ", offset, length)

        resp = self._transfer(pkt, retries=5)

        chunk = memoryview(resp)[11:]

        if into is None:
            return chunk

        try:
            into.write(chunk)
        except AttributeError:
            BytesIO(chunk).readinto(into)

        return len(chunk)

    def send_chunk(self, mp, fname, data, offset=0, length=None,
                   create=True, truncate=False, finalize=True):

        if length is None:
            length = len(data)

        path = mp + ":" + fname
        flags = int(create) << 0
        flags |= int(truncate) << 1
        flags |= int(finalize) << 2

        pkt = struct.pack('<BBBB', P_TYPE, ACT_SET, flags, len(path))
        pkt += bytes(path, 'utf-8')
        pkt += struct.pack("<QQ", offset, length)
        pkt += data.read(length)

        # TODO no-copy sends
        self._transfer(pkt, retries=5)

    def send(self, mp, fname, data):
        try:
            length = len(data)
            buf = BytesIO(data)
        except TypeError:
            length = os.stat(data.fileno()).st_size
            buf = data

        chunks = length // FS_CHUNK_SIZE

        with progress() as prog:
            prog.set_message(f"Uploading {fname}")
            prog.set_progress(0.0)
            for i in range(chunks):
                o = i * FS_CHUNK_SIZE
                self.send_chunk(mp, fname, buf,
                                offset=o, length=min(FS_CHUNK_SIZE, length - o),
                                create=False, truncate=False, finalize=False)
                prog.set_progress(i / chunks)

            o = chunks * FS_CHUNK_SIZE
            self.send_chunk(mp, fname, buf,
                            offset=o, length=length - o,
                            create=False, truncate=True, finalize=True)

    def get(self, mp, fname, into=None):
        size = self.size(mp, fname)

        if into is None:
            buf = bytearray(size)
        else:
            buf = into

        recvd = 0
        for i in range(size // FS_CHUNK_SIZE + 1):
            o = i * FS_CHUNK_SIZE
            recvd += self.get_chunk(mp, fname, into=buf,
                                    length=FS_CHUNK_SIZE, offset=o)

        if into is not None:
            return recvd
        return buf

    def remove(self, mp, fname):
        fname = mp + ":" + fname

        data = struct.pack('<BBB', P_TYPE, ACT_REMOVE, len(fname))
        data += fname.encode('utf8')

        self._transfer(data)

    def csum(self, mp, fname, csum='sha256'):
        # sha256 = hashlib.sha256().digest().hex()
        # csum = hex(zlib.crc32), reversed byte order
        fname = mp + ":" + fname
        flags = CSUM_FLAGS[csum]

        data = struct.pack('<BBBB', P_TYPE, ACT_CSUM, flags, len(fname))
        data += fname.encode('utf8')

        rep = BytesIO(self._transfer(data))
        length = struct.unpack('BBBB', rep.read(4))[3]

        return rep.read(length).decode('utf8')

    def size(self, mp, fname):
        fname = mp + ":" + fname

        data = struct.pack('<BBB', P_TYPE, ACT_SIZE, len(fname))
        data += fname.encode('utf-8')

        rep = BytesIO(self._transfer(data))

        ptype, status, action, length = struct.unpack("<BBBQ", rep.read(11))

        return length

    def list(self, mp, csum='none'):
        flags = CSUM_FLAGS[csum]

        data = struct.pack('<BB', P_TYPE, ACT_LIST)
        data += bytes(mp, 'utf8')
        data += struct.pack('<B', flags)

        reply = BytesIO(self._transfer(data))

        ptype, status, action, length = struct.unpack("<BBBH", reply.read(5))

        names = {}
        for i in range(length):
            if csum in ['sha256', 'sha256_short']:
                chk = reply.read(64)
            elif csum == 'crc32':
                chk = reply.read(8)
            else:
                chk = b''

            bl, fl = struct.unpack("<QB", reply.read(9))

            names[reply.read(fl).decode()] = (chk, bl)

        return names

    def free(self, mp):
        data = struct.pack('<BB', P_TYPE, ACT_FREE) + mp.encode('utf8')
        rep = self._transfer(data)

        _, _, _, t, f = struct.unpack("<BBBQQ", rep)

        return t, f

    def rename(self, smp, sname, dmp, dname, move=False):
        sname = smp + ":" + sname
        dname = dmp + ":" + dname
        pkt = bytearray([len(sname)])
        pkt += sname.encode('utf8')
        pkt += bytearray([len(dname)])
        pkt += dname.encode('utf8')

        flags = 1 if move else 0
        pkt += bytearray([flags])

        rep = self._transfer(8, pkt)

        return rep

    def rename_status(self):
        try:
            dat = self._transfer(9, b'')
            stat = ERR_OK
        except FileServerException as e:
            dat = e.dat
            stat = ERR_BUSY

        size, pc = struct.unpack("<QB", dat)

        return stat, size, pc

    def close(self):
        self.skt.close()

    def __enter__(self):
        return self

    def __exit__(self, exc_type, exc_value, traceback):
        self.close()


class FileTransfer511(FileTransfer):

    def _transfer(self, data, retries=1):
        for i in range(retries):
            try:
                rep = ZMQReq.transfer(self.skt, data)
                break
            except NetworkError:
                if i == retries - 1:
                    raise

        ptype, length, act, status = struct.unpack("<BQBB", rep[:11])

        if not status:
            return rep

        try:
            ex = FileServerException(msgs[status])
        except KeyError:
            ex = FileServerException(
                "Received invalid status ID: {:d}".format(status))

        raise ex

    def _finalise(self, path, length):
        pkt = struct.pack('<BQBB', P_TYPE, 0x00000000, 0x07, len(path))
        pkt += bytes(path, 'utf-8')
        pkt += struct.pack("<Q", length)
        self._transfer(pkt)

    def send_chunk(self, mp, fname, data, offset=0, length=None,
                   create=True, truncate=False, finalize=True):

        if length is None:
            length = len(data)

        path = mp + ":" + fname

        pkt = struct.pack('<BQBB', P_TYPE, 0x00000000, ACT_SET, len(path))
        pkt += bytes(path, 'utf-8')
        pkt += struct.pack("<QQ", offset, length)
        pkt += data.read(length)

        # TODO no-copy sends
        self._transfer(pkt, retries=5)

        if finalize:
            self._finalise(path, offset + length)


class FileTransfer433(FileTransfer511):

    def connect(self, ip_addr):
        return ZMQReq.connect(self, ip_addr)

    def _transfer(self, data, retries=1):
        for i in range(retries):
            try:
                rep = ZMQReq.transfer(self.skt, data)
                break
            except NetworkError:
                if i == retries - 1:
                    raise

        ptype, length, act, status = struct.unpack("<BQBB", rep[:11])

        if not status:
            return rep

        try:
            ex = FileServerException(msgs[status])
        except KeyError:
            ex = FileServerException(
                "Received invalid status ID: {:d}".format(status))

        raise ex
