import unittest
import json

from autark import JsonReader, OnceFrame, P


class TestApiContract(unittest.TestCase):
    def test_all_unaryop_matches_explicit_columns(self) -> None:
        rows = [
            {"v": 1.0, "w": 10.0},
            {"v": 2.0, "w": 20.0},
            {"v": 3.0, "w": 30.0},
        ]
        frame = OnceFrame(JsonReader(rows))
        p = P()
        schema = frame.schema_of_columns(["v", "w"])

        p.columns(None).unaryop("neg").alias("implicit", schema)
        p.columns(["v", "w"]).unaryop("neg").alias("explicit", schema)

        realized = frame.realize(p)
        implicit = realized["implicit"]
        explicit = realized["explicit"]

        self.assertEqual(implicit, explicit)
        self.assertEqual(implicit["v"], [-1.0, -2.0, -3.0])
        self.assertEqual(implicit["w"], [-10.0, -20.0, -30.0])

    def test_realized_mapping_behaviors(self) -> None:
        rows = [
            {"city": "A", "v": 1.0},
            {"city": "B", "v": 2.0},
            {"city": "C", "v": 3.0},
        ]
        frame = OnceFrame(JsonReader(rows))
        p = P()
        p.columns(None).slice(0, 2).alias("head", frame.schema())

        realized = frame.realize(p)

        self.assertEqual(len(realized), 1)
        self.assertIn("head", realized)
        self.assertEqual(realized.get("missing"), None)
        self.assertEqual(realized.get("missing", {"ok": True}), {"ok": True})
        with self.assertRaises(KeyError):
            _ = realized["missing"]

    def test_columns_accepts_positional_indices(self) -> None:
        rows = [
            {"city": "A", "v": 1.0},
            {"city": "B", "v": 2.0},
        ]
        frame = OnceFrame(JsonReader(rows))
        p = P()

        p.columns([1]).alias("positional")
        p.columns(1).alias("legacy_take")

        realized = frame.realize(p)

        positional_values = next(iter(realized["positional"].values()))
        legacy_values = next(iter(realized["legacy_take"].values()))
        self.assertEqual(positional_values, legacy_values)
        self.assertEqual(positional_values, [1.0, 2.0])

    def test_rootless_source_access_defaults_to_input_zero(self) -> None:
        rows = [
            {"city": "A", "v": 1.0},
            {"city": "B", "v": 2.0},
            {"city": "C", "v": 3.0},
        ]
        frame = OnceFrame(JsonReader(rows))
        p = P()

        p.columns("v").reduce("count").alias("count")
        p.columns(None).slice(0, 2).alias("head", frame.schema())

        realized = frame.realize(p)

        count_values = next(iter(realized["count"].values()))
        self.assertEqual(int(count_values[0]), len(rows))
        self.assertEqual(len(realized["head"]["city"]), 2)

    def test_tuple_readers_second_dataframe(self) -> None:
        left_rows = [{"x": 1.0}, {"x": 2.0}]
        right_rows = [{"x": 10.0}, {"x": 20.0}, {"x": 30.0}]
        frame = OnceFrame((JsonReader(left_rows), JsonReader(right_rows)))
        p = P()

        p.columns("x", index=1).reduce("count").alias("right_count")
        realized = frame.realize(p)

        count_table = realized["right_count"]
        count_values = next(iter(count_table.values()))
        self.assertEqual(int(count_values[0]), len(right_rows))

    def test_validation_errors(self) -> None:
        p1 = P()
        p2 = P()

        with self.assertRaises(ValueError):
            _ = p1.columns(None) + p2.columns(None)

        with self.assertRaises(ValueError):
            _ = p1.columns("v").reduce("percentile", x=2.0)

    def test_program_ir_is_json(self) -> None:
        p = P()

        ir = p.columns("v").alias("out").ir()
        value = json.loads(ir)

        self.assertIn("op_pool", value)
        self.assertIn("metadata", value)
        self.assertIn("signature", value)
        self.assertIsNotNone(value["root"])


if __name__ == "__main__":
    unittest.main()
