import unittest

from src.services.workload_templates import (
    CH_TEMPLATE,
    CH_X_TEMPLATE,
    LONGTERM_TPCC_X_TEMPLATE,
    SYSBENCH_TEMPLATE,
    SYSBENCH_X_TEMPLATE,
    TPCC_X_1K_TEMPLATE,
    TPCC_X_5K_TEMPLATE,
    build_premium_workload_config,
    get_workload_templates_for_plan,
    is_premium_workload_config,
    resolve_tcms_workflow_id,
    resolve_workflow_yaml_path,
)


class WorkloadTemplatePlanTests(unittest.TestCase):
    def test_premium_templates_include_all_x_workloads(self):
        self.assertEqual(
            get_workload_templates_for_plan("Premium"),
            [
                TPCC_X_1K_TEMPLATE,
                TPCC_X_5K_TEMPLATE,
                SYSBENCH_X_TEMPLATE,
                CH_X_TEMPLATE,
                LONGTERM_TPCC_X_TEMPLATE,
            ],
        )

    def test_dedicated_templates_exclude_x_workloads(self):
        templates = get_workload_templates_for_plan("Dedicated")

        self.assertIn(SYSBENCH_TEMPLATE, templates)
        self.assertIn(CH_TEMPLATE, templates)
        self.assertNotIn(SYSBENCH_X_TEMPLATE, templates)
        self.assertNotIn(CH_X_TEMPLATE, templates)

    def test_premium_config_preserves_ch_parameters(self):
        config = {
            "template_type": "bench-ch",
            "params": {"threads": "100", "duration": "2h"},
        }

        result = build_premium_workload_config(config, min_rcu=5000, max_rcu=20000)

        self.assertEqual(result["template_type"], "bench-ch")
        self.assertEqual(result["params"], config["params"])
        self.assertEqual(result["plan_type"], "premium")
        self.assertEqual(result["rcu"], {"min": 5000, "max": 20000})
        self.assertNotIn("plan_type", config)

    def test_premium_config_preserves_node_overrides(self):
        config = {"nodes": {"workload": {"workload": {"params": {"threads": 64}}}}}

        result = build_premium_workload_config(config, min_rcu=1000, max_rcu=4000)

        self.assertEqual(result["nodes"], config["nodes"])
        self.assertTrue(is_premium_workload_config(result))

    def test_legacy_premium_config_is_still_recognized(self):
        self.assertTrue(
            is_premium_workload_config(
                {"template_type": "premium", "rcu": {"min": 1000, "max": 4000}}
            )
        )

    def test_premium_longterm_tpcc_mapping(self):
        self.assertEqual(
            resolve_workflow_yaml_path(LONGTERM_TPCC_X_TEMPLATE),
            "config/perf_x_stability_tpcc.yaml",
        )
        self.assertEqual(resolve_tcms_workflow_id(LONGTERM_TPCC_X_TEMPLATE), 210010)


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