"""Data-flow graph (DFG) extraction for GraphCodeBERT.

This is a faithful port of Microsoft's GraphCodeBERT ``parser/`` package
(``utils.py`` + the ``DFG_python`` extractor from ``DFG.py``), adapted to the
modern ``py-tree-sitter`` API (>= 0.22, ``Language(tree_sitter_python.language())``)
instead of the original hand-compiled ``my-languages.so``.

The dataset used in this project contains **Python** snippets (verified in
``preprocess.py``), so only the Python extractor is ported; adding another
language means adding its ``DFG_<lang>`` function and grammar package here.

A DFG entry is the 5-tuple used throughout GraphCodeBERT::

    (variable_name, token_index, edge_type, source_variable_names, source_token_indices)

``edge_type`` is ``"comesFrom"`` (value flows from a previous definition) or
``"computedFrom"`` (value is computed from the right-hand side of an assignment).
"""

from __future__ import annotations

import io
import re
import sys
import tokenize
from typing import Any

from tree_sitter import Language, Node, Parser

__all__ = [
    "get_parser",
    "extract_dataflow",
    "remove_comments_and_docstrings",
    "DataFlowExtractionError",
]

#: tree-sitter recursion is mirrored by the recursive Python walkers below.
#: Competitive-programming snippets can nest deeply, so raise the ceiling but
#: keep it bounded so a pathological file raises RecursionError instead of
#: segfaulting the worker.
_RECURSION_LIMIT = 10_000


class DataFlowExtractionError(RuntimeError):
    """Raised when a snippet cannot be turned into code tokens at all."""


_PARSER_CACHE: dict[str, Parser] = {}


def get_parser(language: str = "python") -> Parser:
    """Return a cached tree-sitter parser for ``language``.

    Cached per process so that ``datasets.map(num_proc=...)`` workers each build
    the parser once rather than once per snippet.
    """
    if language in _PARSER_CACHE:
        return _PARSER_CACHE[language]
    if language != "python":
        raise ValueError(
            f"Only the Python grammar is wired up (requested {language!r}). "
            "Add the matching tree_sitter_<lang> package and DFG_<lang> function."
        )
    try:
        import tree_sitter_python
    except ImportError as exc:  # pragma: no cover - environment problem
        raise ImportError(
            "tree_sitter_python is required for GraphCodeBERT data-flow extraction. "
            "Install it with `pip install tree-sitter tree-sitter-python`."
        ) from exc
    parser = Parser(Language(tree_sitter_python.language()))
    _PARSER_CACHE[language] = parser
    return parser


# --------------------------------------------------------------------------- #
# parser/utils.py
# --------------------------------------------------------------------------- #
def remove_comments_and_docstrings(source: str, lang: str = "python") -> str:
    """Strip comments and docstrings, preserving token columns.

    Column positions are preserved because the DFG indices are ``(row, column)``
    points into the *cleaned* source.
    """
    if lang == "python":
        io_obj = io.StringIO(source)
        out = ""
        prev_toktype = tokenize.INDENT
        last_lineno = -1
        last_col = 0
        for tok in tokenize.generate_tokens(io_obj.readline):
            token_type, token_string = tok[0], tok[1]
            start_line, start_col = tok[2]
            end_line, end_col = tok[3]
            if start_line > last_lineno:
                last_col = 0
            if start_col > last_col:
                out += " " * (start_col - last_col)
            if token_type == tokenize.COMMENT:
                pass
            elif token_type == tokenize.STRING:
                # A string that starts a logical line is a docstring -> drop it.
                if prev_toktype != tokenize.INDENT and prev_toktype != tokenize.NEWLINE:
                    if start_col > 0:
                        out += token_string
            else:
                out += token_string
            prev_toktype = token_type
            last_col = end_col
            last_lineno = end_line
        return "\n".join(x for x in out.split("\n") if x.strip() != "")

    def _replacer(match: re.Match[str]) -> str:
        s = match.group(0)
        return " " if s.startswith("/") else s

    pattern = re.compile(
        r"//.*?$|/\*.*?\*/|\'(?:\\.|[^\\\'])*\'|\"(?:\\.|[^\\\"])*\"",
        re.DOTALL | re.MULTILINE,
    )
    cleaned = re.sub(pattern, _replacer, source)
    return "\n".join(x for x in cleaned.split("\n") if x.strip() != "")


def tree_to_token_index(root_node: Node) -> list[tuple[Any, Any]]:
    """Collect ``(start_point, end_point)`` spans of every leaf token."""
    if (len(root_node.children) == 0 or root_node.type == "string") and root_node.type != "comment":
        return [(root_node.start_point, root_node.end_point)]
    spans: list[tuple[Any, Any]] = []
    for child in root_node.children:
        spans += tree_to_token_index(child)
    return spans


def tree_to_variable_index(root_node: Node, index_to_code: dict) -> list[tuple[Any, Any]]:
    """Collect spans of leaves that are *variables* (token text != node type)."""
    if (len(root_node.children) == 0 or root_node.type == "string") and root_node.type != "comment":
        index = (root_node.start_point, root_node.end_point)
        _, code = index_to_code[index]
        return [] if root_node.type == code else [index]
    spans: list[tuple[Any, Any]] = []
    for child in root_node.children:
        spans += tree_to_variable_index(child, index_to_code)
    return spans


def index_to_code_token(index: tuple[Any, Any], code: list[str]) -> str:
    """Slice the source text covered by a ``(start_point, end_point)`` span."""
    start_point, end_point = index
    if start_point[0] == end_point[0]:
        return code[start_point[0]][start_point[1] : end_point[1]]
    s = code[start_point[0]][start_point[1] :]
    for i in range(start_point[0] + 1, end_point[0]):
        s += code[i]
    s += code[end_point[0]][: end_point[1]]
    return s


# --------------------------------------------------------------------------- #
# parser/DFG.py :: DFG_python
# --------------------------------------------------------------------------- #
_ASSIGNMENT = ("assignment", "augmented_assignment", "for_in_clause")
_IF_STATEMENT = ("if_statement",)
_FOR_STATEMENT = ("for_statement",)
_WHILE_STATEMENT = ("while_statement",)
_DO_FIRST_STATEMENT = ("for_in_clause",)
_DEF_STATEMENT = ("default_parameter",)


def DFG_python(root_node: Node, index_to_code: dict, states: dict) -> tuple[list, dict]:
    """Build the data-flow graph of a Python AST subtree.

    Returns ``(dfg_edges, variable_states)`` where ``variable_states`` maps a
    variable name to the token indices that currently define it.
    """
    states = states.copy()

    if (len(root_node.children) == 0 or root_node.type == "string") and root_node.type != "comment":
        idx, code = index_to_code[(root_node.start_point, root_node.end_point)]
        if root_node.type == code:  # a keyword/operator, not a variable
            return [], states
        if code in states:
            return [(code, idx, "comesFrom", [code], states[code].copy())], states
        if root_node.type == "identifier":
            states[code] = [idx]
        return [(code, idx, "comesFrom", [], [])], states

    if root_node.type in _DEF_STATEMENT:
        name = root_node.child_by_field_name("name")
        value = root_node.child_by_field_name("value")
        dfg: list = []
        if value is None:
            for index in tree_to_variable_index(name, index_to_code):
                idx, code = index_to_code[index]
                dfg.append((code, idx, "comesFrom", [], []))
                states[code] = [idx]
            return sorted(dfg, key=lambda x: x[1]), states
        name_indexs = tree_to_variable_index(name, index_to_code)
        value_indexs = tree_to_variable_index(value, index_to_code)
        temp, states = DFG_python(value, index_to_code, states)
        dfg += temp
        for index1 in name_indexs:
            idx1, code1 = index_to_code[index1]
            for index2 in value_indexs:
                idx2, code2 = index_to_code[index2]
                dfg.append((code1, idx1, "comesFrom", [code2], [idx2]))
            states[code1] = [idx1]
        return sorted(dfg, key=lambda x: x[1]), states

    if root_node.type in _ASSIGNMENT:
        if root_node.type == "for_in_clause":
            right_nodes = [root_node.children[-1]]
            left_nodes = [root_node.child_by_field_name("left")]
        else:
            if root_node.child_by_field_name("right") is None:
                return [], states
            left_nodes = [x for x in root_node.child_by_field_name("left").children if x.type != ","]
            right_nodes = [
                x for x in root_node.child_by_field_name("right").children if x.type != ","
            ]
            if len(right_nodes) != len(left_nodes):
                left_nodes = [root_node.child_by_field_name("left")]
                right_nodes = [root_node.child_by_field_name("right")]
            if len(left_nodes) == 0:
                left_nodes = [root_node.child_by_field_name("left")]
            if len(right_nodes) == 0:
                right_nodes = [root_node.child_by_field_name("right")]
        dfg = []
        for node in right_nodes:
            temp, states = DFG_python(node, index_to_code, states)
            dfg += temp
        for left_node, right_node in zip(left_nodes, right_nodes):
            left_tokens_index = tree_to_variable_index(left_node, index_to_code)
            right_tokens_index = tree_to_variable_index(right_node, index_to_code)
            for token1_index in left_tokens_index:
                idx1, code1 = index_to_code[token1_index]
                dfg.append(
                    (
                        code1,
                        idx1,
                        "computedFrom",
                        [index_to_code[x][1] for x in right_tokens_index],
                        [index_to_code[x][0] for x in right_tokens_index],
                    )
                )
                states[code1] = [idx1]
        return sorted(dfg, key=lambda x: x[1]), states

    if root_node.type in _IF_STATEMENT:
        dfg = []
        current_states = states.copy()
        others_states = []
        tag = "else" in root_node.type
        for child in root_node.children:
            if "else" in child.type:
                tag = True
            if child.type not in ("elif_clause", "else_clause"):
                temp, current_states = DFG_python(child, index_to_code, current_states)
                dfg += temp
            else:
                temp, new_states = DFG_python(child, index_to_code, states)
                dfg += temp
                others_states.append(new_states)
        others_states.append(current_states)
        if tag is False:
            others_states.append(states)
        merged: dict = {}
        for dic in others_states:
            for key in dic:
                merged.setdefault(key, [])
                merged[key] += dic[key]
        for key in merged:
            merged[key] = sorted(set(merged[key]))
        return sorted(dfg, key=lambda x: x[1]), merged

    if root_node.type in _FOR_STATEMENT:
        dfg = []
        # Two passes: loop bodies can consume values defined later in the loop.
        for _ in range(2):
            right_nodes = [x for x in root_node.child_by_field_name("right").children if x.type != ","]
            left_nodes = [x for x in root_node.child_by_field_name("left").children if x.type != ","]
            if len(right_nodes) != len(left_nodes):
                left_nodes = [root_node.child_by_field_name("left")]
                right_nodes = [root_node.child_by_field_name("right")]
            if len(left_nodes) == 0:
                left_nodes = [root_node.child_by_field_name("left")]
            if len(right_nodes) == 0:
                right_nodes = [root_node.child_by_field_name("right")]
            for node in right_nodes:
                temp, states = DFG_python(node, index_to_code, states)
                dfg += temp
            for left_node, right_node in zip(left_nodes, right_nodes):
                left_tokens_index = tree_to_variable_index(left_node, index_to_code)
                right_tokens_index = tree_to_variable_index(right_node, index_to_code)
                for token1_index in left_tokens_index:
                    idx1, code1 = index_to_code[token1_index]
                    dfg.append(
                        (
                            code1,
                            idx1,
                            "computedFrom",
                            [index_to_code[x][1] for x in right_tokens_index],
                            [index_to_code[x][0] for x in right_tokens_index],
                        )
                    )
                    states[code1] = [idx1]
            if root_node.children[-1].type == "block":
                temp, states = DFG_python(root_node.children[-1], index_to_code, states)
                dfg += temp
        return _merge_duplicate_edges(dfg), states

    if root_node.type in _WHILE_STATEMENT:
        dfg = []
        for _ in range(2):
            for child in root_node.children:
                temp, states = DFG_python(child, index_to_code, states)
                dfg += temp
        return _merge_duplicate_edges(dfg), states

    dfg = []
    for child in root_node.children:
        if child.type in _DO_FIRST_STATEMENT:
            temp, states = DFG_python(child, index_to_code, states)
            dfg += temp
    for child in root_node.children:
        if child.type not in _DO_FIRST_STATEMENT:
            temp, states = DFG_python(child, index_to_code, states)
            dfg += temp
    return sorted(dfg, key=lambda x: x[1]), states


def _merge_duplicate_edges(dfg: list) -> list:
    """Collapse the duplicate edges produced by the two-pass loop handling."""
    dic: dict = {}
    for x in dfg:
        key = (x[0], x[1], x[2])
        if key not in dic:
            dic[key] = [x[3], x[4]]
        else:
            dic[key][0] = list(set(dic[key][0] + x[3]))
            dic[key][1] = sorted(set(dic[key][1] + x[4]))
    merged = [(k[0], k[1], k[2], v[0], v[1]) for k, v in sorted(dic.items(), key=lambda t: t[0][1])]
    return sorted(merged, key=lambda x: x[1])


# --------------------------------------------------------------------------- #
# Public entry point (GraphCodeBERT's `extract_dataflow`)
# --------------------------------------------------------------------------- #
def extract_dataflow(code: str, language: str = "python") -> tuple[list[str], list, dict]:
    """Tokenise ``code`` and extract its data-flow graph.

    Returns ``(code_tokens, dfg, status)``. ``status`` records *why* a stage
    degraded so callers can report it instead of hiding it:

    ``comment_strip`` : ``"ok"`` | ``"failed"``
    ``parse``         : ``"ok"`` | ``"failed"``
    ``dfg``           : ``"ok"`` | ``"failed"`` | ``"recursion_limit"``
    ``error``         : ``None`` or ``"<ExcType>: <message>"``

    A degraded DFG yields an **empty** data-flow component -- the snippet is
    still trained on (GraphCodeBERT tolerates zero nodes), it is never dropped.
    """
    status: dict[str, Any] = {"comment_strip": "ok", "parse": "ok", "dfg": "ok", "error": None}

    try:
        cleaned = remove_comments_and_docstrings(code, language)
    except Exception as exc:
        # Syntactically broken snippets are common in the wild; fall back to the
        # raw source rather than discarding the example.
        status["comment_strip"] = "failed"
        status["error"] = f"{type(exc).__name__}: {exc}"
        cleaned = code

    parser = get_parser(language)
    try:
        tree = parser.parse(bytes(cleaned, "utf8"))
        root_node = tree.root_node
    except Exception as exc:
        raise DataFlowExtractionError(f"tree-sitter failed to parse snippet: {exc}") from exc

    old_limit = sys.getrecursionlimit()
    sys.setrecursionlimit(_RECURSION_LIMIT)
    try:
        try:
            tokens_index = tree_to_token_index(root_node)
        except RecursionError as exc:
            status["parse"] = "failed"
            status["dfg"] = "recursion_limit"
            status["error"] = f"{type(exc).__name__}: token index recursion limit"
            raise DataFlowExtractionError("snippet nests deeper than the recursion limit") from exc

        lines = cleaned.split("\n")
        code_tokens = [index_to_code_token(x, lines) for x in tokens_index]
        index_to_code = {
            index: (idx, token) for idx, (index, token) in enumerate(zip(tokens_index, code_tokens))
        }

        try:
            dfg, _ = DFG_python(root_node, index_to_code, {})
        except RecursionError as exc:
            status["dfg"] = "recursion_limit"
            status["error"] = f"{type(exc).__name__}: DFG recursion limit"
            dfg = []
        except Exception as exc:
            status["dfg"] = "failed"
            status["error"] = f"{type(exc).__name__}: {exc}"
            dfg = []
    finally:
        sys.setrecursionlimit(old_limit)

    # Keep only nodes that participate in at least one edge (GraphCodeBERT does
    # the same: isolated nodes carry no data-flow signal).
    dfg = sorted(dfg, key=lambda x: x[1])
    keep: set[int] = set()
    for d in dfg:
        if len(d[-1]) != 0:
            keep.add(d[1])
        keep.update(d[-1])
    dfg = [d for d in dfg if d[1] in keep]

    return code_tokens, dfg, status