from tests._helpers import (
  unittest, parameterized, Fraction, time,
  run_test, run_compiler_error_test,
  ArrayIndex, ArraySlice, Bool, BinOp, BuiltIn, DataArray, DataDict, DataRangeArray,
  Defini, Continva, Designa, DesignaDestructure, DesignaIndex, DumStatement,
  Erumpe, ExpressionStatement, Fvnctio, ID, InterpolatedString, Invoca,
  ModuleCall, Nullus, Numeral, PerStatement, Program, Redi, SiStatement,
  String, TemptaStatement, UnaryMinus, UnaryNot, Fractio, frac_to_fraction,
  fraction_to_frac, num_to_int, int_to_num, make_string,
  ValInt, ValStr, ValBool, ValList, ValDict, ValNul, ValFunc, ValFrac,
  CentvrionError, _RUNTIME_C, _cent_rng,
  Lexer, Parser, compile_program,
  os, subprocess, tempfile, StringIO, patch,
)


# --- Dict (TABVLA) ---

dict_tests = [
  # empty dict
  ("TABVLA {}",
    Program([], [ExpressionStatement(DataDict([]))]),
    ValDict({})),
  # single string key
  ('TABVLA {"a" VT I}',
    Program([], [ExpressionStatement(DataDict([(String("a"), Numeral("I"))]))]),
    ValDict({"a": ValInt(1)})),
  # multiple entries
  ('TABVLA {"a" VT I, "b" VT II}',
    Program([], [ExpressionStatement(DataDict([(String("a"), Numeral("I")), (String("b"), Numeral("II"))]))]),
    ValDict({"a": ValInt(1), "b": ValInt(2)})),
  # integer keys
  ('TABVLA {I VT "one", II VT "two"}',
    Program([], [ExpressionStatement(DataDict([(Numeral("I"), String("one")), (Numeral("II"), String("two"))]))]),
    ValDict({1: ValStr("one"), 2: ValStr("two")})),
  # expression values
  ('TABVLA {"x" VT I + II}',
    Program([], [ExpressionStatement(DataDict([(String("x"), BinOp(Numeral("I"), Numeral("II"), "SYMBOL_PLUS"))]))]),
    ValDict({"x": ValInt(3)})),
  # empty dict with newline
  ('TABVLA {\n}',
    Program([], [ExpressionStatement(DataDict([]))]),
    ValDict({})),
  # single entry with surrounding newlines
  ('TABVLA {\n"a" VT I\n}',
    Program([], [ExpressionStatement(DataDict([(String("a"), Numeral("I"))]))]),
    ValDict({"a": ValInt(1)})),
  # multiple entries with newlines after commas
  ('TABVLA {\n"a" VT I,\n"b" VT II\n}',
    Program([], [ExpressionStatement(DataDict([(String("a"), Numeral("I")), (String("b"), Numeral("II"))]))]),
    ValDict({"a": ValInt(1), "b": ValInt(2)})),
  # newlines around every delimiter
  ('TABVLA {\n"a" VT I,\n"b" VT II,\n"c" VT III\n}',
    Program([], [ExpressionStatement(DataDict([(String("a"), Numeral("I")), (String("b"), Numeral("II")), (String("c"), Numeral("III"))]))]),
    ValDict({"a": ValInt(1), "b": ValInt(2), "c": ValInt(3)})),
]

class TestDict(unittest.TestCase):
  @parameterized.expand(dict_tests)
  def test_dict(self, source, nodes, value):
    run_test(self, source, nodes, value)


dict_index_tests = [
  # string key access
  ('TABVLA {"a" VT X}["a"]',
    Program([], [ExpressionStatement(ArrayIndex(DataDict([(String("a"), Numeral("X"))]), String("a")))]),
    ValInt(10)),
  # integer key access
  ('TABVLA {I VT "one"}[I]',
    Program([], [ExpressionStatement(ArrayIndex(DataDict([(Numeral("I"), String("one"))]), Numeral("I")))]),
    ValStr("one")),
  # access via variable
  ('DESIGNA d VT TABVLA {"x" VT V}\nd["x"]',
    Program([], [
      Designa(ID("d"), DataDict([(String("x"), Numeral("V"))])),
      ExpressionStatement(ArrayIndex(ID("d"), String("x"))),
    ]),
    ValInt(5)),
  # nested dict access
  ('TABVLA {"a" VT TABVLA {"b" VT X}}["a"]["b"]',
    Program([], [ExpressionStatement(
      ArrayIndex(ArrayIndex(DataDict([(String("a"), DataDict([(String("b"), Numeral("X"))]))]), String("a")), String("b"))
    )]),
    ValInt(10)),
]

class TestDictIndex(unittest.TestCase):
  @parameterized.expand(dict_index_tests)
  def test_dict_index(self, source, nodes, value):
    run_test(self, source, nodes, value)


dict_assign_tests = [
  # update existing key
  ('DESIGNA d VT TABVLA {"a" VT I}\nDESIGNA d["a"] VT X\nd["a"]',
    Program([], [
      Designa(ID("d"), DataDict([(String("a"), Numeral("I"))])),
      DesignaIndex(ID("d"), [String("a")], Numeral("X")),
      ExpressionStatement(ArrayIndex(ID("d"), String("a"))),
    ]),
    ValInt(10)),
  # insert new key
  ('DESIGNA d VT TABVLA {"a" VT I}\nDESIGNA d["b"] VT II\nd["b"]',
    Program([], [
      Designa(ID("d"), DataDict([(String("a"), Numeral("I"))])),
      DesignaIndex(ID("d"), [String("b")], Numeral("II")),
      ExpressionStatement(ArrayIndex(ID("d"), String("b"))),
    ]),
    ValInt(2)),
  # original key unaffected after insert
  ('DESIGNA d VT TABVLA {"a" VT I}\nDESIGNA d["b"] VT II\nd["a"]',
    Program([], [
      Designa(ID("d"), DataDict([(String("a"), Numeral("I"))])),
      DesignaIndex(ID("d"), [String("b")], Numeral("II")),
      ExpressionStatement(ArrayIndex(ID("d"), String("a"))),
    ]),
    ValInt(1)),
]

class TestDictAssign(unittest.TestCase):
  @parameterized.expand(dict_assign_tests)
  def test_dict_assign(self, source, nodes, value):
    run_test(self, source, nodes, value)


dict_builtin_tests = [
  # LONGITVDO on dict
  ('LONGITVDO(TABVLA {"a" VT I, "b" VT II})',
    Program([], [ExpressionStatement(BuiltIn("LONGITVDO", [DataDict([(String("a"), Numeral("I")), (String("b"), Numeral("II"))])]))]),
    ValInt(2)),
  # LONGITVDO on empty dict
  ('LONGITVDO(TABVLA {})',
    Program([], [ExpressionStatement(BuiltIn("LONGITVDO", [DataDict([])]))]),
    ValInt(0)),
  # CLAVES
  ('CLAVES(TABVLA {"a" VT I, "b" VT II})',
    Program([], [ExpressionStatement(BuiltIn("CLAVES", [DataDict([(String("a"), Numeral("I")), (String("b"), Numeral("II"))])]))]),
    ValList([ValStr("a"), ValStr("b")])),
  # CLAVES with int keys
  ('CLAVES(TABVLA {I VT "x", II VT "y"})',
    Program([], [ExpressionStatement(BuiltIn("CLAVES", [DataDict([(Numeral("I"), String("x")), (Numeral("II"), String("y"))])]))]),
    ValList([ValInt(1), ValInt(2)])),
]

class TestDictBuiltins(unittest.TestCase):
  @parameterized.expand(dict_builtin_tests)
  def test_dict_builtin(self, source, nodes, value):
    run_test(self, source, nodes, value)


dict_iteration_tests = [
  # PER iterates over keys
  ('DESIGNA r VT ""\nPER k IN TABVLA {"a" VT I, "b" VT II} FAC {\nDESIGNA r VT r & k\n}\nr',
    Program([], [
      Designa(ID("r"), String("")),
      PerStatement(
        DataDict([(String("a"), Numeral("I")), (String("b"), Numeral("II"))]),
        ID("k"),
        [Designa(ID("r"), BinOp(ID("r"), ID("k"), "SYMBOL_AMPERSAND"))],
      ),
      ExpressionStatement(ID("r")),
    ]),
    ValStr("ab")),
]

class TestDictIteration(unittest.TestCase):
  @parameterized.expand(dict_iteration_tests)
  def test_dict_iteration(self, source, nodes, value):
    run_test(self, source, nodes, value)


dict_display_tests = [
  # DIC on dict
  ('DIC(TABVLA {"a" VT I})',
    Program([], [ExpressionStatement(BuiltIn("DIC", [DataDict([(String("a"), Numeral("I"))])]))]),
    ValStr("{a VT I}"), "{a VT I}\n"),
  # DIC on multi-entry dict
  ('DIC(TABVLA {"a" VT I, "b" VT II})',
    Program([], [ExpressionStatement(BuiltIn("DIC", [DataDict([(String("a"), Numeral("I")), (String("b"), Numeral("II"))])]))]),
    ValStr("{a VT I, b VT II}"), "{a VT I, b VT II}\n"),
  # DIC on empty dict
  ('DIC(TABVLA {})',
    Program([], [ExpressionStatement(BuiltIn("DIC", [DataDict([])]))]),
    ValStr("{}"), "{}\n"),
]

class TestDictDisplay(unittest.TestCase):
  @parameterized.expand(dict_display_tests)
  def test_dict_display(self, source, nodes, value, output):
    run_test(self, source, nodes, value, output)