#!/usr/bin/env python3
"""Derive Leo's canonical PostgreSQL read surface from its runtime SQL.

The runtime queries are the source of truth.  This helper extracts canonical
``public`` and ``kb_stage`` relation aliases and their qualified column
references directly from ``cloudsql_memory_tool.py``.  It deliberately rejects
unaliased canonical relations, alias collisions, direct canonical DML, and
unqualified identifiers that are not mechanically known query outputs.
"""

from __future__ import annotations

import argparse
import ast
import json
import re
from collections import defaultdict
from pathlib import Path

CANONICAL_RELATION = re.compile(
    r"\b(?:from|join)\s+(public|kb_stage)\.([a-z_][a-z0-9_]*)",
    re.IGNORECASE,
)
ALIASED_CANONICAL_RELATION = re.compile(
    r"\b(?:from|join)\s+(public|kb_stage)\.([a-z_][a-z0-9_]*)\s+as\s+([a-z_][a-z0-9_]*)",
    re.IGNORECASE,
)
MUTATING_SQL_KEYWORD = re.compile(
    r"\b(?:insert|update|delete|truncate|merge)\b",
    re.IGNORECASE,
)
LARGE_OBJECT_MUTATOR_REFERENCE = re.compile(
    r'(?<![a-z0-9_])(?:(?:"pg_catalog"|pg_catalog)\s*\.\s*)?'
    r'(?:(?:"lo_creat"|lo_creat)|(?:"lo_create"|lo_create)|'
    r'(?:"lo_export"|lo_export)|(?:"lo_from_bytea"|lo_from_bytea)|'
    r'(?:"lo_import"|lo_import)|(?:"lo_open"|lo_open)|'
    r'(?:"lo_put"|lo_put)|(?:"lo_truncate64"|lo_truncate64)|'
    r'(?:"lo_truncate"|lo_truncate)|(?:"lo_unlink"|lo_unlink)|'
    r'(?:"lowrite"|lowrite))\s*\(',
    re.IGNORECASE,
)
PSQL_META_COMMAND = re.compile(r"\\")
CANONICAL_NAMESPACE_REFERENCE = re.compile(r"\b(?:public|kb_stage)\s*\.", re.IGNORECASE)
DYNAMIC_SQL_VALUE = "__teleo_dynamic_sql_value__"
UNREVIEWED_DYNAMIC_SQL = "__teleo_unreviewed_dynamic_sql__"
REVIEWED_SQL_INTERPOLATION_CALLS = frozenset({"sql_array", "sql_integer", "sql_json_array", "sql_literal"})
SQL_IDENTIFIER = re.compile(r"\b[a-z_][a-z0-9_]*\b", re.IGNORECASE)
SQL_KEYWORDS = frozenset(
    """
    all and any array as asc between bigint bool boolean by case cast cross
    current_database current_user desc distinct else end exists false filter
    first following for from full group having ilike in inner interval is join
    json jsonb last lateral left like limit not null nulls offset on or order
    outer over partition preceding range recursive right row rows select some
    symmetric table text then ties time timestamp timestamptz to true union
    unknown using uuid values varchar when where window with zone
    """.split()  # noqa: SIM905 - the compact vocabulary is easier to audit as SQL.
)
NONCANONICAL_QUERY_IDENTIFIERS = frozenset(
    {
        # The status query reads these PostgreSQL metadata columns; they are not
        # part of the canonical application-role contract.
        "extname",
        "pg_extension",
        "table_schema",
        "table_type",
    }
)


class QueryContractError(ValueError):
    """The runtime SQL cannot be safely reduced to an exact read contract."""


def _is_reviewed_sql_interpolation(node: ast.expr) -> bool:
    return (
        isinstance(node, ast.Call)
        and isinstance(node.func, ast.Name)
        and node.func.id in REVIEWED_SQL_INTERPOLATION_CALLS
    )


def _render_sql_expression(node: ast.expr) -> str | None:
    if isinstance(node, ast.Constant) and isinstance(node.value, str):
        return node.value
    if isinstance(node, ast.JoinedStr):
        parts: list[str] = []
        for value in node.values:
            if isinstance(value, ast.Constant) and isinstance(value.value, str):
                parts.append(value.value)
            elif isinstance(value, ast.FormattedValue) and _is_reviewed_sql_interpolation(value.value):
                # Runtime interpolation supplies values, limits, or predicates;
                # canonical relation and column identifiers remain static.
                parts.append(DYNAMIC_SQL_VALUE)
            else:
                parts.append(UNREVIEWED_DYNAMIC_SQL)
        return "".join(parts)
    return None


def extract_canonical_query_literals(source: str) -> tuple[tuple[int, str], ...]:
    """Return canonical SQL literals actually passed to the runtime query runners."""

    tree = ast.parse(source)
    queries: set[tuple[int, str]] = set()
    for function in (node for node in ast.walk(tree) if isinstance(node, (ast.FunctionDef, ast.AsyncFunctionDef))):
        assignments: dict[str, tuple[int, str | None]] = {}
        for node in ast.walk(function):
            if isinstance(node, ast.Assign):
                rendered = _render_sql_expression(node.value)
                for target in node.targets:
                    if isinstance(target, ast.Name):
                        assignments[target.id] = (node.lineno, rendered)
            elif isinstance(node, ast.AnnAssign) and isinstance(node.target, ast.Name):
                assignments[node.target.id] = (node.lineno, _render_sql_expression(node.value))

        for call in (node for node in ast.walk(function) if isinstance(node, ast.Call)):
            runner = call.func.id if isinstance(call.func, ast.Name) else None
            if runner not in {"run_psql", "psql_json_lines"}:
                continue
            if function.name == "psql_json_lines" and runner == "run_psql":
                continue
            if len(call.args) < 2:
                raise QueryContractError(f"line {call.lineno}: {runner} requires a statically inspectable SQL argument")
            expression = call.args[1]
            if isinstance(expression, ast.Name):
                line, sql = assignments.get(expression.id, (call.lineno, None))
            else:
                line, sql = call.lineno, _render_sql_expression(expression)
            if sql is None:
                raise QueryContractError(f"line {call.lineno}: {runner} SQL must remain a local string literal")
            scrubbed = _strip_sql_strings_and_comments(sql)
            if UNREVIEWED_DYNAMIC_SQL in sql:
                raise QueryContractError(
                    f"line {line}: dynamic SQL interpolation is not part of the reviewed runtime query surface"
                )
            if PSQL_META_COMMAND.search(scrubbed):
                raise QueryContractError(
                    f"line {line}: psql meta-commands are not part of the reviewed runtime query surface"
                )
            if MUTATING_SQL_KEYWORD.search(scrubbed) and CANONICAL_NAMESPACE_REFERENCE.search(scrubbed):
                raise QueryContractError(f"line {line}: runtime query performs direct canonical DML")
            if LARGE_OBJECT_MUTATOR_REFERENCE.search(scrubbed):
                raise QueryContractError(
                    f"line {line}: runtime query invokes a provider-owned large-object mutator"
                )
            if CANONICAL_RELATION.search(sql):
                queries.add((line, sql))

    if not queries:
        raise QueryContractError("runtime source contains no canonical SQL query literals")
    return tuple(sorted(queries, key=lambda item: item[0]))


def _strip_sql_strings_and_comments(sql: str) -> str:
    without_strings = re.sub(r"(?:[eE])?'(?:''|[^'])*'", " ", sql)
    without_block_comments = re.sub(r"/\*.*?\*/", " ", without_strings, flags=re.DOTALL)
    return re.sub(r"--[^\n]*", " ", without_block_comments)


def _unqualified_identifier_candidates(sql: str) -> set[str]:
    """Find identifiers that are neither qualified columns nor query structure."""

    scrubbed = _strip_sql_strings_and_comments(sql)
    allowed = {*NONCANONICAL_QUERY_IDENTIFIERS, DYNAMIC_SQL_VALUE}

    # Output aliases and relation aliases may be consumed unqualified by an
    # enclosing query. CTE names/header columns and subquery aliases are also
    # query structure rather than base-relation column reads.
    allowed.update(match.group(1).lower() for match in re.finditer(r"\bas\s+([a-z_][a-z0-9_]*)", scrubbed, re.I))
    allowed.update(match.group(3).lower() for match in ALIASED_CANONICAL_RELATION.finditer(scrubbed))
    allowed.update(
        match.group(1).lower()
        for match in re.finditer(
            r"(?:\bwith|,)\s*([a-z_][a-z0-9_]*)\s*(?:\([^)]*\))?\s+as\s*\(",
            scrubbed,
            re.I,
        )
    )
    for match in re.finditer(
        r"(?:\bwith|,)\s*[a-z_][a-z0-9_]*\s*\(([^)]*)\)\s+as\s*\(",
        scrubbed,
        re.I,
    ):
        allowed.update(identifier.lower() for identifier in SQL_IDENTIFIER.findall(match.group(1)))
    allowed.update(
        match.group(1).lower() for match in re.finditer(r"\)\s+(?:as\s+)?([a-z_][a-z0-9_]*)", scrubbed, re.I)
    )
    allowed.update(
        match.group(1).lower()
        for match in re.finditer(
            r"\b(?:from|join)\s+[a-z_][a-z0-9_]*(?:\s+as)?\s+([a-z_][a-z0-9_]*)",
            scrubbed,
            re.I,
        )
    )
    # A qualified projection becomes a mechanically known output name for an
    # enclosing CTE query (for example ``ce.claim_id`` -> ``claim_id``).
    allowed.update(match.group(1).lower() for match in re.finditer(r"\.\s*([a-z_][a-z0-9_]*)", scrubbed, re.I))

    scrubbed = re.sub(r"\b[a-z_][a-z0-9_]*\s*\.\s*[a-z_][a-z0-9_]*", " ", scrubbed, flags=re.I)
    scrubbed = re.sub(r"::\s*[a-z_][a-z0-9_]*", " ", scrubbed, flags=re.I)

    candidates: set[str] = set()
    for match in SQL_IDENTIFIER.finditer(scrubbed):
        identifier = match.group(0).lower()
        if identifier in SQL_KEYWORDS or identifier in allowed:
            continue
        if re.match(r"\s*\(", scrubbed[match.end() :]):
            continue
        candidates.add(identifier)
    return candidates


def _reject_inexact_canonical_alias_uses(
    sql: str,
    line: int,
    aliases: dict[str, tuple[str, str]],
) -> None:
    """Reject canonical reads that cannot be reduced to named columns."""

    scrubbed = _strip_sql_strings_and_comments(sql)
    for alias in aliases:
        if re.search(
            rf"(?<![a-z0-9_]){re.escape(alias)}\s*\.\s*{re.escape(DYNAMIC_SQL_VALUE)}(?![a-z0-9_])",
            scrubbed,
            re.IGNORECASE,
        ):
            raise QueryContractError(
                f"line {line}: dynamic SQL interpolation cannot supply a canonical column identifier"
            )
        if re.search(
            rf"(?<![a-z0-9_]){re.escape(alias)}\s*\.\s*\*",
            scrubbed,
            re.IGNORECASE,
        ):
            raise QueryContractError(f"line {line}: canonical wildcard reads prevent exact derivation")

        without_relations = ALIASED_CANONICAL_RELATION.sub(" ", scrubbed)
        without_qualified_columns = re.sub(
            rf"(?<![a-z0-9_]){re.escape(alias)}\s*\.\s*[a-z_][a-z0-9_]*",
            " ",
            without_relations,
            flags=re.IGNORECASE,
        )
        if re.search(
            rf"(?<![a-z0-9_]){re.escape(alias)}(?![a-z0-9_])",
            without_qualified_columns,
            re.IGNORECASE,
        ):
            raise QueryContractError(
                f"line {line}: whole-row canonical alias {alias!r} prevents exact column derivation"
            )

    self_alias = re.search(
        r"(?<![a-z0-9_.:])([a-z_][a-z0-9_]*)\s+as\s+\1(?![a-z0-9_])",
        scrubbed,
        re.IGNORECASE,
    )
    if self_alias:
        raise QueryContractError(
            f"line {line}: unqualified SQL identifier {self_alias.group(1).lower()!r} is masked by an output alias"
        )


def derive_runtime_read_columns(source: str) -> dict[tuple[str, str], frozenset[str]]:
    """Mechanically derive canonical relation columns from runtime SQL."""

    observed: dict[tuple[str, str], set[str]] = defaultdict(set)
    for line, sql in extract_canonical_query_literals(source):
        aliases: dict[str, tuple[str, str]] = {}
        explicit_starts: set[int] = set()
        for match in ALIASED_CANONICAL_RELATION.finditer(sql):
            relation = (match.group(1).lower(), match.group(2).lower())
            alias = match.group(3).lower()
            existing = aliases.get(alias)
            if existing is not None and existing != relation:
                raise QueryContractError(
                    f"line {line}: canonical alias {alias!r} maps to both {existing!r} and {relation!r}"
                )
            aliases[alias] = relation
            explicit_starts.add(match.start())

        for match in CANONICAL_RELATION.finditer(sql):
            if match.start() not in explicit_starts:
                relation = f"{match.group(1).lower()}.{match.group(2).lower()}"
                raise QueryContractError(f"line {line}: canonical relation {relation} requires an explicit AS alias")

        _reject_inexact_canonical_alias_uses(sql, line, aliases)

        for alias, relation in aliases.items():
            observed[relation].update(
                column.lower()
                for column in re.findall(
                    rf"(?<![a-z0-9_]){re.escape(alias)}\.([a-z_][a-z0-9_]*)",
                    sql,
                    re.IGNORECASE,
                )
            )

        unqualified = _unqualified_identifier_candidates(sql)
        if unqualified:
            joined = ", ".join(sorted(unqualified))
            raise QueryContractError(f"line {line}: unqualified SQL identifiers prevent exact derivation: {joined}")

    return {relation: frozenset(columns) for relation, columns in observed.items()}


def verify_runtime_read_columns(
    source: str,
    expected: dict[tuple[str, str], frozenset[str]],
) -> dict[tuple[str, str], frozenset[str]]:
    """Fail when the runtime query surface and reviewed allowlist diverge."""

    observed = derive_runtime_read_columns(source)
    normalized_expected = {relation: frozenset(columns) for relation, columns in expected.items()}
    if observed != normalized_expected:
        missing_relations = sorted(normalized_expected.keys() - observed.keys())
        extra_relations = sorted(observed.keys() - normalized_expected.keys())
        column_drift = {
            f"{schema}.{relation}": {
                "missing": sorted(normalized_expected[(schema, relation)] - observed[(schema, relation)]),
                "extra": sorted(observed[(schema, relation)] - normalized_expected[(schema, relation)]),
            }
            for schema, relation in sorted(normalized_expected.keys() & observed.keys())
            if observed[(schema, relation)] != normalized_expected[(schema, relation)]
        }
        detail = json.dumps(
            {
                "missing_relations": [".".join(relation) for relation in missing_relations],
                "extra_relations": [".".join(relation) for relation in extra_relations],
                "column_drift": column_drift,
            },
            sort_keys=True,
        )
        raise QueryContractError(f"runtime query surface differs from reviewed read allowlist: {detail}")
    return observed


def main() -> None:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("runtime", type=Path)
    args = parser.parse_args()
    contract = derive_runtime_read_columns(args.runtime.read_text(encoding="utf-8"))
    payload = {f"{schema}.{relation}": sorted(columns) for (schema, relation), columns in sorted(contract.items())}
    print(json.dumps(payload, indent=2, sort_keys=True))


if __name__ == "__main__":
    main()