"""
novelWriter – Project Tree Class
================================
Data class for the project's tree of project items

File History:
Created: 2020-05-07 [0.4.5]

This file is a part of novelWriter
Copyright 2018–2022, Veronica Berglyd Olsen

This program is free software: you can redistribute it and/or modify
it under the terms of the GNU General Public License as published by
the Free Software Foundation, either version 3 of the License, or
(at your option) any later version.

This program is distributed in the hope that it will be useful, but
WITHOUT ANY WARRANTY; without even the implied warranty of
MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the GNU
General Public License for more details.

You should have received a copy of the GNU General Public License
along with this program. If not, see <https://www.gnu.org/licenses/>.
"""

import os
import logging

from time import time
from lxml import etree
from hashlib import sha256

from novelwriter.enum import nwItemType, nwItemClass, nwItemLayout
from novelwriter.error import logException
from novelwriter.common import checkHandle
from novelwriter.constants import nwConst, nwFiles
from novelwriter.core.item import NWItem

logger = logging.getLogger(__name__)


class NWTree():

    def __init__(self, theProject):

        self.theProject = theProject

        self._projTree    = {}     # Holds all the items of the project
        self._treeOrder   = []     # The order of the tree items on the tree view
        self._treeRoots   = []     # The root items of the tree
        self._trashRoot   = None   # The handle of the trash root folder
        self._archRoot    = None   # The handle of the archive root folder
        self._theIndex    = 0      # The current iterator index
        self._treeChanged = False  # True if tree structure has changed

        self._handleSeed  = None   # Used for generating handles for testing
        self._handleCount = 0      # A counter that is added to the handle generator

        return

    ##
    #  Class Methods
    ##

    def clear(self):
        """Clear the item tree entirely.
        """
        self._projTree  = {}
        self._treeOrder = []
        self._treeRoots = []
        self._trashRoot = None
        self._archRoot  = None
        self._theIndex  = 0
        self._treeChanged = False
        return

    def handles(self):
        """Returns a copy of the list of all the active handles.
        """
        return self._treeOrder.copy()

    def append(self, tHandle, pHandle, nwItem):
        """Add a new item to the end of the tree.
        """
        tHandle = checkHandle(tHandle, None, True)
        pHandle = checkHandle(pHandle, None, True)
        if tHandle is None:
            tHandle = self._makeHandle()

        if tHandle in self._projTree:
            logger.warning("Duplicate handle '%s' detected, skipping", tHandle)
            return False

        logger.verbose("Adding item '%s' with parent '%s'", str(tHandle), str(pHandle))

        nwItem.setHandle(tHandle)
        nwItem.setParent(pHandle)

        if nwItem.itemType == nwItemType.ROOT:
            logger.verbose("Item '%s' is a root item", str(tHandle))
            self._treeRoots.append(tHandle)
            if nwItem.itemClass == nwItemClass.ARCHIVE:
                logger.verbose("Item '%s' is the archive folder", str(tHandle))
                self._archRoot = tHandle

        if nwItem.itemType == nwItemType.TRASH:
            if self._trashRoot is None:
                logger.verbose("Item '%s' is the trash folder", str(tHandle))
                self._trashRoot = tHandle
            else:
                logger.error("Only one trash folder allowed")
                return False

        self._projTree[tHandle] = nwItem
        self._treeOrder.append(tHandle)
        self._setTreeChanged(True)

        return True

    def packXML(self, xParent):
        """Pack the content of the tree into the provided XML object. In
        the order defined by the _treeOrder list.
        """
        xContent = etree.SubElement(xParent, "content", attrib={
            "count": str(len(self._treeOrder))}
        )
        for tHandle in self._treeOrder:
            tItem = self.__getitem__(tHandle)
            tItem.packXML(xContent)
        return

    def unpackXML(self, xContent):
        """Iterate through all items of a content XML object and add
        them to the project tree.
        """
        if xContent.tag != "content":
            logger.error("XML entry is not a NWTree")
            return False

        self.clear()
        for xItem in xContent:
            nwItem = NWItem(self.theProject)
            if nwItem.unpackXML(xItem):
                self.append(nwItem.itemHandle, nwItem.itemParent, nwItem)
                nwItem.saveInitialCount()

        return True

    def writeToCFile(self):
        """Write the convenience table of contents file in the root of
        the project directory.
        """
        tocList = []
        tocLen = 0
        for tHandle in self._treeOrder:
            tItem = self.__getitem__(tHandle)
            if tItem is None:
                continue
            tFile = tHandle+".nwd"
            if os.path.isfile(os.path.join(self.theProject.projContent, tFile)):
                tocLine = "{0:<25s}  {1:<9s}  {2:<8s}  {3:s}".format(
                    os.path.join("content", tFile),
                    tItem.itemClass.name,
                    tItem.itemLayout.name,
                    tItem.itemName,
                )
                tocList.append(tocLine)
                tocLen = max(tocLen, len(tocLine))

        try:
            # Dump the text
            tocText = os.path.join(self.theProject.projPath, nwFiles.TOC_TXT)
            with open(tocText, mode="w", encoding="utf-8") as outFile:
                outFile.write("\n")
                outFile.write("Table of Contents\n")
                outFile.write("=================\n")
                outFile.write("\n")
                outFile.write("{0:<25s}  {1:<9s}  {2:<8s}  {3:s}\n".format(
                    "File Name", "Class", "Layout", "Document Label"
                ))
                outFile.write("-"*max(tocLen, 62) + "\n")
                outFile.write("\n".join(tocList))
                outFile.write("\n")

        except Exception:
            logger.error("Could not write ToC file")
            logException()
            return False

        return True

    def sumWords(self):
        """Loop over all entries and add up the word counts.
        """
        noteWords = 0
        novelWords = 0
        for tHandle in self._treeOrder:
            tItem = self.__getitem__(tHandle)
            if tItem is None:
                continue
            if tItem.itemLayout == nwItemLayout.NO_LAYOUT:
                pass
            elif tItem.itemLayout == nwItemLayout.NOTE:
                noteWords += tItem.wordCount
            else:
                novelWords += tItem.wordCount
        return novelWords, noteWords

    ##
    #  Tree Structure Methods
    ##

    def checkType(self, tHandle, itemType):
        """Return true of item exists and is of the specified item type.
        """
        tItem = self.__getitem__(tHandle)
        if not tItem:
            return False
        return tItem.itemType == itemType

    def trashRoot(self):
        """Returns the handle of the trash folder, or None if there
        isn't one.
        """
        if self._trashRoot:
            return self._trashRoot
        return None

    def isTrashRoot(self, tHandle):
        """Check if a handle is the trash folder.
        """
        if self._trashRoot is None:
            return False
        return tHandle == self._trashRoot

    def archiveRoot(self):
        """Returns the handle of the archive folder, or None if there
        isn't one.
        """
        if self._archRoot:
            return self._archRoot
        return None

    def findRoot(self, theClass):
        """Find the root item for a given class.
        Note: This returns the first item for class CUSTOM.
        """
        for aRoot in self._treeRoots:
            tItem = self.__getitem__(aRoot)
            if tItem is None:
                continue
            if theClass == tItem.itemClass:
                return tItem.itemHandle
        return None

    def checkRootUnique(self, theClass):
        """Checks if there already is a root entry of class 'theClass'
        in the root of the project tree. CUSTOM class is skipped as it
        is not required to be unique.
        """
        if theClass == nwItemClass.CUSTOM:
            return True
        for aRoot in self._treeRoots:
            tItem = self.__getitem__(aRoot)
            if tItem is None:
                continue
            if theClass == tItem.itemClass:
                return False
        return True

    def getRootItem(self, tHandle):
        """Iterate upwards in the tree until we find the item with
        parent None, the root item. We do this with a for loop with a
        maximum depth to make infinite loops impossible.
        """
        tItem = self.__getitem__(tHandle)
        if tItem is not None:
            for i in range(nwConst.MAX_DEPTH + 1):
                if tItem.itemParent is None:
                    return tItem
                else:
                    tHandle = tItem.itemParent
                    tItem = self.__getitem__(tHandle)
        return None

    def getItemPath(self, tHandle):
        """Iterate upwards in the tree until we find the item with
        parent None, the root item, and return the list of handles.
        We do this with a for loop with a maximum depth to make
        infinite loops impossible.
        """
        tTree = []
        tItem = self.__getitem__(tHandle)
        if tItem is not None:
            tTree.append(tHandle)
            for _ in range(nwConst.MAX_DEPTH + 1):
                if tItem.itemParent is None:
                    return tTree
                else:
                    tHandle = tItem.itemParent
                    tItem = self.__getitem__(tHandle)
                    if tItem is None:
                        return tTree
                    else:
                        tTree.append(tHandle)
        return tTree

    ##
    #  Setters
    ##

    def setOrder(self, newOrder):
        """Reorders the tree based on a list of items.
        """
        tmpOrder = []

        # Add all known elements to a new temp list
        for tHandle in newOrder:
            if tHandle in self._projTree:
                tmpOrder.append(tHandle)
            else:
                logger.error("Handle '%s' in new tree order is not in project tree", tHandle)

        # Do a reverse lookup to check for items that will be lost
        # This is mainly for debugging purposes
        for tHandle in self._treeOrder:
            if tHandle not in tmpOrder:
                logger.warning("Handle '%s' in old tree order is not in new tree order", tHandle)

        # Save the temp list
        self._treeOrder = tmpOrder
        self._setTreeChanged(True)
        logger.verbose("Project tree order updated")

        return

    def setSeed(self, theSeed):
        """Used for debugging!
        Sets a seed for generating handles so that they always come out
        in a predictable order.
        """
        self._handleSeed = theSeed
        return

    def setFileItemLayout(self, tHandle, itemLayout):
        """Set the nwItemLayout for a specific file.
        """
        tItem = self.__getitem__(tHandle)
        if tItem is None:
            return False
        if tItem.itemType != nwItemType.FILE:
            logger.error("Item '%s' is not a file", tHandle)
            return False
        if not isinstance(itemLayout, nwItemLayout):
            return False

        tItem.setLayout(itemLayout)

        return True

    ##
    #  Getters
    ##

    def countTypes(self):
        """Count the number of files, folders and roots in the project.
        """
        nRoot = 0
        nFolder = 0
        nFile = 0

        for tHandle in self._treeOrder:
            tItem = self.__getitem__(tHandle)
            if tItem is None:
                continue
            elif tItem.itemType == nwItemType.ROOT:
                nRoot += 1
            elif tItem.itemType == nwItemType.FOLDER:
                nFolder += 1
            elif tItem.itemType == nwItemType.FILE:
                nFile += 1

        return nRoot, nFolder, nFile

    ##
    #  Meta Methods
    ##

    def __len__(self):
        """Return the length counter. Does not check that it is correct!
        """
        return len(self._treeOrder)

    def __bool__(self):
        """Returns True if the tree has any entries.
        """
        return len(self._treeOrder) > 0

    ##
    #  Item Access Methods
    ##

    def __getitem__(self, tHandle):
        """Return a project item based on its handle. Returns None if
        the handle doesn't exist in the project.
        """
        if tHandle in self._projTree:
            return self._projTree[tHandle]
        logger.error("No tree item with handle '%s'", str(tHandle))
        return None

    def __delitem__(self, tHandle):
        """Remove an item from the internal lists and dictionaries.
        """
        if tHandle in self._treeOrder and tHandle in self._projTree:
            self._treeOrder.remove(tHandle)
            del self._projTree[tHandle]
        else:
            logger.warning("Failed to delete item '%s': item not found", tHandle)
            return

        if tHandle in self._treeRoots:
            self._treeRoots.remove(tHandle)
        if tHandle == self._trashRoot:
            self._trashRoot = None
        if tHandle == self._archRoot:
            self._archRoot = None

        self._setTreeChanged(True)

        return

    def __contains__(self, tHandle):
        """Checks if a handle exists in the tree.
        """
        return tHandle in self._treeOrder

    ##
    #  Iterator Methods
    ##

    def __iter__(self):
        """Initiates the iterator.
        """
        self._theIndex = 0
        return self

    def __next__(self):
        """Returns the item from the next entry in the _treeOrder list.
        """
        if self._theIndex < len(self._treeOrder):
            theItem = self.__getitem__(self._treeOrder[self._theIndex])
            self._theIndex += 1
            return theItem
        else:
            raise StopIteration

    ##
    #  Internal Functions
    ##

    def _setTreeChanged(self, theState):
        """Set the changed flag to theState, and if being set to True,
        propagate that state change to the parent NWProject class.
        """
        self._treeChanged = theState
        if theState:
            self.theProject.setProjectChanged(True)
        return

    def _makeHandle(self, addSeed=""):
        """Generate a unique item handle. In the event that the key
        already exists, salt the seed and generate a new handle.
        A key collision is very unlikely to be caused by the truncation
        of the sha256 hash to 13 characters. Assuming it is near-random,
        it will on average happen every 4.5^15 times. However, the clock
        seed is likely to occasionally generate a collision if the
        handle requests come faster than the clock resolution.
        """
        if self._handleSeed is None:
            newSeed = "%s_%d_%s" % (str(time()), self._handleCount, addSeed)
            self._handleCount += 1
        else:
            # This is used for debugging
            newSeed = str(self._handleSeed)
            self._handleSeed += 1

        logger.verbose("Generating handle with seed '%s'", newSeed)
        itemHandle = sha256(newSeed.encode()).hexdigest()[0:13]
        if itemHandle in self._projTree:
            logger.warning("Duplicate handle encountered! Retrying ...")
            itemHandle = self._makeHandle(addSeed+"!")

        return itemHandle

# END Class NWTree
