"""AGType parsing: ANTLR4 parser → Piggie dataclasses.

Wraps the upstream AGE ANTLR4 generated parser and provides a psycopg3 Loader
that converts agtype wire values into Piggie's Vertex/Edge/Path types.
"""

from __future__ import annotations

from decimal import Decimal
from typing import Any

import psycopg.adapt
from antlr4 import CommonTokenStream, InputStream
from antlr4.error.ErrorListener import ErrorListener
from antlr4.tree.Tree import TerminalNode

from piggie._agtype.gen.AgtypeLexer import AgtypeLexer
from piggie._agtype.gen.AgtypeParser import AgtypeParser
from piggie._agtype.gen.AgtypeVisitor import AgtypeVisitor
from piggie.exceptions import AgtypeParseError
from piggie.types import Edge, Path, Vertex


class _ResultVisitor(AgtypeVisitor):
    """ANTLR4 visitor that converts AGType parse trees into Python objects.

    Produces Piggie's frozen dataclasses (Vertex, Edge, Path) instead of the
    mutable model objects from the upstream AGE driver.
    """

    def visitAgType(self, ctx: AgtypeParser.AgTypeContext) -> Any:
        ag_val = ctx.agValue()
        if ag_val is not None:
            return ag_val.accept(self)
        return None

    def visitAgValue(self, ctx: AgtypeParser.AgValueContext) -> Any:
        anno_ctx = ctx.typeAnnotation()
        value_ctx = ctx.value()
        if anno_ctx is not None:
            anno_ctx.accept(self)
            anno = anno_ctx.IDENT().getText()
            return self._handle_annotated(anno, value_ctx)
        return value_ctx.accept(self)

    def visitStringValue(self, ctx: AgtypeParser.StringValueContext) -> str:
        return ctx.STRING().getText().strip('"')

    def visitIntegerValue(self, ctx: AgtypeParser.IntegerValueContext) -> int:
        return int(ctx.INTEGER().getText())

    def visitFloatLiteral(self, ctx: AgtypeParser.FloatLiteralContext) -> float:
        child = ctx.getChild(0)
        tp = child.symbol.type
        text = ctx.getText()
        if tp in (AgtypeParser.RegularFloat, AgtypeParser.ExponentFloat):
            return float(text)
        if text == "NaN":
            return float("nan")
        if text == "-Infinity":
            return float("-inf")
        if text == "Infinity":
            return float("inf")
        raise ValueError(f"Unknown float expression: {text}")

    def visitTrueBoolean(self, ctx: AgtypeParser.TrueBooleanContext) -> bool:
        return True

    def visitFalseBoolean(self, ctx: AgtypeParser.FalseBooleanContext) -> bool:
        return False

    def visitNullValue(self, ctx: AgtypeParser.NullValueContext) -> None:
        return None

    def visitObj(self, ctx: AgtypeParser.ObjContext) -> dict[str, Any]:
        obj: dict[str, Any] = {}
        for c in ctx.getChildren():
            if isinstance(c, AgtypeParser.PairContext):
                name, val_ctx = self.visitPair(c)
                obj[name] = val_ctx.accept(self)
        return obj

    def visitPair(self, ctx: AgtypeParser.PairContext) -> tuple[str, Any]:
        self.visitChildren(ctx)
        return (ctx.STRING().getText().strip('"'), ctx.agValue())

    def visitArray(self, ctx: AgtypeParser.ArrayContext) -> list[Any]:
        items: list[Any] = []
        for c in ctx.getChildren():
            if not isinstance(c, TerminalNode):
                items.append(c.accept(self))
        return items

    def _handle_annotated(self, anno: str, ctx: Any) -> Any:
        if anno == "numeric":
            return Decimal(ctx.getText())

        if anno == "vertex":
            d = ctx.accept(self)
            return Vertex(
                id=d["id"],
                label=d["label"],
                properties=d.get("properties", {}),
            )

        if anno == "edge":
            d = ctx.accept(self)
            return Edge(
                id=d["id"],
                label=d["label"],
                start_id=d["start_id"],
                end_id=d["end_id"],
                properties=d.get("properties", {}),
            )

        if anno == "path":
            arr = ctx.accept(self)
            return Path(entities=tuple(arr))

        return ctx.accept(self)


class _StrictErrorListener(ErrorListener):
    """ANTLR4 error listener that raises on any syntax error."""

    def syntaxError(
        self, recognizer: Any, offendingSymbol: Any,
        line: int, column: int, msg: str, e: Any,
    ) -> None:
        raise AgtypeParseError(msg)


def parse_agtype(value: str) -> Any:
    """Parse an agtype string into Python objects.

    Returns primitive types (str, int, float, bool, None, Decimal) for scalar
    values, and Vertex/Edge/Path for graph objects.
    """
    if not value:
        return None

    try:
        lexer = AgtypeLexer(InputStream(value))
        lexer.removeErrorListeners()
        lexer.addErrorListener(_StrictErrorListener())
        parser = AgtypeParser(CommonTokenStream(lexer))
        parser.removeErrorListeners()
        parser.addErrorListener(_StrictErrorListener())
        tree = parser.agType()
        return tree.accept(_ResultVisitor())
    except Exception as exc:
        raise AgtypeParseError(value, exc) from exc


class AgtypeLoader(psycopg.adapt.Loader):
    """psycopg3 Loader that deserializes agtype columns into Python objects."""

    def load(self, data: bytes | bytearray | memoryview) -> Any:
        if isinstance(data, memoryview):
            raw: bytes | bytearray = data.tobytes()
        else:
            raw = data
        return parse_agtype(raw.decode("utf-8"))
