"""SystemVerilog language parser using tree-sitter.

Supports extraction of: module declarations, function/task definitions,
always blocks, and initial blocks from SystemVerilog / Verilog source code.
"""

from __future__ import annotations

from typing import Any

import tree_sitter_systemverilog
from tree_sitter import Node

from viper.languages.base import Function, TreeSitterParser


class SystemVerilogTreeSitterParser(TreeSitterParser):
    """Tree-sitter based parser for the SystemVerilog / Verilog language."""

    language_name = "systemverilog"
    file_extensions = [".sv", ".svh", ".v", ".vh"]

    @staticmethod
    def _get_ts_language() -> Any:
        return tree_sitter_systemverilog.language()

    @staticmethod
    def _get_function_query() -> str:
        """Query to capture function/task/always/initial/module constructs.

        The grammar uses:
        - ``function_declaration`` for function definitions
        - ``task_declaration`` for task definitions
        - ``always_construct`` for always blocks
        - ``initial_construct`` for initial blocks
        - ``module_declaration`` for module definitions
        """
        return """
        [
            (function_declaration) @func
            (task_declaration) @func
            (always_construct) @func
            (initial_construct) @func
            (module_declaration) @func
        ]
        """

    def _extract_function(self, node: Node, code: bytes) -> Function:
        func = Function()
        func.start = node.start_point[0] + 1
        func.end = node.end_point[0] + 1

        if node.type == "module_declaration":
            self._extract_module(func, node, code)
        elif node.type == "function_declaration":
            self._extract_func_or_task(func, node, code, kind="function")
        elif node.type == "task_declaration":
            self._extract_func_or_task(func, node, code, kind="task")
        elif node.type == "always_construct":
            self._extract_always(func, node, code)
        elif node.type == "initial_construct":
            self._extract_initial(func, node, code)

        return func

    # -- extraction helpers --------------------------------------------------

    def _extract_module(self, func: Function, node: Node, code: bytes) -> None:
        """Extract module name from module_ansi_header or module_nonansi_header."""
        for child in node.children:
            if child.type in ("module_ansi_header", "module_nonansi_header"):
                name_node = child.child_by_field_name("name")
                if name_node:
                    func.name = self._node_text(name_node, code)
                    func.line = name_node.start_point[0] + 1
                # Extract port list as signature
                self._extract_module_ports(func, child, code)
                break

    def _extract_module_ports(
        self, func: Function, header_node: Node, code: bytes
    ) -> None:
        """Extract port list from module header as signature."""
        for child in header_node.children:
            if child.type in (
                "list_of_port_declarations",
                "list_of_ports",
            ):
                func.signature = self._node_text(child, code)
                # Normalise multi-line signatures
                if "\n" in func.signature:
                    func.signature = " ".join(func.signature.split())
                break

    def _extract_func_or_task(
        self, func: Function, node: Node, code: bytes, *, kind: str
    ) -> None:
        """Extract function or task name and signature.

        Grammar structure:
        - function_declaration -> function_body_declaration (has name field)
        - task_declaration -> task_body_declaration (has name field)
        """
        body_type = f"{kind}_body_declaration"
        for child in node.children:
            if child.type == body_type:
                name_node = child.child_by_field_name("name")
                if name_node:
                    func.name = self._node_text(name_node, code)
                    func.line = name_node.start_point[0] + 1
                # Extract return type for functions
                if kind == "function":
                    self._extract_function_return_type(func, child, code)
                # Extract port list as signature
                self._extract_tf_ports(func, child, code)
                break
        # Fallback: try to find name in function_prototype or direct children
        if not func.name:
            self._extract_name_fallback(func, node, code)

    def _extract_function_return_type(
        self, func: Function, body_node: Node, code: bytes
    ) -> None:
        """Extract function return type from data_type_or_void."""
        for child in body_node.children:
            if child.type == "data_type_or_void":
                func.return_type = self._node_text(child, code)
                break

    def _extract_tf_ports(
        self, func: Function, body_node: Node, code: bytes
    ) -> None:
        """Extract task/function port list as signature."""
        for child in body_node.children:
            if child.type in (
                "tf_port_list",
                "list_of_tf_variable_identifiers",
            ):
                func.signature = self._node_text(child, code)
                if "\n" in func.signature:
                    func.signature = " ".join(func.signature.split())
                break

    def _extract_always(self, func: Function, node: Node, code: bytes) -> None:
        """Extract always block; use always keyword + line number as name."""
        keyword = "always"
        for child in node.children:
            if child.type == "always_keyword":
                keyword = self._node_text(child, code)
                func.line = child.start_point[0] + 1
                break
        func.name = f"{keyword}@{func.start}"

    def _extract_initial(self, func: Function, node: Node, code: bytes) -> None:
        """Extract initial block; use ``initial@<line>`` as name."""
        func.name = f"initial@{func.start}"
        func.line = node.start_point[0] + 1

    def _extract_name_fallback(
        self, func: Function, node: Node, code: bytes
    ) -> None:
        """Fallback: search for simple_identifier children as function/task name."""
        for child in node.children:
            if child.type == "simple_identifier":
                func.name = self._node_text(child, code)
                func.line = child.start_point[0] + 1
                return
            # Recurse one level into prototype nodes
            if child.type in ("function_prototype", "task_prototype"):
                name_node = child.child_by_field_name("name")
                if name_node:
                    func.name = self._node_text(name_node, code)
                    func.line = name_node.start_point[0] + 1
                    return

    @staticmethod
    def _node_text(node: Node, code: bytes) -> str:
        return code[node.start_byte:node.end_byte].decode("utf-8")

    # -- cyclomatic complexity -----------------------------------------------

    def _get_decision_node_types(self) -> set[str]:
        """Traditional CCM: each case_item counts as a separate decision."""
        return {
            "conditional_statement",    # if/else
            "case_item",                # each case branch
            "loop_statement",           # for, while, do-while, forever, repeat, foreach
            "conditional_expression",   # ternary ?:
            "&&",
            "||",
        }

    def _get_modified_decision_node_types(self) -> set[str]:
        """Modified CCM: entire case_statement counts as one decision."""
        return {
            "conditional_statement",
            "case_statement",
            "loop_statement",
            "conditional_expression",
            "&&",
            "||",
        }

    def _get_statement_node_types(self) -> set[str]:
        return {
            "blocking_assignment",
            "nonblocking_assignment",
            "conditional_statement",
            "case_statement",
            "loop_statement",
            "jump_statement",           # return, break, continue
            "disable_statement",
            "event_trigger",
            "procedural_timing_control_statement",
            "wait_statement",
            "procedural_assertion_statement",
            "subroutine_call_statement",
            "seq_block",                # begin...end
        }

    # -- FQDN ---------------------------------------------------------------

    def _get_scope_name(self, node: Node) -> str | None:
        """Extract scope name from module, class, or package ancestor nodes."""
        if node.type == "module_declaration":
            return self._find_header_name(node)
        if node.type in ("class_declaration", "package_declaration"):
            name_node = node.child_by_field_name("name")
            if name_node and name_node.text:
                return name_node.text.decode("utf-8")
        return None

    def _find_header_name(self, module_node: Node) -> str | None:
        """Find module name from its header child."""
        for child in module_node.children:
            if child.type in ("module_ansi_header", "module_nonansi_header"):
                name_node = child.child_by_field_name("name")
                if name_node and name_node.text:
                    return name_node.text.decode("utf-8")
        return None

    def _compute_fqdn(self, node: Node, name: str) -> str:
        """Compute FQDN with ``::`` separator (e.g. ``counter::add``)."""
        parts: list[str] = []
        current = node.parent
        while current is not None:
            part = self._get_scope_name(current)
            if part:
                parts.append(part)
            current = current.parent
        parts.reverse()
        parts.append(name)
        return "::".join(parts) if len(parts) > 1 else name

    # -- filter module itself when nested ------------------------------------

    def _query_function_nodes(self, root: Node) -> list[Node]:
        """Return function nodes, filtering out top-level module_declarations
        that are parents of other captured nodes (to avoid double-counting)."""
        captures = self.capture_nodes(root, self._get_function_query())
        return captures.get("func", [])
