XNbt.java

/*
Copyright (c) 2021, David Quintana <gigaherz@gmail.com>
All rights reserved.

Redistribution and use in source and binary forms, with or without
modification, are permitted provided that the following conditions are met:
    * Redistributions of source code must retain the above copyright
      notice, this list of conditions and the following disclaimer.
    * Redistributions in binary form must reproduce the above copyright
      notice, this list of conditions and the following disclaimer in the
      documentation and/or other materials provided with the distribution.
    * Neither the name of the author nor the
      names of the contributors may be used to endorse or promote products
      derived from this software without specific prior written permission.

THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" AND
ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED
WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
DISCLAIMED. IN NO EVENT SHALL THE AUTHOR BE LIABLE FOR ANY
DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES
(INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES;
LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND
ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
(INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS
SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
 */
package dev.gigaherz.xnbt;

import net.minecraft.nbt.*;
import org.w3c.dom.Document;
import org.w3c.dom.Element;
import org.w3c.dom.Node;
import org.w3c.dom.NodeList;
import org.xml.sax.SAXException;

import javax.xml.parsers.DocumentBuilder;
import javax.xml.parsers.DocumentBuilderFactory;
import javax.xml.parsers.ParserConfigurationException;
import javax.xml.transform.*;
import javax.xml.transform.dom.DOMSource;
import javax.xml.transform.stream.StreamResult;
import java.io.*;
import java.nio.file.Files;
import java.nio.file.Path;
import java.util.Arrays;
import java.util.Objects;
import java.util.stream.Collectors;

/**
 * <p>Enables reading and writing NBT as XML.
 *
 * <p>Structure:
 *
 * <pre>{@code
 * <?xml version="1.0" encoding="UTF-8"?>
 * <n:compound xmlns:n="./xnbt.xsd">
 *     <n:byte key="byte">1</n:byte>
 *     <n:short key="short">2</n:short>
 *     <n:int key="int">3</n:int>
 *     <n:long key="long">4</n:long>
 *     <n:float key="float">5.0</n:float>
 *     <n:double key="double">6.0</n:double>
 *     <n:string key="string">seven</n:string>
 *     <n:list key="list">
 *         <n:int>1</n:int>
 *         <n:int>2</n:int>
 *         <n:int>3</n:int>
 *     </n:list>
 *     <n:array of="byte" key="array">1,2,3,4,5</n:array>
 * </n:compound>
 * }</pre>
 *
 * <p>Notes:
 *
 * <ul>
 *  <li>Objects inside a compound must have a {@code n:key} attribute, to designate their key in the compound map.
 *  <li>Objects inside a list must be all of the same type, as indicated by the {@code n:of} attribute.
 * </ul>
 */
public class XNbt
{
    private static final DocumentBuilderFactory DOCUMENT_BUILDER_FACTORY = DocumentBuilderFactory.newInstance();

    public static Tag read(String file) throws IOException, SAXException
    {
        return read(new File(file));
    }

    public static Tag read(File file) throws IOException, SAXException
    {
        try
        {
            DocumentBuilder dBuilder = DOCUMENT_BUILDER_FACTORY.newDocumentBuilder();
            return convert(dBuilder.parse(file));
        }
        catch (ParserConfigurationException e)
        {
            throw new RuntimeException("Error reading document", e);
        }
    }

    public static Tag read(Path path) throws IOException, SAXException
    {
        try(InputStream stream = Files.newInputStream(path))
        {
            return read(stream);
        }
    }

    public static Tag read(InputStream stream) throws IOException, SAXException
    {
        try
        {
            DocumentBuilder dBuilder = DOCUMENT_BUILDER_FACTORY.newDocumentBuilder();
            return convert(dBuilder.parse(stream));
        }
        catch (ParserConfigurationException e)
        {
            throw new RuntimeException("Error reading document", e);
        }
    }

    public static Tag convert(Document doc)
    {
        Element root = doc.getDocumentElement();

        root.normalize();

        return convert(root);
    }

    public static Tag convert(Element element)
    {
        return switch(element.getTagName())
        {
            case "n:byte" -> convertByte(element);
            case "n:short" -> convertShort(element);
            case "n:int" -> convertInt(element);
            case "n:long" -> convertLong(element);
            case "n:float" -> convertFloat(element);
            case "n:double" -> convertDouble(element);
            case "n:string" -> convertString(element);
            case "n:list" -> convertList(element);
            case "n:compound" -> convertCompound(element);
            case "n:array" -> convertArray(element);
            default -> throw new IllegalStateException("Do not know how to convert element with tag name: " + element.getTagName());
        };
    }

    public static Tag convertByte(Element element)
    {
        return ByteTag.valueOf(Byte.parseByte(element.getTextContent()));
    }

    public static Tag convertShort(Element element)
    {
        return ShortTag.valueOf(Short.parseShort(element.getTextContent()));
    }

    public static Tag convertInt(Element element)
    {
        return IntTag.valueOf(Integer.parseInt(element.getTextContent()));
    }

    public static Tag convertLong(Element element)
    {
        return LongTag.valueOf(Long.parseLong(element.getTextContent()));
    }

    public static Tag convertFloat(Element element)
    {
        return FloatTag.valueOf(Float.parseFloat(element.getTextContent()));
    }

    public static Tag convertDouble(Element element)
    {
        return DoubleTag.valueOf(Double.parseDouble(element.getTextContent()));
    }

    public static Tag convertString(Element element)
    {
        return StringTag.valueOf(element.getTextContent());
    }

    public static Tag convertList(Element element)
    {
        ListTag listTag = new ListTag();

        String firstName = null;
        NodeList children = element.getChildNodes();
        for (int i = 0; i < children.getLength(); i++)
        {
            var child = children.item(i);
            if (child.getNodeType() != Node.ELEMENT_NODE)
                continue;

            var childElement = (Element)child;

            if (firstName != null && firstName.equals(childElement.getTagName()))
                throw new IllegalStateException("NBT Lists must be homogenous (all children must be of the same type).");

            firstName = childElement.getTagName();

            listTag.add(convert(childElement));
        }

        return listTag;
    }

    public static Tag convertCompound(Element element)
    {
        CompoundTag compoundTag = new CompoundTag();

        NodeList children = element.getChildNodes();
        for (int i = 0; i < children.getLength(); i++)
        {
            var child = children.item(i);
            if (child.getNodeType() != Node.ELEMENT_NODE)
                continue;

            var childElement = (Element)child;

            if (!childElement.hasAttribute("key"))
                throw new IllegalStateException("Children of n:compound must have a 'key' attribute.");

            var key = childElement.getAttribute("key");

            compoundTag.put(key, convert(childElement));
        }

        return compoundTag;
    }

    public static Tag convertArray(Element element)
    {
        if (!element.hasAttribute("of"))
            throw new IllegalStateException("Arrays must have an 'of' attribute.");

        var of = element.getAttribute("of");

        return switch(of)
        {
            case "byte" -> new ByteArrayTag(Arrays.stream(element.getTextContent().split(",")).map(Byte::valueOf).toList());
            case "int" -> new IntArrayTag(Arrays.stream(element.getTextContent().split(",")).map(Integer::valueOf).toList());
            case "long" -> new LongArrayTag(Arrays.stream(element.getTextContent().split(",")).map(Long::valueOf).toList());
            default -> throw new IllegalStateException("NBT Arrays must be of type 'byte', 'int', or 'long'. Found " + of + " instead.");
        };
    }

    // ============================== SERIALIZER ====================================== //

    private static final TransformerFactory TRANSFORMER_FACTORY = TransformerFactory.newInstance();

    public static void write(Tag tag, String file) throws IOException, TransformerException
    {
        write(tag, new File(file));
    }

    public static void write(Tag tag, File file) throws IOException, TransformerException
    {
        try(FileWriter writer = new FileWriter(file))
        {
            write(tag, writer);
        }
    }

    public static void write(Tag tag, Writer writer) throws TransformerException
    {
        var doc = convert(tag);

        try
        {
            Transformer transformer = TRANSFORMER_FACTORY.newTransformer();
            DOMSource source = new DOMSource(doc);
            StreamResult result = new StreamResult(writer);

            transformer.setOutputProperty(OutputKeys.INDENT, "yes");
            transformer.setOutputProperty("{http://xml.apache.org/xslt}indent-amount", "2");

            transformer.transform(source, result);
        }
        catch (TransformerConfigurationException e)
        {
            throw new RuntimeException("Error encoding XML Document");
        }
    }

    public static Document convert(Tag tag)
    {
        try
        {
            DocumentBuilder dBuilder = DOCUMENT_BUILDER_FACTORY.newDocumentBuilder();

            var doc = dBuilder.newDocument();

            var root = convert(doc, tag);

            root.setAttribute("xmlns:n", "https://dogforce-games.com/xnbt/2021/XNbtSchema");

            doc.appendChild(root);

            return doc;
        }
        catch (ParserConfigurationException e)
        {
            throw new RuntimeException("Error constructing XML Document");
        }
    }

    public static Element convert(Document doc, Tag tag)
    {
        return switch(tag.getId())
                {
                    case Tag.TAG_END -> convertEnd(doc);
                    case Tag.TAG_BYTE -> convertByte(doc, (ByteTag) tag);
                    case Tag.TAG_SHORT -> convertShort(doc, (ShortTag) tag);
                    case Tag.TAG_INT -> convertInt(doc, (IntTag) tag);
                    case Tag.TAG_LONG -> convertLong(doc, (LongTag) tag);
                    case Tag.TAG_FLOAT -> convertFloat(doc, (FloatTag) tag);
                    case Tag.TAG_DOUBLE -> convertDouble(doc, (DoubleTag) tag);
                    case Tag.TAG_STRING -> convertString(doc, (StringTag) tag);
                    case Tag.TAG_LIST -> convertList(doc, (ListTag) tag);
                    case Tag.TAG_COMPOUND -> convertCompound(doc, (CompoundTag) tag);
                    case Tag.TAG_BYTE_ARRAY -> convertByteArray(doc, (ByteArrayTag) tag);
                    case Tag.TAG_INT_ARRAY -> convertIntArray(doc, (IntArrayTag) tag);
                    case Tag.TAG_LONG_ARRAY -> convertLongArray(doc, (LongArrayTag) tag);
                    case Tag.TAG_ANY_NUMERIC -> throw new IllegalStateException("ANY_NUMERIC tag id is not valid in this context.");
                    default -> throw new IllegalStateException("Unknown tag type " + tag.getId() + " with class " + tag.getClass());
                };
    }

    public static Element convertEnd(Document doc)
    {
        return doc.createElement("n:end");
    }

    public static Element convertByte(Document doc, ByteTag tag)
    {
        var element = doc.createElement("n:byte");
        element.setTextContent(Byte.toString(tag.getAsByte()));
        return element;
    }

    public static Element convertShort(Document doc, ShortTag tag)
    {
        var element = doc.createElement("n:short");
        element.setTextContent(Short.toString(tag.getAsShort()));
        return element;
    }

    public static Element convertInt(Document doc, IntTag tag)
    {
        var element = doc.createElement("n:int");
        element.setTextContent(Integer.toString(tag.getAsInt()));
        return element;
    }

    public static Element convertLong(Document doc, LongTag tag)
    {
        var element = doc.createElement("n:long");
        element.setTextContent(Long.toString(tag.getAsLong()));
        return element;
    }

    public static Element convertFloat(Document doc, FloatTag tag)
    {
        var element = doc.createElement("n:float");
        element.setTextContent(Float.toString(tag.getAsFloat()));
        return element;
    }

    public static Element convertDouble(Document doc, DoubleTag tag)
    {
        var element = doc.createElement("n:double");
        element.setTextContent(Double.toString(tag.getAsDouble()));
        return element;
    }

    public static Element convertString(Document doc, StringTag tag)
    {
        var element = doc.createElement("n:string");
        element.setTextContent(tag.getAsString());
        return element;
    }

    public static Element convertList(Document doc, ListTag tag)
    {
        var element = doc.createElement("n:list");

        for(var childTag : tag)
        {
            element.appendChild(convert(doc, childTag));
        }

        return element;
    }

    public static Element convertCompound(Document doc, CompoundTag tag)
    {
        var element = doc.createElement("n:compound");

        for(var key : tag.getAllKeys())
        {
            var childElement = convert(doc, Objects.requireNonNull(tag.get(key)));
            childElement.setAttribute("key", key);
            element.appendChild(childElement);
        }

        return element;
    }

    public static Element convertByteArray(Document doc, ByteArrayTag tag)
    {
        var element = doc.createElement("n:array");
        element.setAttribute("of", "byte");
        element.setTextContent(tag.stream().mapToInt(ByteTag::getAsInt).mapToObj(Integer::toString).collect(Collectors.joining(",")));
        return element;
    }

    public static Element convertIntArray(Document doc, IntArrayTag tag)
    {
        var element = doc.createElement("n:array");
        element.setAttribute("of", "int");
        element.setTextContent(tag.stream().mapToInt(IntTag::getAsInt).mapToObj(Integer::toString).collect(Collectors.joining(",")));
        return element;
    }

    public static Element convertLongArray(Document doc, LongArrayTag tag)
    {
        var element = doc.createElement("n:array");
        element.setAttribute("of", "long");
        element.setTextContent(tag.stream().mapToLong(LongTag::getAsLong).mapToObj(Long::toString).collect(Collectors.joining(",")));
        return element;
    }

}