import unittest

from datainventor import visualize as V


def sample(n=12):
    areas = ["Billing", "Returns", "Shipping"]
    return [{"id": i, "prompt": f"Question number {i} about an order that arrived late and damaged",
             "completion": " ".join(["word"] * (20 + 10 * i)), "area": areas[i % 3], "difficulty": "basic" if i % 2 else "advanced",
             "score": str(i % 5 + 1), "ok": bool(i % 2)} for i in range(n)]


class Prepare(unittest.TestCase):
    def test_types_and_derived_lengths(self):
        rows, cols = V.prepare(sample())
        self.assertEqual(cols["prompt"]["type"], "text")
        self.assertEqual(cols["area"]["type"], "cat")
        self.assertEqual(cols["score"]["type"], "num")
        self.assertEqual(cols["ok"]["type"], "cat")
        self.assertNotIn("id", cols)
        self.assertEqual(cols["completion words"]["derived"], "words")
        self.assertEqual(rows[0]["_id"], 1)
        self.assertEqual(rows[0]["completion words"], 20.0)

    def test_chat_cells_become_text_with_turns(self):
        rows, cols = V.prepare([{"messages": [{"role": "user", "content": "hi " * 30}, {"role": "assistant", "content": "hello " * 30}]}] * 3)
        self.assertEqual(cols["messages"]["type"], "text")
        self.assertEqual(rows[0]["messages turns"], 2.0)


class Charts(unittest.TestCase):
    def setUp(self):
        self.rows, self.cols = V.prepare(sample(30))

    def test_every_row_lands_in_exactly_one_bin(self):
        c = V.chart(self.rows, self.cols, {"type": "hist", "x": "completion words"})
        ids = [i for b in c["marks"] for i in b["ids"]]
        self.assertEqual(sorted(ids), list(range(1, 31)))

    def test_bar_counts_and_ids(self):
        c = V.chart(self.rows, self.cols, {"type": "bar", "x": "area", "agg": "count"})
        self.assertEqual(sum(m["value"] for m in c["marks"]), 30)
        self.assertTrue(all(len(m["ids"]) == m["value"] for m in c["marks"]))

    def test_box_quartiles_are_ordered(self):
        c = V.chart(self.rows, self.cols, {"type": "box", "x": "area", "y": "completion words"})
        for m in c["marks"]:
            self.assertTrue(m["lo"] <= m["q1"] <= m["median"] <= m["q3"] <= m["hi"])

    def test_check_spec_fixes_and_rejects(self):
        self.assertEqual(V.check_spec({"type": "bar", "x": "completion words"}, self.cols)["type"], "hist")
        self.assertEqual(V.check_spec({"type": "bar", "x": "area", "y": "score"}, self.cols)["agg"], "mean")
        with self.assertRaises(V.PlanError):
            V.check_spec({"type": "box", "x": "area", "y": "area"}, self.cols)
        with self.assertRaises(V.PlanError):
            V.check_spec({"type": "bar", "x": "missing"}, self.cols)


class Plans(unittest.TestCase):
    def setUp(self):
        self.rows, self.cols = V.prepare(sample())

    def test_parse_plan_with_fences_and_derived_column(self):
        p = V.parse_plan('Sure!\n```json\n{"title": "Tone by area", "derive": [{"name": "Tone!", "kind": "label", "from": "completion", "labels": ["calm", "curt"]}],'
                         ' "chart": {"type": "stack", "x": "area", "color": "tone"}, "follow_ups": ["a"]}\n```', self.cols)
        self.assertEqual(p["derive"][0]["name"], "tone")
        self.assertEqual(p["chart"]["agg"], "share")

    def test_derived_name_never_overwrites_an_original_column(self):
        p = V.parse_plan('{"title": "t", "derive": [{"name": "area", "kind": "label", "from": "prompt", "labels": ["a", "b"]}],'
                         ' "chart": {"type": "bar", "x": "area"}}', self.cols)
        self.assertEqual(p["derive"][0]["name"], "area (model)")
        self.assertEqual(p["chart"]["x"], "area (model)")

    def test_snake_case_names_still_match_the_chart(self):
        p = V.parse_plan('{"derive": [{"name": "step_by_step", "kind": "pattern", "from": "completion", "pattern": "(?i)step \\\\d"}],'
                         ' "chart": {"type": "bar", "x": "step_by_step"}}', self.cols)
        self.assertEqual(p["derive"][0]["name"], "step by step")
        self.assertEqual(p["chart"]["x"], "step by step")
        self.assertEqual(p["derive"][0]["pattern"], "step \\d")

    def test_dangerous_pattern_rejected(self):
        with self.assertRaises(V.PlanError):
            V.parse_plan('{"derive": [{"name": "x", "kind": "pattern", "from": "prompt", "pattern": "(a+)+$"}], "chart": {"type": "bar", "x": "x"}}', self.cols)

    def test_label_reading(self):
        d = {"kind": "label", "labels": ["Formal", "Friendly"]}
        self.assertEqual(V._read_labels('{"1": "formal", "2": "friendly tone"}', d, 2), {1: "Formal", 2: "Friendly"})
        self.assertIsNone(V._read_labels('{"1": "angry"}', d, 2))

    def test_pattern_fill_needs_no_model(self):
        rows, _ = V.prepare([{"t": "Step 1 do this. Step 2 do that."}, {"t": "Just a plain answer here."}])
        vals = V.fill({"name": "steps", "kind": "pattern", "from": "t", "pattern": r"step \d"}, rows, chat=None)
        self.assertEqual(vals, {1: "yes", 2: "no"})


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