# SPDX-FileCopyrightText: 2026 Bart Broere
# SPDX-License-Identifier: GPL-3.0-or-later
"""Light-weight Cypher text analysis needed to run queries through AGE.

AGE's ``cypher()`` SQL function needs the caller to declare the result columns
up front, does not accept Neo4j-style parameter maps reliably, and lacks a few
constructs neo4j-browser relies on. This module works on the query text only:
it masks out strings and comments, finds the final top-level RETURN clause,
and inlines parameters as Cypher literals.
"""

import math
import re
from dataclasses import dataclass


class CypherError(Exception):
    def __init__(self, code, message):
        super().__init__(message)
        self.code = code
        self.message = message


@dataclass
class Scan:
    text: str
    masked: str  # strings/backticks replaced by 'x', comments by spaces
    depth: list  # bracket nesting depth for every character


def scan(text: str) -> Scan:
    masked = list(text)
    depth = [0] * len(text)
    level = 0
    i, n = 0, len(text)
    while i < n:
        c = text[i]
        if c in "'\"`":
            j = i + 1
            while j < n and text[j] != c:
                j += 2 if text[j] == "\\" and c != "`" else 1
            for k in range(i, min(j + 1, n)):
                masked[k] = "x"
                depth[k] = level
            i = j + 1
            continue
        if text.startswith("//", i):
            j = text.find("\n", i)
            j = n if j == -1 else j
            for k in range(i, j):
                masked[k] = " "
                depth[k] = level
            i = j
            continue
        if text.startswith("/*", i):
            j = text.find("*/", i + 2)
            j = n if j == -1 else j + 2
            for k in range(i, j):
                masked[k] = " "
                depth[k] = level
            i = j
            continue
        if c in "([{":
            depth[i] = level
            level += 1
        elif c in ")]}":
            level = max(level - 1, 0)
            depth[i] = level
        else:
            depth[i] = level
        i += 1
    return Scan(text, "".join(masked), depth)


def _top_level(s: Scan, pattern, start=0, end=None):
    end = len(s.text) if end is None else end
    for m in re.finditer(pattern, s.masked[start:end], re.IGNORECASE):
        if s.depth[start + m.start()] == 0:
            yield start + m.start(), start + m.end()


def strip_statement(text: str) -> str:
    text = text.strip()
    while text.endswith(";"):
        text = text[:-1].rstrip()
    return text


def return_columns(text: str) -> list:
    """Column names of the final top-level RETURN clause ([] if there is none)."""
    s = scan(text)
    returns = list(_top_level(s, r"(?<![\w.$])RETURN(?!\w)"))
    if not returns:
        return []
    clause_start = returns[-1][1]
    clause_end = len(text)
    for start, _ in _top_level(s, r"(?<![\w.$])(ORDER\s+BY|SKIP|LIMIT|UNION)(?!\w)", clause_start):
        clause_end = start
        break
    distinct = re.match(r"\s*DISTINCT(?!\w)", s.masked[clause_start:clause_end], re.IGNORECASE)
    if distinct:
        clause_start += distinct.end()

    items, item_start = [], clause_start
    for comma, _ in _top_level(s, ",", clause_start, clause_end):
        items.append((item_start, comma))
        item_start = comma + 1
    items.append((item_start, clause_end))

    columns = []
    for start, end in items:
        expr = text[start:end].strip()
        if expr == "*":
            raise CypherError("Neo.ClientError.Statement.FeatureNotSupported",
                              "RETURN * is not supported by peo; list the returned variables explicitly.")
        alias = None
        for a_start, a_end in _top_level(s, r"(?<![\w.$])AS\s+(`[^`]*`|\w+)\s*$", start, end):
            alias = re.match(r"AS\s+(.*)", text[a_start:a_end].strip(), re.IGNORECASE | re.DOTALL).group(1)
        if alias is not None:
            name = alias.strip()
            if name.startswith("`") and name.endswith("`"):
                name = name[1:-1].replace("``", "`")
            columns.append(name)
        else:
            columns.append(expr)
    return columns


_PARAM = re.compile(r"\$(`(?:[^`]|``)+`|[A-Za-z_][A-Za-z0-9_]*|\d+)")


def inline_parameters(text: str, params: dict) -> str:
    """Replace ``$name`` references with literal Cypher values."""
    if "$" not in text:
        return text
    s = scan(text)
    pieces, last, missing = [], 0, []
    for m in _PARAM.finditer(s.masked):
        name = m.group(1)
        if name.startswith("`"):
            name = name[1:-1].replace("``", "`")
        if name not in params:
            missing.append(name)
            continue
        pieces.append(text[last:m.start()])
        pieces.append(literal(params[name]))
        last = m.end()
    if missing:
        raise CypherError("Neo.ClientError.Statement.ParameterMissing",
                          "Expected parameter(s): " + ", ".join(dict.fromkeys(missing)))
    pieces.append(text[last:])
    return "".join(pieces)


def _escape_string(v: str) -> str:
    out = []
    for ch in v:
        if ch == "\\":
            out.append("\\\\")
        elif ch == "'":
            out.append("\\'")
        elif ch == "\n":
            out.append("\\n")
        elif ch == "\r":
            out.append("\\r")
        elif ch == "\t":
            out.append("\\t")
        elif ord(ch) < 0x20:
            out.append(f"\\u{ord(ch):04x}")
        else:
            out.append(ch)
    return "'" + "".join(out) + "'"


def literal(v) -> str:
    if v is None:
        return "null"
    if v is True:
        return "true"
    if v is False:
        return "false"
    if isinstance(v, int):
        return str(v)
    if isinstance(v, float):
        if math.isnan(v):
            return "'NaN'::float"
        if math.isinf(v):
            return "'Infinity'::float" if v > 0 else "'-Infinity'::float"
        return repr(v)
    if isinstance(v, str):
        return _escape_string(v)
    if isinstance(v, (list, tuple)):
        return "[" + ", ".join(literal(x) for x in v) + "]"
    if isinstance(v, dict):
        return "{" + ", ".join(
            "`" + str(k).replace("`", "``") + "`: " + literal(x) for k, x in v.items()) + "}"
    raise CypherError("Neo.ClientError.Statement.TypeError",
                      f"Parameters of type {type(v).__name__} are not supported by peo")


# neo4j-browser's "expand node" query uses a pattern comprehension, which AGE
# does not support. Rewrite it into an equivalent OPTIONAL MATCH + count.
_NEIGHBOUR_COUNT = re.compile(
    r"WITH\s+(\w+)\s*,\s*size\(\s*\[\s*\(\s*\1\s*\)\s*--\s*\(\s*\)\s*\|\s*1\s*\]\s*\)\s+AS\s+(\w+)",
    re.IGNORECASE)


def rewrite_for_age(text: str) -> str:
    return _NEIGHBOUR_COUNT.sub(
        lambda m: f"OPTIONAL MATCH ({m.group(1)})-[peo_r]-() "
                  f"WITH {m.group(1)}, count(peo_r) AS {m.group(2)}",
        text)


# AGE fails with "not a common type" on `id(x) IN [a, b, ...]` when the list has
# more than one element, which neo4j-browser uses to fetch relationships between
# displayed nodes. Expand such lists of integer literals into an OR chain.
_ID_IN_LIST = re.compile(r"(?<![\w.])(id\(\s*\w+\s*\))\s+IN\s+\[([\s\d,+\-]*)\]", re.IGNORECASE)


def expand_id_in_lists(text: str) -> str:
    s = scan(text)
    pieces, last = [], 0
    for m in _ID_IN_LIST.finditer(s.masked):
        values = [v.strip() for v in m.group(2).split(",") if v.strip()]
        if len(values) == 1:
            continue
        expr = text[m.start(1):m.end(1)]
        pieces.append(text[last:m.start()])
        pieces.append("(" + " OR ".join(f"{expr} = {v}" for v in values) + ")" if values else "false")
        last = m.end()
    pieces.append(text[last:])
    return "".join(pieces)


def first_keyword(text: str) -> str:
    m = re.match(r"\s*(\w+)", scan(text).masked)
    return m.group(1).upper() if m else ""