# UGO.py by pbsds
# Hatenatools3 Python 3 port by RekuNote aka rekushiiiiii
# AGPL3 licensed
#
# Python 3 port. This class reads/writes Flipnote Hatena UGO files and
# imports/exports the project's UGOXML representation.

from __future__ import annotations

import os
import sys
from base64 import b64decode, b64encode
import xml.etree.ElementTree as ET

try:
    # Import the actual PPM submodule. Importing ``PPM`` from the package
    # namespace can resolve to the convenience-exported PPM class instead.
    from .PPM import TMB as _TMB
    HasPPM = True
except (ImportError, ValueError):
    try:
        from PPM import TMB as _TMB  # type: ignore
        HasPPM = True
    except ImportError:
        _TMB = None
        HasPPM = False


def AscDec(data, LittleEndian=False):
    """Convert a byte string to an integer (legacy Hatenatools helper)."""
    if isinstance(data, str):
        data = data.encode("latin-1")
    return int.from_bytes(bytes(data), "little" if LittleEndian else "big")


def DecAsc(dec, length=None, LittleEndian=False):
    """Convert an integer to bytes (legacy Hatenatools helper)."""
    if dec < 0:
        raise ValueError("DecAsc does not support negative values")
    if length is None:
        length = max(1, (dec.bit_length() + 7) // 8) if dec else 0
    if length == 0:
        return b""
    mask = (1 << (length * 8)) - 1
    dec &= mask
    return dec.to_bytes(length, "little" if LittleEndian else "big")


def zipalign(length, r=4):
    return length + (r - length % r) if length % r else length


def indentXML(elem, level=0):
    i = "\n" + level * "\t"
    if len(elem):
        if not elem.text or not elem.text.strip():
            elem.text = i + "\t"
        if not elem.tail or not elem.tail.strip():
            elem.tail = i
        for child in elem:
            indentXML(child, level + 1)
        if not child.tail or not child.tail.strip():
            child.tail = i
    elif level and (not elem.tail or not elem.tail.strip()):
        elem.tail = i


class UGO:
    def __init__(self):
        self.Loaded = False
        self.Items = []
        self.Files = []

    def ReadFile(self, path):
        with open(path, "rb") as f:
            return self.Read(f.read())

    def Read(self, data):
        global HasPPM
        if isinstance(data, str):
            data = data.encode("latin-1")
        data = bytes(data)

        if data[:4] != b"UGAR":
            return False

        sections = AscDec(data[4:8], True)
        self.TableLength = AscDec(data[8:12], True) if sections >= 1 else 0
        self.ExtraLength = AscDec(data[12:16], True) if sections >= 2 else 0
        if sections > 2:
            print("Warning: This UGO file has more than the 2 known sections:", sections)
            print("Please send this UGO file to pbsds over at pbsds.net")
            print("This file could possibly be read incorrectly...")

        headerlength = 8 + sections * 4

        if sections >= 1:
            table_bytes = data[headerlength : headerlength + self.TableLength]
            table_text = table_bytes.decode("utf-8")
            TableOfContents = tuple(line.split("\t") for line in table_text.split("\n") if line != "")
        else:
            TableOfContents = []

        if sections >= 2:
            extra_start = zipalign(headerlength + self.TableLength)
            ExtraData = data[extra_start : extra_start + self.ExtraLength]
        else:
            ExtraData = b""

        self.Items = []
        self.Files = []
        pos = 0
        tmbcount = 1
        ntftcount = 1
        names = []

        for entry in TableOfContents:
            if not entry or not entry[0]:
                continue
            item_type = int(entry[0])
            if item_type == 0:
                self.Items.append(("layout", list(map(int, entry[1:]))))
                continue
            if item_type == 1:
                num = int(entry[1])
                labels = [b64decode(entry[n]).decode("UTF-16LE") for n in range(2, 7)]
                self.Items.append(("topscreen text", labels, num))
                continue
            if item_type == 2:
                link = entry[1]
                label = b64decode(entry[2]).decode("UTF-16LE")
                selected = int(entry[3]) != 0
                self.Items.append(("category", link, label, selected))
                continue
            if item_type == 3:
                link = entry[1]
                label = b64decode(entry[2]).decode("UTF-16LE")
                self.Items.append(("post", link, label))
                continue
            if item_type == 4:
                link = entry[1]
                trait = int(entry[2])
                label = b64decode(entry[3]).decode("UTF-16LE")
                other = list(entry[4:])

                embedded = None
                if trait < 100 and ExtraData:
                    if ExtraData == b"\x20":
                        pass
                    elif ExtraData[pos : pos + 4] == b"PARA":
                        filedata = ExtraData[pos : pos + 0x6A0]
                        pos += 0x6A0
                        if HasPPM and _TMB is not None:
                            tmb = _TMB().Read(filedata)
                            name = tmb.CurrentFilename[:-4] if tmb else "embedded tmb #%i" % tmbcount
                            if not tmb:
                                tmbcount += 1
                        else:
                            name = "embedded tmb #%i" % tmbcount
                            tmbcount += 1

                        if name + ".tmb" in names:
                            j = 2
                            while "%s_%i.tmb" % (name, j) in names:
                                j += 1
                            name = "%s_%i" % (name, j)
                        embedded = (name + ".tmb", filedata)
                        names.append(name + ".tmb")
                    else:
                        name = label.encode("ascii", "ignore").decode("ascii")
                        if not name:
                            name = "nameless ntft %i" % ntftcount
                            ntftcount += 1
                        if name + ".ntft" in names:
                            j = 2
                            while "%s_%i.ntft" % (name, j) in names:
                                j += 1
                            name = "%s_%i" % (name, j)
                        embedded = (name + ".ntft", ExtraData[pos : pos + 2048])
                        names.append(name + ".ntft")
                        pos += 2048

                self.Items.append(("button", trait, label, link, other, embedded))
                continue

            self.Items.append(("unknown", list(entry)))
            print("Unknown UGO item discovered:", entry)

        self.Loaded = True
        return self

    def WriteFile(self, path):
        if not self.Loaded:
            return False
        out = self.Pack()
        if out is False:
            return False
        with open(path, "wb") as f:
            f.write(out)
        return True

    def Pack(self):
        if not self.Loaded:
            return False

        table_rows = []
        extra_parts = []

        for item in self.Items:
            kind = item[0]
            if kind == "unknown":
                table_rows.append("\t".join(map(str, item[1])))
            elif kind == "layout":
                table_rows.append("\t".join(["0"] + [str(value) for value in item[1]]))
            elif kind == "topscreen text":
                labels, num = item[1:]
                encoded_labels = [b64encode(label.encode("UTF-16LE")).decode("ascii") for label in labels]
                # The format expects exactly five labels.
                encoded_labels = (encoded_labels + [""] * 5)[:5]
                table_rows.append("\t".join(["1", str(num)] + encoded_labels))
            elif kind == "category":
                link, label, selected = item[1:]
                encoded = b64encode(label.encode("UTF-16LE")).decode("ascii")
                table_rows.append("\t".join(("2", link, encoded, str(int(bool(selected))))))
            elif kind == "post":
                link, label = item[1:]
                encoded = b64encode(label.encode("UTF-16LE")).decode("ascii")
                table_rows.append("\t".join(("3", link, encoded)))
            elif kind == "button":
                trait, label, link, other, embedded = item[1:]
                encoded = b64encode(label.encode("UTF-16LE")).decode("ascii")
                table_rows.append("\t".join(["4", link, str(trait), encoded] + [str(v) for v in other]))
                if embedded:
                    extra_parts.append(bytes(embedded[1]))
            else:
                print("Unrecognized entry in self.Items:", item)

        table = "\n".join(table_rows).encode("utf-8")
        extra = b"".join(extra_parts)

        sections = 0
        lengths = []
        if table:
            sections += 1
            lengths.append(len(table))
        if extra:
            sections += 1
            lengths.append(len(extra))

        header = b"UGAR" + DecAsc(sections, 4, True) + b"".join(DecAsc(length, 4, True) for length in lengths)
        if table:
            table += b"\0" * ((-len(table)) % 4)
        if extra:
            extra += b"\0" * ((-len(extra)) % 4)
        return header + table + extra

    def WriteXML(self, xmlname="content.ugoxml", folder="content.ugoxml embedded"):
        if not self.Loaded:
            return False

        path, xml_basename = os.path.split(xmlname)
        ugo_xml = ET.Element("ugo_xml")
        files = []

        for item in self.Items:
            kind = item[0]
            if kind == "unknown":
                elem = ET.SubElement(ugo_xml, "raw", type=str(item[1][0]))
                for value in item[1][1:]:
                    ET.SubElement(elem, "value").text = str(value)
            elif kind == "layout":
                elem = ET.SubElement(ugo_xml, "layout")
                for value in item[1]:
                    ET.SubElement(elem, "value").text = str(value)
            elif kind == "topscreen text":
                elem = ET.SubElement(ugo_xml, "title")
                labels, num = item[1:]
                for label in labels:
                    ET.SubElement(elem, "label").text = label
                ET.SubElement(elem, "num").text = str(num)
            elif kind == "category":
                elem = ET.SubElement(ugo_xml, "category")
                link, label, selected = item[1:]
                ET.SubElement(elem, "label").text = label
                ET.SubElement(elem, "address").text = link
                ET.SubElement(elem, "selected").text = str(bool(selected)).lower()
            elif kind == "post":
                elem = ET.SubElement(ugo_xml, "post")
                link, label = item[1:]
                ET.SubElement(elem, "label").text = label
                ET.SubElement(elem, "address").text = link
            elif kind == "button":
                elem = ET.SubElement(ugo_xml, "button")
                trait, label, link, other, embedded = item[1:]
                ET.SubElement(elem, "label").text = label
                ET.SubElement(elem, "address").text = link
                ET.SubElement(elem, "trait").text = str(trait)
                for n, value in enumerate(other):
                    entry = ET.SubElement(elem, "value")
                    entry.text = str(value)
                    if n == 0 and trait == 3:
                        entry.attrib["tip"] = "stars"
                if embedded:
                    ET.SubElement(elem, "embedded_file").text = os.path.join(folder, embedded[0])
                    files.append((os.path.join(folder, embedded[0]), bytes(embedded[1])))

        indentXML(ugo_xml)
        files.append((xml_basename, ET.tostring(ugo_xml, encoding="UTF-8")))

        folder_path = os.path.join(path, folder)
        if folder and not os.path.isdir(folder_path):
            os.makedirs(folder_path, exist_ok=True)
        for name, data in files:
            full_path = os.path.join(path, name)
            parent = os.path.dirname(full_path)
            if parent:
                os.makedirs(parent, exist_ok=True)
            with open(full_path, "wb") as f:
                f.write(data)
        return True

    def ReadXML(self, xmlfile, silent=True):
        ugo_xml = ET.parse(xmlfile).getroot()
        xmlpath = os.path.split(xmlfile)[0]

        items = []
        for elem in ugo_xml:
            if elem.tag == "raw":
                if "type" not in elem.attrib:
                    if not silent:
                        print('Invalid formatting. <raw> without "type" attribute')
                    return False
                values = [elem.attrib["type"]]
                for value in elem:
                    if value.tag != "value":
                        if not silent:
                            print("Invalid formatting. <%s> found within <unknown>" % value.tag)
                        return False
                    values.append(value.text or "")
                items.append(("unknown", values))

            elif elem.tag == "layout":
                values = []
                for value in elem:
                    if value.tag != "value":
                        if not silent:
                            print("Invalid formatting. <%s> found within <layout>" % value.tag)
                        return False
                    if not (value.text or "").isdigit():
                        if not silent:
                            print("Invalid entry. <value> in <layout> is not a number")
                        return False
                    values.append(int(value.text))
                items.append(("layout", values))

            elif elem.tag == "title":
                labels = ["", "", "", "", ""]
                num = 0
                pos = 0
                numset = False
                for value in elem:
                    if value.tag not in ("label", "num"):
                        if not silent:
                            print("Invalid formatting. <%s> found within <title>" % value.tag)
                        return False
                    if value.tag == "label":
                        if pos >= 5:
                            if not silent:
                                print("Invalid formatting. More than 5 <labels> in <title>")
                            return False
                        labels[pos] = value.text or ""
                        pos += 1
                    else:
                        if numset:
                            if not silent:
                                print("Invalid formatting. Multiple <num> in <title>")
                            return False
                        if not (value.text or "").isdigit():
                            if not silent:
                                print("Invalid entry. <num> in <title> is not a number!")
                            return False
                        num = int(value.text)
                        numset = True
                items.append(("topscreen text", labels, num))

            elif elem.tag == "category":
                link = label = selected = None
                for value in elem:
                    if value.tag not in ("label", "address", "selected"):
                        if not silent:
                            print("Invalid formatting. <%s> found within <category>" % value.tag)
                        return False
                    if value.tag == "address":
                        if link is not None:
                            return False
                        link = value.text or ""
                    elif value.tag == "label":
                        if label is not None:
                            return False
                        label = value.text or ""
                    else:
                        if selected is not None:
                            return False
                        selected = bool(value.text) and value.text[0].lower() in "t1"
                if link is None or label is None or selected is None:
                    return False
                items.append(("category", link, label, selected))

            elif elem.tag == "post":
                label = link = None
                for value in elem:
                    if value.tag == "label":
                        if label is not None:
                            return False
                        label = value.text or ""
                    elif value.tag == "address":
                        if link is not None:
                            return False
                        link = value.text or ""
                    else:
                        if not silent:
                            print("Invalid formatting. <%s> found within <post>" % value.tag)
                        return False
                if link is None or label is None:
                    return False
                items.append(("post", link, label))

            elif elem.tag == "button":
                trait = label = link = None
                other = []
                embedded = None
                for value in elem:
                    if value.tag not in ("label", "address", "trait", "value", "embedded_file"):
                        if not silent:
                            print("Invalid formatting. <%s> found within <button>" % value.tag)
                        return False
                    if value.tag == "label":
                        if label is not None:
                            return False
                        label = value.text or ""
                    elif value.tag == "address":
                        if link is not None:
                            return False
                        link = value.text or ""
                    elif value.tag == "trait":
                        if trait is not None or not (value.text or "").isdigit():
                            return False
                        trait = int(value.text)
                    elif value.tag == "value":
                        other.append(value.text or "")
                    else:
                        if embedded is not None:
                            return False
                        if not value.text:
                            return False
                        embedded_path = os.path.join(xmlpath, value.text)
                        if not os.path.isfile(embedded_path):
                            if not silent:
                                print('Invalid entry. Embedded file "%s" not found!' % value.text)
                            return False
                        with open(embedded_path, "rb") as f:
                            embedded = (os.path.split(value.text)[1], f.read())
                if trait is None or label is None or link is None:
                    return False
                items.append(("button", trait, label, link, other, embedded))

            else:
                if not silent:
                    print("Invalid formatting: <%s> found within <ugo_xml>" % elem.tag)
                return False

        self.Items = items
        self.Loaded = True
        return self


def _main():
    print("              ==      UGO.py      ==")
    print("             ==      by pbsds      ==")
    print("              ==       v0.93      ==")
    print()
    if len(sys.argv) < 2:
        print("Usage:")
        print("      UGO.py [<mode>] <input> [<output> [<foldername>]]")
        print()
        print("      <Mode>:")
        print("          -d: Converts the UGO file in <input> to a UGOXML file with the same")
        print("              name, unless <output> is specified. Any embedded files will be")
        print('              written to a folder called UGOXML-filename + " embedded" unless')
        print("              <foldername> is given.")
        print("              <foldername> is relative to the XML.")
        print("          -e: Converts the UGOXML file in <input> to a UGO file with the same")
        print("              name, unless <output> is specified.")
        print("          If mode is not specified, it will try to find out for itself")
        return 0

    mode = sys.argv[1]
    if mode not in ("-d", "-e"):
        if os.path.exists(mode):
            with open(mode, "rb") as f:
                magic = f.read(4)
            mode = "-d" if magic == b"UGAR" else "-e"
            print("No mode specified. %s chosen" % ("UGO -> UGOXML" if mode == "-d" else "UGOXML -> UGO"))
            sys.argv.insert(1, mode)
        else:
            print("Invalid <mode> given!")
            return 1

    input_path = sys.argv[2]
    if mode == "-d":
        output = sys.argv[3] if len(sys.argv) >= 4 else input_path + "xml"
        foldername = sys.argv[4] if len(sys.argv) >= 5 else os.path.split(output)[1] + " embedded"
        print("Reading %s..." % os.path.split(input_path)[1])
        ugo = UGO().ReadFile(input_path)
        if not ugo:
            print("Error!\n The given file is not a UGO file!")
            return 1
        print("Writing XML...")
        ugo.WriteXML(output, foldername)
        print("Done!")
        return 0

    output = sys.argv[3] if len(sys.argv) >= 4 else ".".join(input_path.split(".")[:-1]) + ".ugo"
    print("Reading %s..." % os.path.split(input_path)[1])
    try:
        ugo = UGO().ReadXML(input_path, False)
    except ET.ParseError:
        print("Error!\nThe given file is not in the XML format!")
        return 1
    if not ugo:
        return 1
    print("Writing UGO...")
    ugo.WriteFile(output)
    print("Done")
    return 0


if __name__ == "__main__":
    raise SystemExit(_main())