# vhdl.py: A vhdl domain for the Sphinx documentation system
# Copyright (C) 2021 CESNET z.s.p.o.
# Author(s): Jindrich Dite <xditej01@stud.fit.vutbr.cz>
#            Jakub Cabal <cabal@cesnet.cz>
#
# SPDX-License-Identifier: BSD-3-Clause

from collections import defaultdict
from typing import Iterable, Tuple, List, Optional, Union

from docutils import nodes
from docutils.statemachine import StringList
from docutils.parsers.rst import directives

from sphinx import addnodes
from sphinx.addnodes import desc_signature, pending_xref
from sphinx.application import Sphinx
from sphinx.directives import ObjectDescription, ObjDescT
from sphinx.domains import Domain, Index, IndexEntry
from sphinx.roles import XRefRole
from sphinx.util.docutils import SphinxDirective
from sphinx.util.nodes import make_refnode
from sphinx.util import logging

from . import autodoc

logger = logging.getLogger(__name__)

def init_autodoc(domain: Domain):
    if not domain.data['autodoc_initialized']:
        domain.data['autodoc_initialized'] = True
        autodoc.init(domain.env.app.config.vhdl_autodoc_source_path)
        logger.info('SPHINX-VHDL: Parsing of VHDL files completed.')


class VHDLEnumTypeDirective(ObjectDescription):
    has_content = True
    required_arguments = 1

    def handle_signature(self, sig: str, signode: desc_signature):
        signode += addnodes.desc_sig_keyword(text='TYPE ')
        signode += addnodes.desc_name(text=sig)
        signode += addnodes.desc_sig_keyword(text=' IS')

        return sig

    def add_target_and_index(self, name: ObjDescT, sig: str, signode: desc_signature) -> None:
        name = f'vhdl-enum-{sig.lower()}'
        signode['ids'].append(name)
        if 'noindex' not in self.options:
            self.env.domaindata['vhdl']['types'].append((name, sig, 'Enumeration', self.env.docname))
            self.env.domaindata['vhdl']['refs']['types'][sig.split('.')[-1].lower()].append(
                (sig.lower(), (self.env.docname, name)))


class VHDLRecordTypeDirective(ObjectDescription):
    has_content = True
    required_arguments = 1

    def handle_signature(self, sig: str, signode: desc_signature):
        signode += addnodes.desc_sig_keyword(text='TYPE ')
        signode += addnodes.desc_name(text=sig)
        signode += addnodes.desc_sig_keyword(text=' IS RECORD')

        return sig

    def add_target_and_index(self, name: ObjDescT, sig: str, signode: desc_signature) -> None:
        name = f'vhdl-record-{sig.lower()}'
        signode['ids'].append(name)
        if 'noindex' not in self.options:
            self.env.domaindata['vhdl']['types'].append((name, sig, 'Record', self.env.docname))
            self.env.domaindata['vhdl']['refs']['types'][sig.split('.')[-1].lower()].append(
                (sig.lower(), (self.env.docname, name)))


class VHDLGeneralTypeDirective(ObjectDescription):
    has_content = True
    required_arguments = 1

    def handle_signature(self, sig: str, signode: desc_signature) -> ObjDescT:
        signode += addnodes.desc_sig_keyword(text='TYPE ')
        signode += addnodes.desc_name(text=sig.split(':')[0].strip())
        signode += addnodes.desc_sig_keyword(text=' : ')
        tempnode = nodes.entry('')
        self.state.nested_parse(StringList([sig.split(':', 1)[-1].strip()]), 0, tempnode)
        type_name = addnodes.desc_name()
        signode += type_name
        type_name += tempnode[0][0]
        return sig

    def add_target_and_index(self, name: ObjDescT, sig: str, signode: desc_signature) -> None:
        name = f'vhdl-type-{sig.lower()}'
        signode['ids'].append(name)
        if 'noindex' not in self.options:
            self.env.domaindata['vhdl']['types'].append((name, sig, 'Type', self.env.docname))
            self.env.domaindata['vhdl']['refs']['types'][sig.split('.')[-1].lower()].append(
                (sig.lower(), (self.env.docname, name)))

class VHDLEnumValDirective(ObjectDescription):
    has_content = True
    required_arguments = 1

    def handle_signature(self, sig: str, signode: desc_signature) -> ObjDescT:
        signode += addnodes.desc_name(text=sig)
        return sig


class VHDLRecordElementDirective(ObjectDescription):
    has_content = True
    required_arguments = 1

    def handle_signature(self, sig: str, signode: desc_signature) -> ObjDescT:
        signode += addnodes.desc_name(text=sig.split(':')[0].strip())
        signode += addnodes.desc_sig_keyword(text=' : ')
        tempnode = nodes.entry('')
        self.state.nested_parse(StringList([sig.split(':', 1)[-1].strip()]), 0, tempnode)
        type_name = addnodes.desc_name()
        signode += type_name
        type_name += tempnode[0][0]
        return sig


class VHDLFunctionDirective(ObjectDescription):
    has_content = True
    required_arguments = 1

    def handle_signature(self, sig: str, signode: desc_signature) -> ObjDescT:
        signode += addnodes.desc_sig_keyword(text='FUNCTION ')
        signode += addnodes.desc_name(text=sig.split()[0])
        signode += addnodes.desc_sig_keyword(text=' RETURNS ')
        signode += addnodes.desc_name(text=sig.split()[1])
        return sig


class VHDLEntityDirective(ObjectDescription):
    has_content = True
    required_arguments = 1

    def handle_signature(self, sig: str, signode: desc_signature) -> ObjDescT:
        signode += addnodes.desc_sig_keyword(text='ENTITY ')
        signode += addnodes.desc_name(text=sig)
        signode += addnodes.desc_sig_keyword(text=' IS')
        return sig

    def add_target_and_index(self, name: ObjDescT, sig: str, signode: desc_signature) -> None:
        name = f'vhdl-entity-{sig.lower()}'
        signode['ids'].append(name)
        if 'noindex' not in self.options:
            self.env.domaindata['vhdl']['refs']['entity'][sig.split('.')[-1].lower()].append(
                (sig.lower(), (self.env.docname, name))
            )


class VHDLEntityIOGenericDirective(SphinxDirective):
    has_content = True
    required_arguments = 1
    table_headers: Tuple[str, str, str, str]
    title: str
    id_title: str

    def get_fields_from_definition(self, definition: str) -> Union[Tuple[str, str, str], Tuple[str, str, str, str]]:
        raise NotImplementedError

    def run(self):
        table = nodes.table()
        group = nodes.tgroup()
        table += group

        head = nodes.thead()
        body = nodes.tbody()

        index = 0
        has_groups = False
        has_group_desc = False

        # Get all fields from ports or generics
        # Then check if ports or generics have groups and group description
        while len(self.content) > index:
            if len(self.content[index]) > 0 and not self.content[index][0].isspace():
                fields = self.get_fields_from_definition(self.content[index])
            index += 1
        index = 0

        row: Optional[nodes.row]
        row = nodes.row()
        # Define first line of table
        row += nodes.entry('', nodes.paragraph('', nodes.Text(self.table_headers[1])))
        row += nodes.entry('', nodes.paragraph('', nodes.Text(self.table_headers[2])))
        row += nodes.entry('', nodes.paragraph('', nodes.Text(self.table_headers[3])))
        row += nodes.entry('', nodes.paragraph('', nodes.Text(self.table_headers[4])))
        head += row

        row = None
        current_description_lines = StringList()
        description_entry = nodes.entry('')

        # Define table fields and their dimensions
        group += nodes.colspec(colwidth=10)
        group += nodes.colspec(colwidth=10)
        group += nodes.colspec(colwidth=10)
        group += nodes.colspec(colwidth=100)

        # Insert group header to table
        group += head
        group += body
        current_group = ''

        # Fill the table with content
        while len(self.content) > index:
            if len(self.content[index]) > 0 and not self.content[index][0].isspace():

                if row is not None:
                    body += row
                row = nodes.row()
                fields = self.get_fields_from_definition(self.content[index])
                has_groups = has_groups or len(fields) >= 4

                # If there is a group then fill first line of table with name of group, separators and description
                if has_groups:
                    if current_group != (fields[0]):
                        current_group = (fields[0])
                        has_group_desc = len(autodoc.groups_desc[current_group]) != 0 and has_groups

                        # Create nodes that contains name and description of group
                        group_name = nodes.entry('')
                        group_desc = nodes.entry('')
                        self.state.nested_parse(StringList(initlist=[fields[0].split(' ', 1)[1]]), 0, group_name)
                        self.state.nested_parse(StringList(autodoc.groups_desc[current_group]), 0, group_desc)

                        # Create row that contains information about group (name, description and separators)
                        separator = "====="
                        par = [
                            nodes.entry('', nodes.paragraph('', nodes.Text(separator))),
                            group_name,
                            nodes.entry('', nodes.paragraph('', nodes.Text(separator))),
                            group_desc if has_group_desc else nodes.entry('', nodes.paragraph('', nodes.Text(separator))),
                        ]
                        for p in par:
                            row += p
                        body += row
                        pass
                    row = nodes.row()

                # Fill the table with content
                if len(fields) == 3:
                    fields = "", fields[0], fields[1], fields[2]
                row_id = f'vhdl-{self.id_title}-{self.arguments[0].lower()}-{fields[1].lower()}'
                self.env.domaindata['vhdl']['refs'][self.id_title][fields[1].lower()].append(
                    (self.arguments[0].lower(), (self.env.docname, row_id)))
                row['ids'].append(row_id)

                row += nodes.entry('', nodes.paragraph('', nodes.Text(fields[1])))

                refnode = pending_xref(refdomain='vhdl', reftype='type', reftarget=fields[2])
                refnode += nodes.Text(fields[2])
                type_node = nodes.entry('')
                self.state.nested_parse(StringList(initlist=[fields[2]]), 0, type_node)
                row += type_node

                row += nodes.entry('', nodes.paragraph('', nodes.Text(fields[3])))

                self.state.nested_parse(current_description_lines.get_indented()[0], 0, description_entry)
                description_entry = nodes.entry('')
                row += description_entry
                current_description_lines = StringList()
            else:
                current_description_lines.append(self.content[index], source=self.content.info(index))
            index += 1
        if row is not None:
            self.state.nested_parse(current_description_lines.get_indented()[0], 0, description_entry)
            body += row
        return [addnodes.desc_name(text=self.title), table]


class VHDLPortsDirective(VHDLEntityIOGenericDirective):
    id_title = 'portsignal'
    title = 'Ports'
    table_headers = 'Group', 'Port', 'Type', 'Mode', 'Description'

    def get_fields_from_definition(self, definition: str) -> Union[Tuple[str, str, str], Tuple[str, str, str, str]]:
        try:
            if definition.strip().startswith("SPHINXGRP "):
                return definition.split(maxsplit=1)[1].strip(), "", "", ""
            elif '}' not in definition:
                return (
                    definition.split(":")[0].strip(),
                    definition.split(":", 1)[1].strip().split(maxsplit=1)[1],
                    definition.split(":", 1)[1].strip().split()[0]
            )
            else:
                return(
                    definition.split("}")[0].strip(),
                    definition.split("}")[1].split(":")[0].strip(),
                    definition.split("}")[1].split(":", 1)[1].strip().split(maxsplit=1)[1],
                    definition.split("}")[1].split(":", 1)[1].strip().split()[0]
                )
        except:
            raise ValueError(
                f'Malformed port definition, must be in the form `name : mode type`, got {definition}'
            )

class VHDLConstantsDirective(VHDLEntityIOGenericDirective):
    id_title = 'genconstant'
    title = 'Constants'
    table_headers = 'Group', 'Constant', 'Type', 'Value', 'Description'

    def get_fields_from_definition(self, definition: str) -> Tuple[str, str, str]:
        try:
            return (
                definition.split(':')[0].strip(),
                definition.split(':', 1)[1].split(':=')[0].strip(),
                definition.split(':=')[1].strip()
            )
        except:
            raise ValueError(
                f"Malformed constant definition, must be in the form `name : type := defaultValue`, got {definition}"
            )


class VHDLParametersDirective(VHDLEntityIOGenericDirective):
    id_title = 'parameters'
    title = 'Parameters'
    table_headers = 'Group', 'Parameter', 'Type', 'Mode', 'Description'

    def get_fields_from_definition(self, definition: str) -> Tuple[str, str, str]:
        try:
            return (
                definition.split(":")[0].strip(),
                definition.split(":", 1)[1].strip().split(maxsplit=1)[1],
                definition.split(":", 1)[1].strip().split()[0]
            )
        except:
            raise ValueError(
                f'Malformed port definition, must be in the form `name : mode type`, got {definition}'
            )


class VHDLGenericsDirective(VHDLEntityIOGenericDirective):
    id_title = 'gengeneric'
    title = 'Generics'
    table_headers = 'Group', 'Generic', 'Type', 'Default', 'Description'

    def get_fields_from_definition(self, definition: str) -> Union[Tuple[str, str, str], Tuple[str, str, str, str]]:
        try:
            if '}' not in definition:
                return (
                    definition.split(':')[0].strip(),
                    definition.split(':', 1)[1].split(':=')[0].strip(),
                    definition.split(':=')[1].strip()
            )
            else:
                return(
                    definition.split("}")[0].strip(),
                    definition.split("}")[1].split(":")[0].strip(),
                    definition.split("}")[1].split(':', 1)[1].split(':=')[0].strip(),
                    definition.split("}")[1].split(':=')[1].strip()
                )
        except:
            raise ValueError(
                f"Malformed generic definition, must be in the form `name : type := defaultValue`, got {definition}"
            )


class VHDLPackagesDirective(ObjectDescription):
    has_content = True
    required_arguments = 1

    def handle_signature(self, sig: str, signode: desc_signature) -> ObjDescT:
        signode += addnodes.desc_sig_keyword(text='PACKAGE ')
        signode += addnodes.desc_name(text=sig)
        signode += addnodes.desc_sig_keyword(text=' IS')
        return sig


class VHDLAutoEntityDirective(VHDLEntityDirective):
    option_spec = {
        'noautoports': directives.flag,
        'noautogenerics': directives.flag
    }

    def handle_signature(self, sig: str, signode: desc_signature) -> ObjDescT:
        init_autodoc(self.env.domains['vhdl'])
        try:
            my_entity = autodoc.entities[sig.lower()]
            self.content = self.content + StringList(['', ''] + autodoc.entities[sig.lower()])
            if 'noautogenerics' not in self.options:
                self.content = self.content + StringList(['', f'.. vhdl:autogenerics:: {sig}', ''])
            if 'noautoports' not in self.options:
                self.content = self.content + StringList(['', f'.. vhdl:autoports:: {sig}', ''])
        except:
            logger.warning(f"SPHINX-VHDL: Entity {sig.lower()} was not found in parsed VHDL files!", location=self.get_location())
            self.content = self.content + StringList([f"SPHINX-VHDL: Entity was not found in parsed VHDL files!"])

        return super().handle_signature(sig, signode)


class VHDLAutoRecordDirective(VHDLRecordTypeDirective):
    def handle_signature(self, sig: str, signode: desc_signature) -> ObjDescT:
        init_autodoc(self.env.domains['vhdl'])
        self.content = self.content + StringList(['', ''] + autodoc.records[sig])
        for key, value in autodoc.record_elements[sig].items():
            self.content = self.content + StringList(['', '', f'.. vhdl:recordelem:: {key}', ''] + ['  ' + x for x in value])

        return super().handle_signature(sig, signode)


class VHDLAutoFunctionDirective(VHDLFunctionDirective):
    def handle_signature(self, sig: str, signode: desc_signature) -> ObjDescT:
        init_autodoc(self.env.domains['vhdl'])
        identifier = get_closest_identifier(sig.lower(), list(autodoc.functions.items()))
        if identifier is None:
            logger.warning(f"SPHINX-VHDL: Function {sig.lower()} was not found in parsed VHDL files!", location=self.get_location())
            self.content = StringList([f"SPHINX-VHDL: Function was not found in parsed VHDL files!"]) + self.content
            sig = f'{sig.lower()} Unknown'
        else:
            self.content = self.content + StringList(['', ''] + identifier[1])
            sig = f'{identifier[0].split(".")[-1]} {identifier[0].split(".")[0]}'

        return super().handle_signature(sig, signode)


class VHDLAutoEnumDirective(VHDLEnumTypeDirective):
    def handle_signature(self, sig: str, signode: desc_signature) -> ObjDescT:
        init_autodoc(self.env.domains['vhdl'])
        self.content = self.content + StringList(['', ''] + autodoc.enums[sig])
        for key, value in autodoc.enumvals[sig].items():
            self.content = self.content + StringList(['', '', f'.. vhdl:enumval:: {key}', ''] + ['  ' + x for x in value])
        return super().handle_signature(sig, signode)


class VHDLAutoPackageDirective(VHDLPackagesDirective):
    def handle_signature(self, sig: str, signode: desc_signature) -> ObjDescT:
        init_autodoc(self.env.domains['vhdl'])
        identifier = get_closest_identifier(sig.lower(), list(autodoc.packages.items()))
        if identifier is None:
            logger.warning(f"SPHINX-VHDL: Package {sig.lower()} was not found in parsed VHDL files!", location=self.get_location())
            self.content = StringList([f"SPHINX-VHDL: Package was not found in parsed VHDL files!"]) + self.content
        else:
            self.content = StringList(identifier[1] + ['', '']) + self.content
        return super().handle_signature(sig, signode)


class VHDLAutoPortsDirective(VHDLPortsDirective):
    has_content = False

    def run(self):
        init_autodoc(self.env.domains['vhdl'])
        self.content = StringList(
            [item for subitem in [[key, *[f'  {x}' for x in autodoc.portsignals[self.arguments[0].lower()][key]]]
                                  for key in autodoc.portsignals[self.arguments[0].lower()].keys()] for item in subitem]
        )
        return super().run()


class VHDLAutoGenericsDirective(VHDLGenericsDirective):
    has_content = False

    def run(self):
        init_autodoc(self.env.domains['vhdl'])
        self.content = StringList(
            [item for subitem in [[key, *[f'  {x}' for x in autodoc.generics[self.arguments[0].lower()][key]]]
                                  for key in autodoc.generics[self.arguments[0].lower()].keys()] for item in subitem]
        )
        return super().run()

class VHDLAutoConstantsDirective(VHDLConstantsDirective):
    has_content = False

    def run(self):
        init_autodoc(self.env.domains['vhdl'])
        self.content = StringList(
            [item for subitem in [[key, *[f'  {x}' for x in autodoc.constants[self.arguments[0].lower()][key]]]
                                  for key in autodoc.constants[self.arguments[0].lower()].keys()] for item in subitem]
        )
        return super().run()

class VHDLAutoTypeDirective(VHDLGeneralTypeDirective):

    def handle_signature(self, sig: str, signode: desc_signature) -> ObjDescT:
        init_autodoc(self.env.domains['vhdl'])
        identifier = get_closest_identifier(sig.lower(), autodoc.types.items())
        if identifier is None:
            logger.warning(f"SPHINX-VHDL: Type {sig.lower()} was not found in parsed VHDL files!", location=self.get_location())
            self.content = StringList([f"SPHINX-VHDL: Type was not found in parsed VHDL files!"]) + self.content
            return super().handle_signature(sig + " : Unknown", signode)
        else:
            self.content = self.content + StringList(['', ''] + identifier[1][1])
            return super().handle_signature(sig + " : " + identifier[1][0], signode)


class VHDLTypeIndex(Index):
    name = 'typeindex'
    localname = "Type Index"
    shortname = 'Types'

    def __init__(self, *args, **kwargs):
        super(VHDLTypeIndex, self).__init__(*args, **kwargs)

    def generate(self, docnames: Iterable[str] = None) -> Tuple[List[Tuple[str, List[IndexEntry]]], bool]:
        if docnames is not None:
            type_list = sorted([x for x in self.domain.data['types'] if x[3] in docnames], key=lambda x: x[1])
        else:
            type_list = sorted(self.domain.data['types'], key=lambda x: x[1])

        result: List[Tuple[str, List[IndexEntry]]] = []
        for name, sig, kind, docname in type_list:
            result.append((sig[0], [IndexEntry(sig, 0, docname, name, f'{kind} Type', '', '')]))

        return result, True


def get_closest_identifier(target_identifier: str, search_through: List[Tuple[str, ObjDescT]]):
    """
    Finds the item with the closes matching identifier to a target one in a list
    :param target_identifier: an identifier to match against
    :param search_through: List of pairs of an identifier and any other bound data
    :return: The tuple with closes match or None
    """
    identifier_part = target_identifier.split('.')
    option_list = []
    match = False
    for x in search_through:
        a = 0
        for y in x[0].split('.'):
            if y in identifier_part:
                match = True
                a += 1
        option_list.append((a, x))

    if match:
        return max(option_list, key=lambda z: z[0])[1]
    else:
        return None


class VHDLDomain(Domain):
    name = 'vhdl'
    label = 'VHDL Language'
    directives = {
        'enum': VHDLEnumTypeDirective,
        'enumval': VHDLEnumValDirective,
        'entity': VHDLEntityDirective,
        'autoentity': VHDLAutoEntityDirective,
        'autoports': VHDLAutoPortsDirective,
        'autopackage': VHDLAutoPackageDirective,
        'autogenerics': VHDLAutoGenericsDirective,
        'autoconstants': VHDLAutoConstantsDirective,
        'autorecord': VHDLAutoRecordDirective,
        'autoenum': VHDLAutoEnumDirective,
        'autotype': VHDLAutoTypeDirective,
        'autofunction': VHDLAutoFunctionDirective,
        'function': VHDLFunctionDirective,
        'parameters': VHDLParametersDirective,
        'ports': VHDLPortsDirective,
        'generics': VHDLGenericsDirective,
        'constants': VHDLConstantsDirective,
        'package': VHDLPackagesDirective,
        'record': VHDLRecordTypeDirective,
        'recordelem': VHDLRecordElementDirective,
        'type': VHDLGeneralTypeDirective,
    }
    initial_data = {
        'types': [],
        'refs': {
            'types': defaultdict(list),
            'portsignal': defaultdict(list),
            'gengeneric': defaultdict(list),
            'genconstant': defaultdict(list),
            'parameters': defaultdict(list),
            'entity': defaultdict(list),
        },
        'autodoc_initialized': False
    }
    indices = {
        VHDLTypeIndex
    }
    roles = {
        'portsignal': XRefRole(),
        'genconstant': XRefRole(),
        'gengeneric': XRefRole(),
        'type': XRefRole(),
        'entity': XRefRole(),
    }

    def resolve_xref(self, env: "BuildEnvironment", fromdocname: str, builder: "Builder", typ: str, target: str,
                     node: pending_xref, contnode: nodes.Element) -> Optional[nodes.Element]:
        if typ == 'type':
            index = self.data['refs']['types']
        elif typ == 'portsignal':
            index = self.data['refs']['portsignal']
        elif typ == 'gengeneric':
            index = self.data['refs']['gengeneric']
        elif typ == 'genconstant':
            index = self.data['refs']['genconstant']
        elif typ == 'entity':
            index = self.data['refs']['entity']
        elif True:
            raise NotImplementedError
        simple_name = target.split('.')[-1].lower()
        if simple_name in index:
            target_address = get_closest_identifier(target.lower(), index[simple_name])
            if target_address is None:
                logger.warning(f"SPHINX-VHDL: Unknown reference {target} discovered by resolve_xref function!")
            else:
                result = make_refnode(builder, fromdocname,
                                    target_address[1][0],
                                    target_address[1][1],
                                    contnode)
                return result


def setup(app: Sphinx):
    app.add_domain(VHDLDomain)
    app.add_config_value('vhdl_autodoc_source_path', '.', 'env', [str, list])
    logger.verbose('The sphinx-vhdl extension has been activated.')

    return {
        'version': '0.2'
    }
