"""Extract RETURN clause columns from Cypher statements.

This is used to build the SQL AS clause for the ag_catalog.cypher() wrapper.
AGE requires knowing the output columns and their types at SQL level.
"""

from __future__ import annotations

import re

# Match the RETURN clause, accounting for DISTINCT and ending at ORDER BY, SKIP, LIMIT, or end.
_RETURN_RE = re.compile(
    r"\bRETURN\b\s+(?:DISTINCT\s+)?(.*?)(?:\bORDER\s+BY\b|\bSKIP\b|\bLIMIT\b|;|\Z)",
    re.IGNORECASE | re.DOTALL,
)

# Match function calls like count(...), collect(...), etc.
_FUNC_RE = re.compile(r"^(\w+)\s*\(")


def extract_return_columns(cypher: str) -> list[str]:
    """Extract column names from a Cypher RETURN clause.

    Parses the RETURN clause to determine what column names should appear
    in the SQL AS clause. Each column becomes ``name agtype``.

    Rules:
    - ``RETURN n`` → ``["n"]``
    - ``RETURN n.name, n.age`` → ``["name", "age"]``  (last dotted segment)
    - ``RETURN n.name AS person_name`` → ``["person_name"]``
    - ``RETURN count(n)`` → ``["count"]``
    - ``RETURN count(n) AS total`` → ``["total"]``
    - ``RETURN *`` → ``["v"]``  (fallback single column)

    Returns:
        List of column name strings.

    Raises:
        ValueError: If no RETURN clause is found.
    """
    m = _RETURN_RE.search(cypher)
    if not m:
        # Write-only queries (CREATE, SET, DELETE, MERGE) have no RETURN.
        # Use a single dummy column that AGE expects.
        return ["v"]

    body = m.group(1).strip()
    if body == "*":
        return ["v"]

    columns: list[str] = []
    # Split on commas, respecting parentheses depth.
    parts = _split_columns(body)

    for part in parts:
        part = part.strip()
        if not part:
            continue

        # Check for AS alias
        as_match = re.search(r"\bAS\s+(\w+)\s*$", part, re.IGNORECASE)
        if as_match:
            columns.append(as_match.group(1))
            continue

        # Check for function call: count(n) → "count"
        func_match = _FUNC_RE.match(part)
        if func_match:
            columns.append(func_match.group(1))
            continue

        # Dotted property: n.name → "name"
        if "." in part:
            columns.append(part.rsplit(".", 1)[1].strip())
            continue

        # Plain identifier
        columns.append(part.strip())

    return columns


def _split_columns(text: str) -> list[str]:
    """Split comma-separated column expressions, respecting parentheses."""
    parts: list[str] = []
    depth = 0
    current: list[str] = []

    for ch in text:
        if ch == "(":
            depth += 1
            current.append(ch)
        elif ch == ")":
            depth -= 1
            current.append(ch)
        elif ch == "," and depth == 0:
            parts.append("".join(current))
            current = []
        else:
            current.append(ch)

    if current:
        parts.append("".join(current))

    return parts
