import tempfile
import unittest
from pathlib import Path

import yaml

from packed_backup_workflow_spec import (
    SpecError,
    load_workflow_spec,
    parse_workflow_spec,
)


def valid_payload():
    return {
        "schemaVersion": 1,
        "nodeGroup": "general",
        "ttlSecondsAfterFinished": 86400,
        "mysql": {
            "secretName": "mysql-secret",
            "secretKey": "mysql.cnf",
        },
        "pack": {
            "shards": 2,
            "cseCtlImage": "pack-image",
            "resources": {
                "cpuRequest": "4",
                "memoryRequest": "8Gi",
            },
            "generatorArgs": ["--cpu.limit", "8"],
            "commandArgs": ["--some-pack-option", "value"],
        },
        "export": {
            "outPrefix": "s3://bucket/shared/export",
            "cseCtlImage": "export-image",
            "dumplingImage": "dumpling-image",
            "resources": {
                "cpuRequest": "4",
                "memoryRequest": "8Gi",
            },
            "workerArgs": ["--parallel", "4"],
            "dumplingArgs": ["--threads", "4"],
        },
        "watermark": {
            "image": "watermark-image",
            "awsRegion": "us-east-1",
            "resources": {
                "cpuRequest": "500m",
                "memoryRequest": "1Gi",
            },
        },
    }


class WorkflowSpecTests(unittest.TestCase):
    def test_valid_spec_is_normalized(self):
        first = parse_workflow_spec(valid_payload())
        reordered = {
            key: valid_payload()[key]
            for key in reversed(list(valid_payload()))
        }
        second = parse_workflow_spec(reordered)

        self.assertEqual(first.as_dict(), second.as_dict())
        self.assertEqual(first.node_group, "general")
        self.assertEqual(first.ttl_seconds_after_finished, 86400)
        self.assertEqual(first.watermark.worker_args, ())

    def test_example_spec_loads(self):
        path = (
            Path(__file__).resolve().parents[1]
            / "example"
            / "workflow-spec.example.yaml"
        )
        spec = load_workflow_spec(str(path))
        self.assertEqual(spec.schema_version, 1)
        self.assertEqual(spec.ttl_seconds_after_finished, 2592000)
        self.assertEqual(
            spec.export.out_prefix,
            "s3://<bucket>/<shared-export-prefix>",
        )

    def test_job_ttl_uses_default_and_validates_yaml_integer(self):
        payload = valid_payload()
        del payload["ttlSecondsAfterFinished"]
        self.assertEqual(
            parse_workflow_spec(payload).ttl_seconds_after_finished,
            2592000,
        )

        for value in (-1, True, "2592000"):
            with self.subTest(value=value):
                payload = valid_payload()
                payload["ttlSecondsAfterFinished"] = value
                with self.assertRaisesRegex(SpecError, "non-negative integer"):
                    parse_workflow_spec(payload)

    def test_unknown_field_is_rejected(self):
        payload = valid_payload()
        payload["pack"]["surprise"] = True
        with self.assertRaisesRegex(SpecError, "unknown field"):
            parse_workflow_spec(payload)

    def test_stage_specific_node_group_is_rejected(self):
        for stage in ("pack", "export", "watermark"):
            with self.subTest(stage=stage):
                payload = valid_payload()
                payload[stage]["nodeGroup"] = "other"
                with self.assertRaisesRegex(SpecError, "unknown field"):
                    parse_workflow_spec(payload)

    def test_non_string_argv_is_rejected(self):
        payload = valid_payload()
        payload["export"]["workerArgs"] = ["--parallel", 4]
        with self.assertRaisesRegex(SpecError, r"workerArgs\[1\]"):
            parse_workflow_spec(payload)

    def test_dumpling_timeout_can_be_forwarded_as_worker_argument(self):
        payload = valid_payload()
        payload["export"]["workerArgs"] = [
            "--parallel",
            "4",
            "--dumpling-timeout",
            "3600",
        ]

        spec = parse_workflow_spec(payload)

        self.assertEqual(
            spec.export.worker_args,
            ("--parallel", "4", "--dumpling-timeout", "3600"),
        )

    def test_controller_owned_argument_is_rejected(self):
        payload = valid_payload()
        payload["export"]["workerArgs"] = [
            "--out-prefix=s3://other/value",
        ]
        with self.assertRaisesRegex(SpecError, "controller-owned"):
            parse_workflow_spec(payload)

    def test_cluster_id_object_input_cannot_be_overridden(self):
        payload = valid_payload()
        payload["pack"]["generatorArgs"] = [
            "--cluster-id-dfs-object-key=other/object",
        ]
        with self.assertRaisesRegex(SpecError, "controller-owned"):
            parse_workflow_spec(payload)

    def test_controller_owned_job_ttl_cannot_be_overridden_per_stage(self):
        for stage in ("pack", "export", "watermark"):
            with self.subTest(stage=stage):
                payload = valid_payload()
                payload[stage]["generatorArgs"] = [
                    "--ttl-seconds-after-finished=1"
                ]
                with self.assertRaisesRegex(SpecError, "controller-owned"):
                    parse_workflow_spec(payload)

    def test_controller_owned_dumpling_argument_is_rejected(self):
        payload = valid_payload()
        payload["export"]["dumplingArgs"] = [
            "--output=s3://other/value",
        ]
        with self.assertRaisesRegex(SpecError, "controller-owned"):
            parse_workflow_spec(payload)

    def test_additional_dumpling_filters_remain_supported(self):
        payload = valid_payload()
        payload["export"]["dumplingArgs"] = [
            "--filter",
            "database.table",
        ]
        spec = parse_workflow_spec(payload)
        self.assertEqual(
            spec.export.dumpling_args,
            ("--filter", "database.table"),
        )

    def test_ebs_storage_class_requires_request(self):
        payload = valid_payload()
        payload["pack"]["ebs"] = {"storageClass": "gp3"}
        with self.assertRaisesRegex(SpecError, "requires pack.ebs.request"):
            parse_workflow_spec(payload)

    def test_duplicate_yaml_key_is_rejected(self):
        text = """
schemaVersion: 1
schemaVersion: 1
mysql: {}
pack: {}
export: {}
watermark: {}
"""
        with tempfile.TemporaryDirectory() as directory:
            path = Path(directory) / "workflow.yaml"
            path.write_text(text, encoding="utf-8")
            with self.assertRaisesRegex(SpecError, "duplicate key"):
                load_workflow_spec(str(path))

    def test_yaml_scalar_types_are_not_silently_coerced(self):
        payload = valid_payload()
        payload["pack"]["resources"]["cpuRequest"] = 4
        with self.assertRaisesRegex(SpecError, "non-empty string"):
            parse_workflow_spec(payload)

    def test_yaml_round_trip_keeps_normalized_spec(self):
        payload = valid_payload()
        with tempfile.TemporaryDirectory() as directory:
            path = Path(directory) / "workflow.yaml"
            path.write_text(
                yaml.safe_dump(payload, sort_keys=False),
                encoding="utf-8",
            )
            loaded = load_workflow_spec(str(path))
        self.assertEqual(loaded.as_dict(), parse_workflow_spec(payload).as_dict())


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