#import pyfuse3

# TODO: Milestones always written into new block
# 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: Move iota-push-code from TangleBlob to chunk? and perform when sealing?
# 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

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

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 Inode():
    def __init__(self, name: str, iotaApi: Iota, type: str = None) -> None:
        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 = {}
        if size!=None:
            self.size = size
            delta[b's'] = size
        if size!=None:
            self.size = 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:
            delta = atom.cont
            if b's' in delta:
                self.size = delta[b's']
            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.elem = TangleFileTreeElement(self.name, self.milestoneIndex, self, self.iotaApi)
            elif self.type=="file":
                self.ref = TangleFile(self.name, self, self.api)
            else:
                raise Exception("Cannot get reference of an inode of type "+self.type)
        return self.ref

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 append(self, data: bytes, newBlock: bool = False) -> None:
        self._requireFetched()
        if len(self.chunks) and not newBlock:
            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:]

    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()

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.getRef().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:
        ## TODO:
        pass

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

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

    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, "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
        atom = Atom()
        atom.type="file"
        atom.name = name
        atom.token = self.genToken()
        atom.size = 0
        self._newAtom(atom)
        self._applyAtom(atom)
        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:
        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)
        # inform parent about milestone (when merged)

    def _updateChildMilestone(self, name: str, milestoneIndex: int):
        self.inodes[name].milestoneIndex = milestoneIndex
        self._newAtom(self.inodes[name])

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 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

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