from __future__ import annotations

import unittest

from sqlalchemy import JSON, Column, Integer, MetaData, Table, create_engine, select
from sqlalchemy.dialects import mysql, postgresql
from sqlalchemy.exc import CompileError

from govoplan_core.db.json_predicates import (
    json_array_contains_object_strings,
    json_array_contains_string,
    json_object_matches_strings,
)


class JsonPredicateTests(unittest.TestCase):
    def setUp(self):
        self.engine = create_engine("sqlite:///:memory:")
        self.metadata = MetaData()
        self.records = Table("records", self.metadata, Column("id", Integer, primary_key=True), Column("value", JSON))
        self.metadata.create_all(self.engine)

    def tearDown(self):
        self.engine.dispose()

    def matched(self, values, predicate):
        with self.engine.begin() as connection:
            connection.execute(self.records.insert(), [{"id": index, "value": value} for index, value in enumerate(values)])
            return list(connection.scalars(select(self.records.c.id).where(predicate).order_by(self.records.c.id)))

    def test_string_membership_is_array_and_type_exact(self):
        self.assertEqual([0], self.matched(
            [["1"], [1], [True], [None], {"key": "1"}, "1", None, ["11"]],
            json_array_contains_string(self.records.c.value, "1"),
        ))

    def test_object_fields_are_exact_strings_not_coerced_or_substrings(self):
        self.assertEqual([0], self.matched(
            [{"kind": "account", "id": "1", "label": "Extra field allowed"},
             {"kind": "account", "id": 1}, {"kind": "account", "id": "11"},
             "not-json", None, ["account", "1"]],
            json_object_matches_strings(self.records.c.value, {"kind": "account", "id": "1"}),
        ))

    def test_array_objects_do_not_match_encoded_objects_or_scalar_elements(self):
        self.assertEqual([0], self.matched(
            [[{"kind": "account", "id": "1"}], ['{"kind":"account","id":"1"}'],
             [{"kind": "account", "id": 1}], ["not-json", None, True, 1],
             {"kind": "account", "id": "1"}, None],
            json_array_contains_object_strings(self.records.c.value, {"kind": "account", "id": "1"}),
        ))

    def test_values_remain_bound_for_both_dialects(self):
        value = "x' OR 1=1 --"
        predicates = [
            json_array_contains_string(self.records.c.value, value),
            json_object_matches_strings(self.records.c.value, {"id": value}),
            json_array_contains_object_strings(self.records.c.value, {"id": value}),
        ]
        for predicate in predicates:
            for dialect in (self.engine.dialect, postgresql.dialect()):
                with self.subTest(predicate=type(predicate).__name__, dialect=dialect.name):
                    compiled = select(self.records.c.id).where(predicate).compile(dialect=dialect)
                    self.assertNotIn(value, str(compiled))
                    self.assertIn(value, compiled.params.values())
            with self.assertRaises(CompileError):
                select(self.records.c.id).where(predicate).compile(dialect=mysql.dialect())

    def test_invalid_field_names_and_non_string_matches_are_rejected(self):
        for fields in ({"id": 1}, {"not.a.field": "1"}, {}):
            with self.subTest(fields=fields), self.assertRaises((TypeError, ValueError)):
                json_array_contains_object_strings(self.records.c.value, fields)
        with self.assertRaises(TypeError):
            json_array_contains_string(self.records.c.value, True)