from typing import Mapping, MutableMapping, Sequence, Iterable, List, Set
import pyfuse3

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

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

class Atom():
    def __init__(self) -> None:
        self.type = "null"

    def load(self, data) -> None:
        if data[0]==True:
            self.type = "milestone"
            self.milestone = data[1:]
        else:
            if data[1]==True:
                self.type = "dir"
            else:
                self.type = "file"
                self.size = data[5]
            self.name = data[2]
            self.token = data[3]
            self.milestoneIndex = data[4]

    def dump(self) -> bytes:
        # TODO: Delimiter?
        # TODO: compression?
        return msgpack.dump(self._dumpAsArray())

    def _dumpAsArray(self):
        if self.type=="milestone":
            return [True]+self.milestone
        else:
            return [False, self.type, self.name, self.token, self.milestoneIndex, self.size]

class BlobChunk():
    def __init__(self, data: bytes = b'', sealed: bool = False) -> None:
        self.data = data
        self.sealded = 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

    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)
        self.adressGen = AddressGenerator(Seed(m.digest()))
        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, security_level=2) -> 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, security_level=security_level)['addresses'][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, 2)
        return data

    def fetch(self, security_level=2, skipChunks: int = 0) -> None:
        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, security_level=security_level)['addresses'][0]
            txHash = self.api.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() -> None:
            return

class TangleFileTreeElement(TangleBlob):
    def __init__(self, token: bytes, iotaApi: Iota) -> None:
        super(TangleFileTreeElement, self).__init__(token)
        self.inodes = {}
        self.atomStack = -1

    def _afterFetch(self) -> None:
        data = msgpack.load(self.read())
        newAtoms = []
        for i, elem in enumerate(reversed(data)):
            atom = Atom()
            atom.load(elem)
            if atom.type == "milestone":
                self._applyMilestone(atom.milestone)
                break
            newAtoms.append(atom)
        self.atomStack = len(newAtoms)
        for atom in reversed(newAtoms):
            self._applyAtom(atom)

    def _applyMilestone(self, milestone) -> None:
        self.atomStack = 0
        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.inode:
            self.atomStack += 1
        self.inodes[atom.name] = atom
        if atom.type=="dir":
            self.inodes[atom.name].elem = TangleFileTreeElement(atom.token, self.iotaApi)
        elif atom.type=="file":
            self.inodes[atom.name].elem = TangleFile(atom.token, self.iotaApi)
        else:
            raise Exception("How did such an atom get here?")

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

    def getNameList(self) -> List(str):
        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
        atom = Atom()
        atom.type="dir"
        atom.name = name
        atom.token = "GENTOKENHERE" # TODO: <
        atom.milestoneIndex = 0
        self._newAtom(atom)
        self._applyAtom(atom)
        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 = "GENTOKENHERE" # TODO: <
        atom.milestoneIndex = 0
        self._newAtom(atom)
        self._applyAtom(atom)
        return True

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

    def _updateFileToken(self, name: str, token: bytes, size: int):
        self._requireFetched()
        self.inodes[name].token = token
        self.inodes[name].size = size
        self._newAtom(self.inodes[name])

    def performMilestone(self):
        stones = []
        for atom in self.inodes:
            stones.append(atom._dumpAsArray())
        self.atomStack = 0
        # TODO: Delimiter ?
        # TODO: compression ?
        data = msgpack.dump(stones)
        self.append(data)

class TangleFile():
    def __init__(self, name: str, parent: TangleFileTreeElement, iotaApi: Iota) -> None:
        self.name = name
        self.parent = parent
        self.reflexiveAtom = parent.inodes[name]
        self.size = self.reflexiveAtom.size
        self.token = self.reflexiveAtom.token
        super(TangleFile, self).__init__(self.token, iotaApi)

    def write(self, offset: int, data: bytes):
        if offset == self.size:
            self.append(data)
        else:
            pass


class IotaFs():
    def __init__(self, token) -> None:
        self.api = Iota('https://nodes.thetangle.org:443', local_pow=True)
        self.genesis = TangleFileTreeElement(token, 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