#!/usr/bin/env python3

import sys
import logging
from typing import List, Dict, Set

import xml.etree.ElementTree as ET


###############################################################################
def rename(name: str):
    return name


###############################################################################
class Registry:
    def __init__(self):
        self.types = {}
        self.extensions = {}
        self.features = {}

        self.types['API Constants'] = Unscoped('API Constants')

        self.applied = set()

    def _serialise(self, name: str, queued: Set[str]):
        if name in queued:
            return []

        result = []
        obj = self.types[name]

        for d in obj.depends:
            if d == name:
                continue
            result += self._serialise(d, queued)

        assert name not in queued
        queued.add(name)
        result += [obj]

        return result

    def serialise(self, platform: Set[str]):
        required = []

        for (_, f) in self.features.items():
            required += f.apply(self)
            required.append(f.name)

        for e in self.extensions:
            required += self.extensions[e].apply(self, platform)

        queued = set()
        result = []

        for r in required:
            result += self._serialise(r, queued)

        return result


###############################################################################
class Type(object):
    """
    The base class for all object defined in the Vulkan API.

    This includes (but is not limited to) types, like structures; and values,
    like constants.
    """
    def __init__(self, name: str, depends: List[str] = None):
        assert name

        self.name = name
        self.depends = depends or []

        assert isinstance(self.depends, list)
        for i in self.depends:
            assert isinstance(i, str)

    def depends(self):
        return self.depends

    def declare(self):
        return ""

    def define(self, reg: Registry):
        return ""


###############################################################################
class AliasType(Type):
    """
    A type that is an alias for another type.

    May be serialised using an appropriate host language facility
    (eg, a typedef)
    """
    def __init__(self, name: str, target: str, depends: List[str]=None):
        depends = depends or []
        super().__init__(name, depends=depends+[target])
        self.target = target

    def declare(self):
        return f"using {rename(self.name)} = {rename(self.target)};"


# -----------------------------------------------------------------------------
class AliasValue(Type):
    """
    A value that is an alias for another value.

    May be serialised using an appropriate host language facility.
    """
    def __init__(self, name: str, target: str):
        super().__init__(name, depends=[target])
        self.target = target
        self.value = target

    def declare(self):
        return "constexpr auto %(name)s = %(target)s;" % {
            "name": rename(self.name),
            "target": rename(self.target),
        }


# -----------------------------------------------------------------------------
class Placeholder(Type):
    def __init__(self, name: str):
        super().__init__(name)



# -----------------------------------------------------------------------------
class Unscoped(Type):
    def __init__(self, name: str):
        super().__init__(name)
        self.values = []

    def declare(self):
        return "\n".join(t.declare() for t in self.values)

    def define(self,reg):
        return "\n".join(t.define(reg.types) for t in self.values)



###############################################################################
class Include(Type):
    def __init__(self, node):
        assert node.tag == 'type'
        assert node.attrib['category'] == 'include'

        super().__init__(node.attrib['name'])

        self.directive = node.text

    def declare(self):
        return self.directive or "#include <%s>" % self.name


class Define(Type):
    def __init__(self, node):
        assert node.tag == 'type'
        assert node.attrib['category'] == 'define'

        name = node.attrib.get('name') or node.find('name').text
        super().__init__(name)

        self.directive = "".join(node.itertext())

    def declare(self):
        return self.directive


class Bitmask(Type):
    def __init__(self, node):
        assert node.tag == 'type'
        assert node.attrib['category'] == 'bitmask'

        name = node.find('name').text
        type = node.find('type').text

        super().__init__(name,depends=[type])

        self.type = type
        self.requires = node.attrib.get('requires')
        if self.requires:
            self.depends.append(self.requires)

    def declare(self):
        return "using %(name)s = %(type)s;" % {
            "name": self.name,
            "type": self.type
        }

    def define(self, reg: Registry):
        return self.declare()

        if not self.requires:
            return self.declare()

        return "using %(name)s = %(requires)s;" % {
            "name": self.name,
            "requires": self.requires
        }

        source = reg.types[self.requires]
        members = ["%(k)s = %(v)s" % {"k":k, "v":v.value} for (k,v) in source.values.items()]

        return """enum %(name)s : %(type)s {
            %(members)s
        }""" % {
            "name": self.name,
            "type": self.type,
            "members": ",\n".join(members)
        }


class Handle(Type):
    parents: List[str]
    type: str

    def __init__(self, node):
        assert node.tag == 'type'
        assert node.attrib['category'] == 'handle'

        name = node.find('name').text
        type = node.find('type').text

        super().__init__(name, depends=[type])
        self.type = type

        parents = node.attrib.get('parent', None)
        self.parents = parents.split(',') if parents else []

        assert type

    def declare(self):
        return "struct %(name)s_t; using %(name)s = %(name)s_t*;" % {
            "name": self.name,
            "type": self.type
        }

    def has_parent(self, name: str, reg: Registry) -> bool:
        """
        Recursively check if this type is derived from a given parent type.
        """
        assert name
        assert reg

        if self.name == name:
            return True

        if not self.parents:
            return False
        if name in self.parents:
            return True

        for p in self.parents:
            if reg.types[p].has_parent(name, reg):
                return True
        return False


class Enum(Type):
    def __init__(self, node):
        assert node.tag == 'type'
        assert node.attrib['category'] == 'enum'

        name = node.attrib['name']
        super().__init__(name, depends=["VkEnum"])

        self.values = {}

    def __setitem__(self, key: str, value):
        assert isinstance(value, Constant) or isinstance(value, AliasValue)
        self.values[key] = value

    def declare(self):
        return ""
        return "enum %(name)s : int32_t;" % {
            "name": self.name
        }

    def define(self, reg: Registry):
        values = ("%(name)s = %(value)s" % { "name": k, "value": v.value } for (k,v) in self.values.items())

        return "enum %(name)s : int32_t { %(values)s };" % {
            "name": self.name,
            "values": ", ".join(values)
        }


class BaseType(AliasType):
    """
    Represents fundamental types that aliases of system provided types and used
    extensively by the base API. eg, VkBool32
    """

    def __init__(self, node):
        assert node.tag == 'type'
        assert node.attrib['category'] == 'basetype'

        super().__init__(
            node.find('name').text,
            node.find('type').text
        )


class FuncPointer(Type):
    def __init__(self, node):
        assert node.tag == 'type'
        assert node.attrib['category'] == 'funcpointer'

        name = node.find('name').text

        self.params = list(map(lambda x: x.text, node.findall('./type')))
        self.text = "".join(node.itertext())

        super().__init__(name, depends=['VkBool32']+self.params)

    def declare(self):
        return self.text


class POD(Type):
    def __init__(self, node):
        assert node.tag == 'type'
        assert node.attrib['category'] in ['struct', 'union']

        super().__init__(node.attrib['name'])
        self._node = node
        self._category = node.attrib['category']

        # sometimes there are enums hiding in the member fields being used as array sizes
        self.depends += list(e.text for e in node.findall('.//enum'))
        self.depends += list(t.text for t in node.findall('.//type'))

        self._members = []

        for member in node.findall('./member'):
            type = member.find('type').text
            name = member.find('name').text

            comment = member.find('comment')
            if not comment is None:
                member.remove(comment)
            code = " ".join(member.itertext())

            #code = member.iter()
            #code = filter(lambda x: x.tag != 'comment', code)
            #code = map(lambda x: x.itertext(), code)
            #code = map(lambda x: "".join(x), code)
            #code = "".join(code)

            self._members.append({'code': code, 'type': type, 'name': name})

    def declare(self):
        return "%(category)s %(name)s;" % {
            'category': self._category,
            'name': rename(self.name)
        }

    def define(self, reg: Registry):
        return "%(category)s %(name)s {\n%(members)s\n};" % {
            'category': self._category,
            'name': rename(self.name),
            'members': "\n".join(m['code'] + ';' for m in self._members)
        }


class Struct(POD):
    def __init__(self, node):
        super().__init__(node)


class Union(POD):
    def __init__(self, node):
        super().__init__(node)


class Constant(Type):
    def __init__(self, node, **kwargs):
        assert node.tag == 'enum'

        name = node.attrib['name']
        super().__init__(name)

        if 'offset' in node.attrib:
            assert 'extends' in node.attrib

            number = int(kwargs['extnumber'])
            offset = int(node.attrib['offset'])
            self.value = 1000000000 + 1000 * number + offset

            if 'dir' in node.attrib:
                self.value *= -1
        elif 'value' in node.attrib:
            self.value = node.attrib['value']
        elif 'bitpos' in node.attrib:
            self.value = "1 << %s"  % node.attrib['bitpos']
        else:
            raise RuntimeError("Unknown constant value type")

    def declare(self):
        return "constexpr auto %(name)s = %(value)s;" % {
            "name": self.name,
            "value": self.value
        }


class Command(Type):
    class Param(Type):
        def __init__(self, node, **kwargs):
            assert node.tag == 'param'

            super().__init__(
                name = node.find('name').text,
                depends=[node.find('type').text],
                **kwargs
            )

            self.type = node.find('type').text

            self.param = ""
            for i in node.iter():
                self.param += i.text or ""
                self.param += i.tail or ""
            # normalise whitespace
            self.param = " ".join(self.param.split())

    def __init__(self, node):
        assert node.tag == "command"
        proto = node.find('proto')
        name = proto.find('name').text
        super().__init__(name)

        self.result = proto.find('type').text
        self.params = [self.Param(p) for p in node.findall('./param')]
        self.depends += [self.result]
        for p in self.params:
            self.depends += p.depends

    def declare(self):
        return 'extern "C" %(result)s %(name)s (%(params)s) noexcept;' % {
            'name': rename(self.name),
            'result': self.result,
            'params': ", ".join(p.param for p in self.params)
        }

    def is_instance(self, reg: Registry):
        assert reg

        if not self.params:
            return True

        first_arg = self.params[0].type
        if first_arg == 'VkInstance':
            return True

        first_obj = reg.types[first_arg]

        # If the first type isn't a handle of any description then it should
        # be an instance function.
        if not isinstance(first_obj, Handle):
            return True

        # Both VkInstance and VkPhysicalDevice are listed as possible instance
        # parameters.
        #
        # Test that the handle is derived from VkInstance, and not derived from
        # VkDevice. The second test is required because VkDevice is indirectly
        # derived from VkInstance. This approach buys us a little more
        # generality.

        if not first_obj.has_parent('VkInstance', reg):
            return False
        if first_arg == 'VkDevice' or first_obj.has_parent('VkDevice', reg):
            return False
        return True

    def is_device(self, reg:Registry):
        return not self.is_instance(reg)


class Require(object):
    def __init__(self, root):
        self.values = []
        self.depends = []

        for node in root:
            if node.tag == 'enum':
                self.values.append(node)
            elif node.tag in ['command', 'type']:
                self.depends.append(node.attrib['name'])
            elif node.tag in ['comment']:
                pass
            else:
                raise RuntimeError("Unknown requires node")

    def apply(self, reg: Registry, extnumber=None):
        required = []
        required += self.depends

        for value in self.values:
            name = value.attrib['name']

            if len(value.attrib) == 1:
                assert 'name' in value.attrib
                required.append(name)
                continue

            if not 'extends' in value.attrib:
                obj = Constant(value)
                owner = reg.types['API Constants']
                owner.values.append(obj)
                continue

            owner = reg.types[value.attrib['extends']]

            if 'alias' in value.attrib:
                owner[name] = AliasValue(name, value.attrib['alias'])
                required.append(owner.name)
            elif value.tag == 'enum':
                owner[name] = Constant(value, extnumber=extnumber or int(value.attrib.get('extnumber', '0')))
                required.append(owner.name)
            elif value.tag == 'command':
                required.append(name)
            else:
                raise RuntimeError("Unknown type")

        return required


class Feature(Type):
    def __init__(self, root):
        assert root.tag == 'feature'

        name = root.attrib['name']
        super().__init__(name)

        self.requires = []
        for node in root:
            if 'require' == node.tag:
                self.requires.append(Require(node))
            else:
                raise RuntimeError("Unhandled feature node")

    def define(self, reg: Registry):
        return "#define %s" % self.name

    def apply(self, reg: Registry):
        logging.info("Applying feature:", self.name, file=sys.stderr)

        result = []
        for r in self.requires:
            result += r.apply(reg)

        return result


class Extension(Type):
    def __init__(self, root):
        assert root.tag == 'extension'

        name = root.attrib['name']
        super().__init__(name)

        if 'requires' in root.attrib:
            self.depends += root.attrib['requires'].split(',')

        self.number = int(root.attrib['number'])
        self.platform = root.attrib.get('platform')

        self.requires = []

        for node in root:
            if node.tag == 'require':
                self.requires.append(Require(node))
            else:
                raise RuntimeError("Unknown extension node")

    def apply(self, reg: Registry, platform: Set[str]):
        if self.name in reg.applied:
            return []
        reg.applied.add(self.name)

        if self.platform and self.platform not in platform:
            return []

        required = []

        for dep in self.depends:
            required = reg.extensions[dep].apply(reg, platform)

        logging.info("Applying extension:", self.name, file=sys.stderr)
        for node in self.requires:
            required += node.apply(reg,extnumber=self.number)
        return required



###############################################################################
def ignore_node(types:Dict[str, Type], root):
    pass

parse_comment = ignore_node
parse_vendorids = ignore_node
parse_platforms = ignore_node
parse_tags = ignore_node


def parse_types(reg: Registry, root):
    assert root.tag == 'types'

    for t in root.findall('type'):
        name = t.attrib.get ('name') or t.find('name').text
        assert name not in reg.types

        if 'alias' in t.attrib:
            name = t.attrib['name']
            target = t.attrib['alias']

            reg.types[name] = AliasType(name, target)
            continue

        category = t.attrib.get ('category')

        # if we don't have a category we should have a bare type that has a
        # dependency on something like a header.
        #
        # eg, 'Display' depends on 'X11/Xlib.h'
        if not category:
            reg.types[name] = Placeholder (name)
        else:
            # Whitelist the known types so we don't accidentally instantiate
            # something whacky

            supported_categories = {
                'include': Include,
                'define': Define,
                'bitmask': Bitmask,
                'basetype': BaseType,
                'handle': Handle,
                'enum': Enum,
                'funcpointer': FuncPointer,
                'struct': Struct,
                'union': Union,
            }

            concrete = supported_categories.get(category, None)
            if concrete:
                reg.types[name] = concrete(t)
            else:
                raise RuntimeError('unhandled type')

        if 'requires' in t.attrib:
            reg.types[name].depends.append(t.attrib['requires'])


# -----------------------------------------------------------------------------
def parse_enums(reg: Registry, root):
    assert root.tag == 'enums'
    ownername = root.attrib['name']
    owner = reg.types[ownername] if ownername != 'API Constants' else reg.types

    for node in root.findall('./enum'):
        valuename = node.attrib.get('name')

        assert 'requires' not in node.attrib

        if 'alias' in node.attrib:
            owner[valuename] = AliasValue(valuename, node.attrib['alias'])
        else:
            owner[valuename] = Constant(node)


# -----------------------------------------------------------------------------
def parse_commands(reg: Registry, root):
    assert root.tag == 'commands'

    for node in root.findall('./command'):
        name = node.attrib.get('name') or node.find('./proto/name').text
        assert name not in reg.types

        if 'alias' in node.attrib:
            reg.types[name] = AliasValue(name, node.attrib['alias'])
            continue

        reg.types[name] = Command(node)


# -----------------------------------------------------------------------------
def parse_feature(reg: Registry, root):
    assert root.tag == 'feature'

    name = root.attrib['name']
    assert name not in reg.features

    reg.features[name] = Feature(root)
    reg.types[name] = reg.features[name]


# -----------------------------------------------------------------------------
def parse_extensions(reg: Registry, root):
    assert root.tag == 'extensions'

    for node in root.findall('./extension'):
        name = node.attrib['name']
        assert name not in reg.extensions

        reg.extensions[name] = Extension(node)


###############################################################################
def write_header(path: str, q: List[Type], reg: Registry):
    with open(path, 'w') as dst:
        dst.write("#pragma once\n")

        # Write the declarations and definitions for all types.
        for obj in q:
            dst.write(obj.declare())
            dst.write('\n')
            dst.write(obj.define(reg))
            dst.write('\n')


        # Define the default case for device and instance type traits.
        dst.write("""
        #include <type_traits>

        /// A type trait that tests if a Vulkan type is an instance type
        template <typename NativeT>
        struct is_instance:
            public std::false_type
        {};


        /// A type trait that tests if a Vulkan type is a device type
        template <typename NativeT>
        struct is_device:
            public std::false_type
        {};


        template <typename T>
        constexpr auto is_instance_v = is_instance<T>::value;

        template <typename T>
        constexpr auto is_device_v = is_device<T>::value;
        """)

        # Specialise traits for device and instance types.
        for obj in q:
            if not isinstance(obj, Handle):
                continue

            device_value = "true_type" if obj.has_parent("VkDevice", reg) else "false_type"
            instance_value = "true_type" if obj.has_parent("VkInstance", reg) else "false_type"

            dst.write(f"""
            template <> struct is_instance<{obj.name}>: public std::{instance_value} {{ }};
            template <> struct is_device<{obj.name}>: public std::{device_value} {{ }};
            """)


# -----------------------------------------------------------------------------
def write_icd(path: str, q: List[Type], reg: Registry):
    with open(path, 'w') as icd:
        commands = [i for i in q if isinstance(i, Command)]
        instance_commands = [i for i in commands if i.is_instance(reg)]
        device_commands = [i for i in commands if i.is_device(reg)]

        assert len(instance_commands) + len(device_commands) == len(commands)

        icd.write(f"""
        #include "vk.hpp"

        #include <cruft/util/preprocessor.hpp>

        #define MAP_COMMANDS(FUNC) MAP0(FUNC,{",".join(i.name for i in commands)})
        #define MAP_INSTANCE_COMMANDS(FUNC) MAP0(FUNC,{",".join(i.name for i in instance_commands)})
        #define MAP_DEVICE_COMMANDS(FUNC) MAP0(FUNC,{",".join(i.name for i in device_commands)})

        namespace cruft::vk::icd {{
        class vendor;

        struct func {{
            void *handle;
            void const *table;
        }};

        struct instance_table {{
            instance_table (vendor &);
        """)

        for obj in instance_commands:
            icd.write(f"{obj.result} (*{obj.name}) ({','.join(p.param for p in obj.params)}) = nullptr;\n")

        icd.write("""};

        struct device_table {
        """)

        for obj in device_commands:
            icd.write(f"{obj.result} (*{obj.name}) ({','.join(p.param for p in obj.params)}) = nullptr;\n")

        icd.write("""
        };
        }
        """)


# -----------------------------------------------------------------------------
def write_dispatch(path: str, q: List[Type], reg: Registry):
    with open(path, 'w') as dispatch:
        dispatch.write("""
        #include "../vk.hpp"
        #include "vtable.hpp"

        #include "icd/dispatch.hpp"

        #include <cruft/util/debug.hpp>

        #pragma GCC diagnostic ignored "-Wunused-parameter"

        static cruft::vk::icd::instance_table const *i_table [[maybe_unused]] = nullptr;
        static cruft::vk::icd::device_table const *d_table [[maybe_unused]] = nullptr;

        void (*cruft_vk_icdGetInstanceProcAddr) (
            VkInstance instance,
            const char* pName
        ) = nullptr;

        void cruft::vk::icd::init (vendor const &impl)
        {
            cruft_vk_icdGetInstanceProcAddr = impl.vtable.GetInstanceProc;
        }

        """)

        for obj in (i for i in q if isinstance(i, Command)):
            first_arg = reg.types[obj.params[0].type]

            if not isinstance(first_arg, Handle):
                dispatch.write(f"""
                extern "C" {obj.result} {rename(obj.name)} ({", ".join(p.param for p in obj.params)}) noexcept {{
                    unimplemented ();
                }}""")
                continue

            if first_arg.has_parent('VkDevice', reg):
                table = "d_table";
            elif first_arg.has_parent('VkInstance', reg):
                table = 'i_table'
            else:
                raise Exception("Unknown param type")


            dispatch.write(f"""
            extern "C" {obj.result} {rename(obj.name)} ({", ".join(p.param for p in obj.params)}) noexcept {{
                using first_arg_t = std::decay_t<decltype({obj.params[0].name})>;

                if constexpr (is_instance_v<first_arg_t>) {{
                    auto const entry = reinterpret_cast<cruft::vk::icd::func const*> ({obj.params[0].name});
                    auto const *table = reinterpret_cast<decltype({table})> (entry->table);

                    return (table->{obj.name})(
                        reinterpret_cast<decltype({obj.params[0].name})> (entry->handle)
                        {", ".join([''] + [p.name for p in obj.params[1:]])}
                    );
                }} else {{
                    unimplemented ();
                }}
            }}
            """)




###############################################################################
import argparse


# -----------------------------------------------------------------------------
def main():
    logging.getLogger().setLevel(logging.WARNING)

    parser = argparse.ArgumentParser(description='Transform XML API specification into C++ headers')
    parser.add_argument('--src', type=str, help='the path to the XML file to transform')
    parser.add_argument('--dst', type=str, help='the output path for the result')
    parser.add_argument('--icd', type=str, help='the output path for the icd loading routines')
    parser.add_argument('--dispatch', type=str, help="the output path for function dispatch")
    parser.add_argument(
        '--platform',
        type=str,
        action='append',
        help='a platform to generate output for. may be specific multiple times"'
    )

    args = parser.parse_args()

    # Get a copy of the specification XML
    with open(args.src, 'r') as src:
        tree = ET.parse(src)
    root = tree.getroot()

    # Find a parser for each of the nodes in the XML
    reg = Registry()
    for node in root:
        target = "parse_%s" % node.tag
        globals()[target](reg, node)

    # Override some requested system types so that they fit more naturally in
    # our environment. eg, use appropriate C++ types, or cruft library
    # wrappers.
    reg.types['windows.h'].name = 'cruft/util/win32/windows.hpp'
    reg.types['void*'] = Placeholder('void*')
    reg.types['nullptr'] = Placeholder('nullptr')
    reg.types['VkEnum'] = AliasType('VkEnum', 'int32_t')
    reg.types['VK_DEFINE_NON_DISPATCHABLE_HANDLE'] = AliasType("VK_DEFINE_NON_DISPATCHABLE_HANDLE", "uint64_t")
    reg.types['VK_DEFINE_HANDLE'] = AliasType("VK_DEFINE_HANDLE", "void*")
    reg.types['VK_NULL_HANDLE'] = AliasValue("VK_NULL_HANDLE", "nullptr");

    # Request serialisation of all features
    features = [Feature(n) for n in root.findall('./feature')]
    features = dict((f.name,f) for f in features)

    #reg.extensions['VK_KHR_surface'].apply(reg, platform='xcb')

    # Request a minimal set of extensions that will almost certainly be
    # required for all applications.
    extensions = ["VK_KHR_swapchain", "VK_EXT_debug_report", "VK_KHR_external_memory"]
    q = reg.serialise(args.platform)

    # Finally write out the header, icd, and dispatch code.
    write_header(args.dst, q, reg)
    write_icd(args.icd, q, reg)
    write_dispatch(args.dispatch, q, reg)


# -----------------------------------------------------------------------------
if __name__ == '__main__':
    main()