"""Bound JSON predicates for SQLite, PostgreSQL and MySQL.

Adapted from upstream #2865. MySQL business tables store JSON in TEXT, so
these predicates use JSON functions rather than assuming a native JSON column.
Unsupported keys/values must stay in the repository's residual filter; never
silently discard a predicate (in particular a shared task's taskId filter).
"""

from __future__ import annotations

import math
import re
from typing import Any

from sqlalchemy import BigInteger, Float, String, bindparam
from sqlalchemy.ext.compiler import compiles
from sqlalchemy.sql.elements import ColumnElement
from sqlalchemy.sql.visitors import InternalTraversal
from sqlalchemy.types import Boolean

_SAFE_KEY = re.compile(r"[A-Za-z0-9_-]+\Z")
_STRIP_CHARS = "".join(chr(i) for i in range(0x3001) if chr(i).isspace())


def supports_json_filter(key: Any, value: Any) -> bool:
    if not isinstance(key, str) or not _SAFE_KEY.fullmatch(key):
        return False
    if value is None or isinstance(value, (str, bool)):
        return True
    if isinstance(value, int):
        return -(2**63) <= value < 2**63
    return isinstance(value, float) and math.isfinite(value)


class JsonPredicate(ColumnElement):
    """A scalar equality, Python truthiness, or nonblank-string predicate."""

    inherit_cache = True
    type = Boolean()
    _is_implicitly_boolean = True
    _traverse_internals = [
        ("column", InternalTraversal.dp_clauseelement),
        ("key", InternalTraversal.dp_string),
        ("operation", InternalTraversal.dp_string),
        ("value", InternalTraversal.dp_plain_obj),
        ("value_kind", InternalTraversal.dp_string),
    ]

    def __init__(self, column, key: str, operation: str = "equal", value: Any = None):
        if not supports_json_filter(key, value):
            raise ValueError("Unsupported JSON filter key or value")
        if operation not in {"equal", "truthy", "nonblank"}:
            raise ValueError("Unsupported JSON predicate operation")
        self.column = column
        self.key = key
        self.operation = operation
        self.value = value
        # Keep Python value types separate in SQLAlchemy's compiled cache.
        self.value_kind = type(value).__name__
        super().__init__()


def _compile(element, compiler, family, **kw):
    if not _SAFE_KEY.fullmatch(element.key):
        raise ValueError("Unsafe JSON key")
    column = compiler.process(element.column, **kw)

    def bound(value, type_=String()):
        return compiler.process(bindparam(None, value, type_=type_), **kw)

    if family == "postgres":
        raw = f"({column} -> '{element.key}')"
        kind = f"json_typeof({raw})"
        text = f"({column} ->> '{element.key}')"
        string_kind, number_kinds = "string", "'number'"
    else:
        path = f"'$.\"{element.key}\"'"
        raw = f"json_extract({column}, {path})"
        if family == "mysql":
            kind, text = f"lower(json_type({raw}))", f"json_unquote({raw})"
            string_kind, number_kinds = "string", "'integer', 'double', 'decimal'"
        else:
            kind, text = f"json_type({column}, {path})", raw
            string_kind, number_kinds = "text", "'integer', 'real'"

    if element.operation == "nonblank":
        if family == "mysql":
            # MySQL 5.7 regexes are byte-oriented; a Unicode whitespace regex
            # can misclassify CJK text. Removing each complete whitespace
            # character preserves str.strip()'s nonempty test on 5.7 and 8.x.
            stripped = text
            for char in _STRIP_CHARS:
                stripped = f"replace({stripped}, {bound(char)}, '')"
            has_text = f"length({stripped}) > 0"
        else:
            trim = "btrim" if family == "postgres" else "trim"
            # SQLite length(TEXT) stops at NUL, unlike Python str.strip().
            has_text = f"{trim}({text}, {bound(_STRIP_CHARS)}) <> ''"
        clause = f"({kind} = '{string_kind}' AND {has_text})"
    elif element.operation == "truthy":
        if family == "sqlite":
            containers = f"json({raw}) NOT IN ('[]', '{{}}')"
            boolean = f"{kind} = 'true'"
        elif family == "mysql":
            containers = f"json_length({raw}) > 0"
            boolean = f"{text} = 'true'"
        else:
            containers = f"CAST({raw} AS jsonb) NOT IN ('[]'::jsonb, '{{}}'::jsonb)"
            boolean = f"{text} = 'true'"
        cast = "NUMERIC" if family == "postgres" else "DOUBLE PRECISION"
        number = f"({text} + 0e0)" if family == "mysql" else f"CAST({text} AS {cast})"
        nonempty = f"{text} <> ''" if family == "sqlite" else f"length({text}) > 0"
        clause = (
            f"CASE WHEN {kind} IN ({number_kinds}) THEN {number} <> 0 "
            f"WHEN {kind} = '{string_kind}' THEN {nonempty} "
            f"WHEN {kind} IN ('array', 'object') THEN {containers} "
            f"WHEN {kind} IN ('boolean', 'true', 'false') THEN {boolean} ELSE FALSE END"
        )
    elif element.value is None:
        # Preserve this project's md.get(key) == None semantics for old rows.
        clause = f"({kind} IS NULL OR {kind} = 'null')"
    elif isinstance(element.value, (bool, int, float)):
        if family != "sqlite":
            raise NotImplementedError("Server numeric metadata requires the exact residual filter")
        value = int(element.value) if isinstance(element.value, bool) else element.value
        # Preserve md.get(key) == value: Python considers True == 1 == 1.0.
        # Text ids remain distinct from numbers. SQLite compares its extracted
        # integers/reals natively; server numeric filters use the residual path.
        param = bound(value, Float() if isinstance(value, float) else BigInteger())
        clause = f"({kind} IN ({number_kinds}, 'true', 'false') AND {text} = {param})"
    else:
        # Binary comparison is essential for task/agent ids: MySQL TEXT can
        # inherit a case-insensitive collation and ignores trailing spaces.
        extract = f"CAST({text} AS BINARY)" if family == "mysql" else text
        param = bound(element.value)
        if family == "mysql":
            param = f"CAST({param} AS BINARY)"
        clause = f"({kind} = '{string_kind}' AND {extract} = {param})"
    # Missing/null values must remain false under NOT (system exclusions).
    return f"COALESCE(({clause}), FALSE)"


@compiles(JsonPredicate, "sqlite")
def _sqlite(element, compiler, **kw):
    return _compile(element, compiler, "sqlite", **kw)


@compiles(JsonPredicate, "postgresql")
def _postgres(element, compiler, **kw):
    return _compile(element, compiler, "postgres", **kw)


@compiles(JsonPredicate, "mysql")
def _mysql(element, compiler, **kw):
    return _compile(element, compiler, "mysql", **kw)


@compiles(JsonPredicate)
def _unsupported(element, compiler, **kw):
    raise NotImplementedError(f"JSON predicates unsupported for {compiler.dialect.name}")