#import pyfuse3

# TODO: Implement File COW (except for append) (+ version / token updates caused by this)
# TODO: ? Stop using tokens for dirs, use hashed name + parent token
# TODO: ? Stop using tokens for files, use hashed name + short ?version-number + parent token instead
# TODO: Store version-numbers for files in parent-dir
# TODO: Chane Milestone-Format from [stone,stone,...] to [[dirStne,dirStone,...],[fileStone,,...]] and dont save type for the stones
# TODO: Decide how / when / from which class to push new blocks
# TODO: Unload 'overwritten' blobs
# TODO: Close blobs when they become unknown to the kernel or when we unmount (genesis only on unmount)
# TODO: When unmounting walk throught tree and seal all blobs

from iota import Iota, ProposedTransaction, Address, TryteString, Tag
from iota.crypto.addresses import AddressGenerator
from iota.crypto.types import Seed
from iota.codecs import TrytesDecodeError

from Crypto.Cipher import AES
from Crypto.Util.Padding import pad, unpad

import math
from pprint import pprint
import hashlib
import sys
import random
import time

import msgpack
import asyncio
import copy

import gzip
import secrets

CHUNKSIZE = 2187
SYNCWRITES = True

def log(txt):
    print("[-] "+str(txt))

def sendEmOff(bundles, api):
    print("SENDING (")
    pprint(bundles)
    print(")")
    return
    api.send_trytes(
        trytes=bundles
    )

class Atom():
    def __init__(self, milestone: bool, cont, name: str = None) -> None:
        self.milestone = milestone
        self.delta = not milestone
        self.name = name
        self.cont = cont

    def dump(self):
        if self.milestone:
            return msgpack.dumps([True, self.cont])
        else:
            return msgpack.dumps([False, self.name, self.cont])

class BlobChunk():
    def __init__(self, data: bytes = b'', sealed: bool = False) -> None:
        self.data = data
        self.sealed = sealed

    def getData(self) -> bytes:
        return self.data

    def append(self, data: bytes) -> None:
        if len(data)+len(self.data) > CHUNKSIZE:
            raise Exception("That to big!")
        self.data += data
        if len(self.data) == CHUNKSIZE:
            self.seal()

    def getBytesLeft(self) -> int:
        if self.sealed:
            return 0
        return CHUNKSIZE - len(self.data)

    def seal(self) -> None:
        self.sealed = True

    def isSealed(self) -> bool:
        return self.sealed

class TangleBlob():
    def __init__(self, token: bytes, iotaApi: Iota) -> None:
        self.token = token
        self.iotaApi = iotaApi
        self.preChunks = 0
        self.chunks = []
        m = hashlib.sha3_512()
        m.update(self.token)
        trSeed = TryteString.from_bytes(m.digest())[:81]
        self.adressGen = AddressGenerator(Seed(trSeed))
        self.fetched = False
        self.pushedNum = 0

    def _requireFetched(self):
        if not self.fetched:
            self.fetch()

    def _getKey(self, chunkNum: int) -> bytes:
        m = hashlib.sha3_384()
        m.update(self.token)
        m.update(chunkNum.to_bytes(8, "little")) # 64 bits should be enought...
        m.update(self.token)
        return m.digest()

    def _genBundle(self, data, addr) -> str:
        txMsg = TryteString.from_bytes(data)
        trans = ProposedTransaction(
            address = addr,
            value = 0,
            tag = Tag("IOTAFS"),
            message = txMsg
            )
        return self.iotaApi.prepare_transfer(
            transfers = [trans],
            inputs = [addr]
            )['trytes']

    def _dumpChunk(self, chunkNum: int) -> str:
        key = self._getKey(chunkNum + self.preChunks)
        data = self.chunks[chunkNum].getData()
        cipher = AES.new(key[16:][:16], AES.MODE_CBC, key[:16])
        ct_bytes = cipher.encrypt(pad(data, AES.block_size))
        addr = self.adressGen.get_addresses(start=chunkNum + self.preChunks, count=1)[0]
        return self._genBundle(ct_bytes, addr)

    def dumpAllSealed(self):
        bundles = []
        for i in range(len(self.chunks)-self.pushedNum):
            c = i + self.pushedNum
            chunk = self.chunks[c]
            if chunk.isSealed():
                bundles.append(self._dumpChunk(c))
                self.pushedNum+=1
        return bundles

    def sealAndDump(self):
        # When unmounting / closing / ...
        self.chunks[-1].seal()
        return self.dumpAllSealed()

    def append(self, data: bytes, newBlock: bool = False) -> None:
        self._requireFetched()
        if newBlock:
            self.chunks[-1].seal()
        elif len(self.chunks):
            bytesLeft = self.chunks[-1].getBytesLeft()
            if bytesLeft:
                leftChunk = data[:bytesLeft]
                data = data[bytesLeft:]
                self.chunks[-1].append(leftChunk)
        while len(data):
            chunk = data[:CHUNKSIZE]
            self.chunks.append(BlobChunk(chunk))
            data = data[CHUNKSIZE:]
        if SYNCWRITES:
            bundles = self.dumpAllSealed()
            if bundles:
                sendEmOff(bundles, self.iotaApi)

    def getChunkLen(self) -> int:
        return self.preChunks + len(self.chunks)

    def getSize(self) -> int:
        if len(self.chunks):
            return self.getChunkLen()*CHUNKSIZE - self.chunks[-1].getBytesLeft()
        return self.preChunks

    def read(self) -> bytes:
        data = b''
        for chunk in self.chunks:
            data += chunk.getData()
        return data

    def _dump(self) -> str:
        self.chunks[-1].seal()
        data = ""
        for c in range(len(self.chunks)-self.pushedNum):
            num = c + self.pushedNum
            data += self._dumpChunk(num) # num is without preChunks
        self.pushedNum = self.getChunkLen()
        return data

    def fetch(self) -> None:
        skipChunks = self.preChunks
        chunkNum = self.getChunkLen() + skipChunks
        while True:
            key = self._getKey(chunkNum)
            cipher = AES.new(key[16:][:16], AES.MODE_CBC, key[:16])
            addr = self.adressGen.get_addresses(start=chunkNum, count=1)[0]
            txHash = self.iotaApi.find_transactions(tags=[Tag("IOTAFS")], addresses=[addr])["hashes"]
            if len(txHash)==0:
                break
            bundles = self.api.get_bundles(txHash[0])["bundles"]
            for bundle in bundles:
                for tx in bundle.transactions:
                    # TODO: Can we just strip the 9s and call it a day?
                    tryteStr = TryteString(str(tx.signature_message_fragment).rstrip("9"))
                    try:
                        ct_bytes = tryteStr.as_bytes()
                    except TrytesDecodeError:
                        ct_bytes = (tryteStr+"9").as_bytes()
                    self.chunks.append(BlobChunk(unpad(cipher.decrypt(ct_bytes), AES.block_size), True))
            chunkNum += 1
        self.pushedNum = len(self.chunks)
        self.fetched = True
        self._afterFetch()

    def _afterFetch(self) -> None:
        return

    def genToken(self) -> bytes:
        return secrets.token_bytes(32)

    def sealLastChunk(self) -> None:
        self.chunks[-1].seal()

    def chunkLayout(self, width=50):
        fac = width / CHUNKSIZE
        lines = []
        for c,chunk in enumerate(self.chunks):
            bytesWritten = len(chunk.data)
            bytesEmpty = CHUNKSIZE - bytesWritten
            if chunk.isSealed():
                lines.append("["+"#"*int(bytesWritten*fac)+"="*int(bytesEmpty*fac)+"] (SEALED)")
            else:
                lines.append("["+"#"*int(bytesWritten*fac)+" "*int(bytesEmpty*fac)+"] ("+str(bytesWritten)+"/"+str(CHUNKSIZE)+")")
            if self.pushedNum < c+1:
                lines[-1]+=" {+}"
        return "\n".join(lines)

    def close(self):
        bundles = self.sealAndDump()
        if bundles:
            sendEmOff(bundles, self.iotaApi)

class TangleFileTreeElement(TangleBlob):
    def __init__(self, name: str, lastMilestoneIndex: int, parent, iotaApi: Iota) -> None:
        if isinstance(parent, bytes):
            self.token = hashlib.sha3_384(parent + name.encode()).digest()
        else:
            self.token = hashlib.sha3_384(parent.token + name.encode()).digest()
        super(TangleFileTreeElement, self).__init__(self.token, iotaApi)
        self.name = name
        self.inodes = {}
        self.parent = parent
        self.milestoneIndex = lastMilestoneIndex
        self.preChunks = self.milestoneIndex

    def _afterFetch(self) -> None:
        raw = self.read()
        if raw==b'':
            return
        unpacker = msgpack.Unpacker(raw=True)
        unpacker.feed(raw)
        for i, elem in enumerate(reversed(unpacker)):
            if elem[0]:
                # Is a milestone
                # TODO: Update our known milestoneIndex, if we find one
                #       might have to rewrite .fetch() and merge it here...
                self.milestoneIndex = self.getChunkLen()
                self._applyMilestone(elem[1])
                break
            else:
                if elem[1] in self.inodes: #name
                    self.inodes[elem[1]].applyDelta(elem[2])
                else:
                    # new file
                    self.inodes[elem[1]] = Inode(elem[1], self.iotaApi)
                    self.inodes[elem[1]].applyDelta(elem[2])

    def _getSkipChunks(self):
        return

    def _applyMilestone(self, milestone) -> None:
        self.inodes = {}
        for stone in milestone:
            atom = Atom()
            atom.load(stone)
            self.inodes[atom.name] = atom

    def _applyAtom(self, atom: Atom) -> None:
        if atom.name in self.inodes:
            self.inodes[atom.name].applyAtom(atom)
        else:
            cont = atom.cont
            type = ["dir","file"][cont[b't']]
            inode = Inode(atom.name, self.iotaApi, self, type)
            self.inodes[atom.name] = inode

    def _newAtom(self, atom: Atom) -> None:
        self.append(atom.dump())

    def getNameList(self):
        self._requireFetched()
        return list(self.inodes.keys())

    def _tree(self):
        dirs = {}
        files = []
        for inode in self.inodes:
            if self.inodes[inode].type=="file":
                files.append(inode)
            elif self.inodes[inode].type=="dir":
                dirs[inode] = self.inodes[inode].getRef()._tree()
            else:
                files.append("["+str(self.inodes[inode].type)+"]/"+inode)
        lines = ["{"+self.name+"}"]
        for i,d in enumerate(dirs):
            dir = dirs[d]
            if len(files)==0 and i==len(dirs)-1:
                lines.append(" └──"+dir[0]+"")
                for l in range(len(dir)-1):
                    lines.append("    "+dir[l+1])
            else:
                lines.append(" ├──"+dir[0]+"")
                for l in range(len(dir)-1):
                    lines.append(" │  "+dir[l+1])
        if len(files):
            for f in range(len(files)-1):
                lines.append(" ├──"+files[f])
            lines.append(" └──"+files[-1])
        return lines

    def tree(self):
        return "\n".join(self._tree())

    def getInode(self, name: str) -> Atom:
        self._requireFetched()
        return self.inodes[name]

    def mkdir(self, name: str) -> bool:
        self._requireFetched()
        if name in self.getNameList():
            return False
        inode = Inode(name, self.iotaApi, self, "dir")
        atom = inode.change(milestoneIndex=0)
        self._newAtom(atom)
        self.inodes[name] = inode
        return True

    def mkfile(self, name: str) -> bool:
        self._requireFetched()
        if name in self.getNameList():
            return False
        inode = Inode(name, self.iotaApi, self, "file")
        atom = inode.change(size=0, milestoneIndex=0, hash=b'NULL')
        self._newAtom(atom)
        self.inodes[name] = inode
        return True

    def _updateFileSize(self, name: str, size: int) -> None:
        self._requireFetched()
        self.inodes[name].size = size
        self._newAtom(self.inodes[name])

    def _updateFileToken(self, name: str, token: bytes, size: int) -> None:
        log("New FileToken for file '"+name+"' registered")
        self._requireFetched()
        self.inodes[name].token = token
        self.inodes[name].size = size
        self._newAtom(self.inodes[name])

    def performMilestone(self) -> None:
        if isinstance(self.parent, bytes):
            raise Exception("You cant create a milestone of the genesis block, you idiot!")
        stones = {}
        for a in self.inodes:
            stones[a] = self.inodes[a].toStone()
        self.atomStack = 0
        # TODO: Delimiter ?
        # TODO: compression ?
        milestoneAtom = Atom(True, stones)
        data = milestoneAtom.dump()
        self.milestoneIndex = self.getSize()
        if self.parent!=None:
            self.parent._updateChildMilestone(self.name, self.milestoneIndex)
        self.append(data, True)

    def _updateChildMilestone(self, name: str, milestoneIndex: int):
        if isinstance(self.parent, bytes):
            # We are the genesis-block
            self.append(milestoneIndex)
        else:
            atom = self.inodes[name].change(milestoneIndex = milestoneIndex)
            self._newAtom(atom)

class TangleFile():
    def __init__(self, name: str, parent: TangleFileTreeElement, iotaApi: Iota) -> None:
        self.api = iotaApi
        self.name = name
        self.parent = parent
        self.reflexiveAtom = parent.inodes[name]
        self.size = self.reflexiveAtom.size
        self.hash = self.reflexiveAtom.hash
        self.token = hashlib.sha3_384(b'f' + parent.getRef().token + self.hash).digest()
        self.blob = TangleBlob(self.token, iotaApi)

    def write(self, offset: int, data: bytes):
        if offset == self.size:
            self.append(data)
        else:
            self.token = self.genToken()
            oldData = self.blob.read()
            newData = oldData[:offset] + data + oldData[offset+len(data):]
            self.blob = TangleBlob(self.token)
            self.blob.append(newData)
            self.size = self.blob.getSize()
            self.parent._updateFileToken(self.name, self.token, self.size)

class Inode():
    def __init__(self, name: str, iotaApi: Iota, parent: TangleFileTreeElement = None, type: str = None) -> None:
        self.parent = parent
        self.name = name
        self.type = type
        self.ref = None
        self.iotaApi = iotaApi

    def setType(self, type: str) -> None:
        self.type = type

    def change(self, size: int=None, hash: int=None, milestoneIndex: int=None) -> Atom:
        delta = {}
        delta[b't'] = (self.type=="file")
        if size!=None:
            self.size = size
            delta[b's'] = size
        if hash!=None:
            self.hash = hash
            delta[b'h'] = hash
        if milestoneIndex!=None:
            self.milestoneIndex = milestoneIndex
            delta[b'm'] = milestoneIndex
        return Atom(False, delta, self.name)

    def applyAtom(self, atom: Atom) -> None:
        if atom.name != self.name:
            raise Exception("Cannot apply atom ment for a different inode (names differ)")
        if atom.milestone:
            stones = atom.cont
            if self.name in stones:
                self.applyAtom(Atom(False, stones[self.name], self.name))
        else:
            if (self.type=="file") != atom.cont[b't']:
                raise Exception("I am a "+self.type+"; this atom is for the other thing")
            delta = atom.cont
            if b's' in delta:
                self.size = delta[b's']
            if b'h' in delta:
                self.hash = delta[b'h']
            if b'm' in delta:
                self.milestoneIndex = delta[b'm']

    def toStone(self) -> None:
        if self.type=="file":
            return [self.size, self.hash]
        else:
            return [self.milestoneIndex]

    def getRef(self):
        if self.name == "*":
            return None
        if not self.ref:
            if self.type=="dir":
                self.ref = TangleFileTreeElement(self.name, self.milestoneIndex, self.parent, self.iotaApi)
            elif self.type=="file":
                self.ref = TangleFile(self.name, self, self.iotaApi)
            else:
                raise Exception("Cannot get reference of an inode of type "+self.type)
        return self.ref

class IotaFs():
    def __init__(self, token) -> None:
        self.api = Iota('https://nodes.thetangle.org:443', local_pow=True)
        # TODO Cache last known milestone-Index of genesis locally
        self.genesis = TangleFileTreeElement("*", 0, token, None, self.api)
        log("Fetching Genesis...")
        self.genesis.fetch()
        log("Retrieving reference to root")
        if self.genesis.getNameList()!=["/"]:
            if len(self.genesis.getNameList()):
                # theres another directory in our genesis chain... WTF?!
                raise Exception("Corrupted Genesis-Chain:"
                              + "Unknown records for no root-directory in Genesis Chain: "+str(self.genesis.getNameList()))
            else:
                # we dont have a root yet, lets create one...
                log("Unable to reference to root: Creating new root")
                self.genesis.mkdir("/")
        log("Successfully Mounted!")


    def createNewFile(self, name) -> None:
        pass

api = Iota('https://nodes.thetangle.org:443', local_pow=True)
token = b'testToken'
genesis = TangleFileTreeElement("*", 0, token, api)
genesis.mkdir("/")
root = genesis.inodes["/"].getRef()

root.mkdir("dir1")
root.mkdir("dir2")
root.mkdir("dir3")
root.mkdir("dir4")
root.mkfile("file.txt")
d2 = root.inodes["dir2"].getRef()
d2.mkfile("a.txt")
d3 = root.inodes["dir3"].getRef()
d3.mkfile("b.txt")
d3.mkfile("c.txt")
d3.mkfile("d.txt")
d3.mkdir("subDir")
d4 = root.inodes["dir4"].getRef()
d4.mkdir("sub")
sub = d4.inodes["sub"].getRef()
sub.mkdir("subsub")

print(genesis.tree())