"""Statement mapping (spec §4.2): maps diff line numbers to AST statement
nodes using tree-sitter C, then computes the graded localization targets
L1 (exact patch statements) and L2 (L1 + backward def-use slice, via
src/cpg/defuse.py). L3 (whole function) is computed trivially by callers.
"""
from __future__ import annotations

from dataclasses import dataclass

import yaml
from tree_sitter import Language, Node, Parser
import tree_sitter_c as tsc

_STATEMENT_TYPES_PATH = "configs/cpg/statement_types.yml"
_LANGUAGE = Language(tsc.language())


def _load_statement_types(path: str = _STATEMENT_TYPES_PATH) -> set[str]:
    with open(path) as f:
        return set(yaml.safe_load(f)["statement_types"])


STATEMENT_TYPES = _load_statement_types()


@dataclass
class StatementNode:
    # (start_byte, end_byte): a node's start_byte ALONE is not unique — a node
    # whose first child starts at the same byte (e.g. `declaration`'s
    # `primitive_type` child) collides with it. end_byte disambiguates.
    node_id: tuple[int, int]
    node_type: str
    start_line: int  # 1-based, inclusive
    end_line: int  # 1-based, inclusive


def parse_function(code: str) -> Node:
    parser = Parser(_LANGUAGE)
    tree = parser.parse(code.encode("utf-8"))
    return tree.root_node


def _iter_nodes(node: Node):
    yield node
    for child in node.children:
        yield from _iter_nodes(child)


def all_statement_nodes(root: Node, statement_types: set[str] = STATEMENT_TYPES) -> list[StatementNode]:
    out = []
    for node in _iter_nodes(root):
        if node.type in statement_types:
            out.append(StatementNode(
                node_id=(node.start_byte, node.end_byte),
                node_type=node.type,
                start_line=node.start_point[0] + 1,
                end_line=node.end_point[0] + 1,
            ))
    return out


def smallest_statement_for_line(
    root: Node, line: int, statement_types: set[str] = STATEMENT_TYPES
) -> StatementNode | None:
    """Smallest AST node whose span covers `line` (1-based) and whose type is
    in the statement allowlist (spec §4.2 step 2)."""
    best: Node | None = None
    for node in _iter_nodes(root):
        if node.type not in statement_types:
            continue
        start_line = node.start_point[0] + 1
        end_line = node.end_point[0] + 1
        if start_line <= line <= end_line:
            if best is None or (end_line - start_line) < (best.end_point[0] - best.start_point[0]):
                best = node
    if best is None:
        return None
    return StatementNode(
        node_id=(best.start_byte, best.end_byte),
        node_type=best.type,
        start_line=best.start_point[0] + 1,
        end_line=best.end_point[0] + 1,
    )


def statement_for_line_with_fallback(
    root: Node, line: int, statement_types: set[str] = STATEMENT_TYPES, sibling_window: int = 2
) -> tuple[StatementNode | None, str]:
    """spec amendment §2 (L1 coverage fix): if `line` maps to no statement
    node directly, fall back to the nearest statement (by start/end-line
    distance) anywhere in the function within +/-`sibling_window` lines.
    Returns (StatementNode | None, reason) so callers can log why a line
    resolved the way it did (spec: "Log per-reason rejection counts").

    **Note on the spec amendment's "(1) expand upward to parent; (2) nearest
    sibling" two-step**: step (1) is provably redundant with the "exact"
    search and was removed after implementing and testing it. Proof: a node
    N's ancestors all share N's line-coverage by construction (tree nesting
    is monotonic — if a child's span covers `line`, every ancestor's span
    covers it too). `smallest_statement_for_line` already scans *every* node
    in the tree for a type-and-coverage match, not just one path from root to
    a single "deepest" node — so if any ancestor of any node had an allowed
    type, that same node would already have been found by the exact search.
    Confirmed empirically against three constructed cases (a case-label line,
    a blank line inside an if-block, and a deliberately narrowed allowlist)
    where a hand-written ancestor-walk always returned exactly what "exact"
    already found, never something new. Logged in DECISIONS.md rather than
    shipping dead code silently."""
    exact = smallest_statement_for_line(root, line, statement_types)
    if exact is not None:
        return exact, "exact"

    all_stmts = all_statement_nodes(root, statement_types)
    nearest: StatementNode | None = None
    nearest_dist = sibling_window + 1
    for stmt in all_stmts:
        dist = min(abs(stmt.start_line - line), abs(stmt.end_line - line))
        if dist <= sibling_window and dist < nearest_dist:
            nearest = stmt
            nearest_dist = dist
    if nearest is not None:
        return nearest, "nearest_sibling_within_window"

    return None, "no_statement_found"


def map_patch_lines_to_statements(code: str, patch_lines: list[int]) -> list[StatementNode]:
    """spec §4.2 step 3: output per sample = L1 = set of statement node ids
    + their line ranges, for the given patch diff lines. Uses the
    exact->expand-to-parent->nearest-sibling fallback chain (spec amendment
    §2); use `map_patch_lines_to_statements_detailed` to also see per-line
    rejection reasons."""
    return [stmt for stmt, _reason in map_patch_lines_to_statements_detailed(code, patch_lines) if stmt is not None]


def map_patch_lines_to_statements_detailed(
    code: str, patch_lines: list[int]
) -> list[tuple[StatementNode | None, str]]:
    """Same as `map_patch_lines_to_statements` but returns one
    (StatementNode | None, reason) pair per input line (not deduped), so
    callers can log exactly which lines fell back and why."""
    root = parse_function(code)
    seen_ids: set[tuple[int, int]] = set()
    out: list[tuple[StatementNode | None, str]] = []
    for line in patch_lines:
        stmt, reason = statement_for_line_with_fallback(root, line)
        if stmt is not None and stmt.node_id in seen_ids:
            continue  # already emitted this statement for an earlier line in the same call
        if stmt is not None:
            seen_ids.add(stmt.node_id)
        out.append((stmt, reason))
    return out


def statement_lines(statements: list[StatementNode]) -> set[int]:
    lines: set[int] = set()
    for s in statements:
        lines.update(range(s.start_line, s.end_line + 1))
    return lines