#!/usr/bin/env python3

"""Load and validate the declarative packed-backup workflow specification."""

import re
from dataclasses import dataclass
from pathlib import Path
from typing import Any, Dict, Mapping, Optional, Sequence, Tuple
from urllib.parse import urlsplit

import yaml
from yaml.constructor import ConstructorError


SCHEMA_VERSION = 1
MAX_SPEC_BYTES = 64 * 1024
DEFAULT_TTL_SECONDS_AFTER_FINISHED = 30 * 24 * 60 * 60


class SpecError(ValueError):
    """The workflow spec is missing, malformed, or internally inconsistent."""


class UniqueKeySafeLoader(yaml.SafeLoader):
    """SafeLoader variant that rejects duplicate mapping keys."""


def _construct_unique_mapping(
    loader: UniqueKeySafeLoader,
    node: yaml.MappingNode,
    deep: bool = False,
) -> Dict[Any, Any]:
    mapping: Dict[Any, Any] = {}
    for key_node, value_node in node.value:
        key = loader.construct_object(key_node, deep=deep)
        try:
            duplicate = key in mapping
        except TypeError as exc:
            raise ConstructorError(
                "while constructing a mapping",
                node.start_mark,
                "found an unhashable mapping key",
                key_node.start_mark,
            ) from exc
        if duplicate:
            raise ConstructorError(
                "while constructing a mapping",
                node.start_mark,
                f"found duplicate key {key!r}",
                key_node.start_mark,
            )
        mapping[key] = loader.construct_object(value_node, deep=deep)
    return mapping


UniqueKeySafeLoader.add_constructor(
    yaml.resolver.BaseResolver.DEFAULT_MAPPING_TAG,
    _construct_unique_mapping,
)


@dataclass(frozen=True)
class ResourceSpec:
    cpu_request: str
    memory_request: str

    def as_dict(self) -> Dict[str, str]:
        return {
            "cpuRequest": self.cpu_request,
            "memoryRequest": self.memory_request,
        }


@dataclass(frozen=True)
class MysqlSpec:
    secret_name: str
    secret_key: str

    def as_dict(self) -> Dict[str, str]:
        return {
            "secretName": self.secret_name,
            "secretKey": self.secret_key,
        }


@dataclass(frozen=True)
class PackSpec:
    shards: int
    cse_ctl_image: str
    resources: ResourceSpec
    ebs_request: Optional[str]
    ebs_storage_class: Optional[str]
    generator_args: Tuple[str, ...]
    command_args: Tuple[str, ...]

    def as_dict(self) -> Dict[str, Any]:
        payload: Dict[str, Any] = {
            "shards": self.shards,
            "cseCtlImage": self.cse_ctl_image,
            "resources": self.resources.as_dict(),
            "generatorArgs": list(self.generator_args),
            "commandArgs": list(self.command_args),
        }
        if self.ebs_request is not None or self.ebs_storage_class is not None:
            ebs: Dict[str, str] = {}
            if self.ebs_request is not None:
                ebs["request"] = self.ebs_request
            if self.ebs_storage_class is not None:
                ebs["storageClass"] = self.ebs_storage_class
            payload["ebs"] = ebs
        return payload


@dataclass(frozen=True)
class ExportSpec:
    out_prefix: str
    cse_ctl_image: str
    dumpling_image: str
    runner_image: Optional[str]
    resources: ResourceSpec
    exclude_meta_name_regexp: Optional[str]
    generator_args: Tuple[str, ...]
    worker_args: Tuple[str, ...]
    dumpling_args: Tuple[str, ...]

    def as_dict(self) -> Dict[str, Any]:
        payload: Dict[str, Any] = {
            "outPrefix": self.out_prefix,
            "cseCtlImage": self.cse_ctl_image,
            "dumplingImage": self.dumpling_image,
            "resources": self.resources.as_dict(),
            "generatorArgs": list(self.generator_args),
            "workerArgs": list(self.worker_args),
            "dumplingArgs": list(self.dumpling_args),
        }
        if self.runner_image is not None:
            payload["runnerImage"] = self.runner_image
        if self.exclude_meta_name_regexp is not None:
            payload["excludeMetaNameRegexp"] = self.exclude_meta_name_regexp
        return payload


@dataclass(frozen=True)
class WatermarkSpec:
    image: str
    aws_region: str
    resources: ResourceSpec
    generator_args: Tuple[str, ...]
    worker_args: Tuple[str, ...]

    def as_dict(self) -> Dict[str, Any]:
        return {
            "image": self.image,
            "awsRegion": self.aws_region,
            "resources": self.resources.as_dict(),
            "generatorArgs": list(self.generator_args),
            "workerArgs": list(self.worker_args),
        }


@dataclass(frozen=True)
class WorkflowSpec:
    node_group: str
    ttl_seconds_after_finished: int
    mysql: MysqlSpec
    pack: PackSpec
    export: ExportSpec
    watermark: WatermarkSpec
    schema_version: int = SCHEMA_VERSION

    def as_dict(self) -> Dict[str, Any]:
        return {
            "schemaVersion": self.schema_version,
            "nodeGroup": self.node_group,
            "ttlSecondsAfterFinished": self.ttl_seconds_after_finished,
            "mysql": self.mysql.as_dict(),
            "pack": self.pack.as_dict(),
            "export": self.export.as_dict(),
            "watermark": self.watermark.as_dict(),
        }

def _mapping(value: Any, path: str) -> Mapping[str, Any]:
    if not isinstance(value, dict):
        raise SpecError(f"{path} must be a mapping")
    if not all(isinstance(key, str) for key in value):
        raise SpecError(f"{path} keys must be strings")
    return value


def _keys(
    value: Mapping[str, Any],
    path: str,
    required: Sequence[str],
    optional: Sequence[str] = (),
) -> None:
    required_set = set(required)
    allowed = required_set | set(optional)
    missing = sorted(required_set - set(value))
    unknown = sorted(set(value) - allowed)
    if missing:
        raise SpecError(f"{path} is missing required field(s): {', '.join(missing)}")
    if unknown:
        raise SpecError(f"{path} contains unknown field(s): {', '.join(unknown)}")


def _string(value: Any, path: str) -> str:
    if not isinstance(value, str) or not value:
        raise SpecError(f"{path} must be a non-empty string")
    return value


def _optional_string(value: Any, path: str) -> Optional[str]:
    if value is None:
        return None
    return _string(value, path)


def _positive_integer(value: Any, path: str) -> int:
    if isinstance(value, bool) or not isinstance(value, int) or value <= 0:
        raise SpecError(f"{path} must be a positive integer")
    return value


def _non_negative_integer(value: Any, path: str) -> int:
    if isinstance(value, bool) or not isinstance(value, int) or value < 0:
        raise SpecError(f"{path} must be a non-negative integer")
    return value


def _argv(value: Any, path: str) -> Tuple[str, ...]:
    if value is None:
        return ()
    if not isinstance(value, list):
        raise SpecError(f"{path} must be an array of strings")
    result = []
    for index, item in enumerate(value):
        if not isinstance(item, str):
            raise SpecError(f"{path}[{index}] must be a string")
        if not item:
            raise SpecError(f"{path}[{index}] must not be empty")
        result.append(item)
    return tuple(result)


def _option_name(argument: str) -> Optional[str]:
    if not argument.startswith("-") or argument == "-":
        return None
    return argument.split("=", 1)[0]


def _reject_reserved(
    arguments: Sequence[str],
    path: str,
    reserved: Sequence[str],
) -> None:
    reserved_set = set(reserved)
    for index, argument in enumerate(arguments):
        option = _option_name(argument)
        if argument == "--":
            raise SpecError(f"{path}[{index}] must not introduce an argv boundary")
        if option in reserved_set:
            raise SpecError(
                f"{path}[{index}] uses controller-owned option {option!r}"
            )


def _resources(value: Any, path: str) -> ResourceSpec:
    payload = _mapping(value, path)
    _keys(payload, path, ("cpuRequest", "memoryRequest"))
    return ResourceSpec(
        cpu_request=_string(payload["cpuRequest"], f"{path}.cpuRequest"),
        memory_request=_string(payload["memoryRequest"], f"{path}.memoryRequest"),
    )


def _kubernetes_name(value: Any, path: str) -> str:
    parsed = _string(value, path)
    if (
        len(parsed) > 63
        or re.fullmatch(r"[a-z0-9]([-a-z0-9]*[a-z0-9])?", parsed) is None
    ):
        raise SpecError(f"{path} must be a Kubernetes DNS label")
    return parsed


def _label_value(value: Any, path: str) -> str:
    parsed = _string(value, path)
    pattern = r"[A-Za-z0-9](?:[-A-Za-z0-9_.]*[A-Za-z0-9])?"
    if len(parsed) > 63 or re.fullmatch(pattern, parsed) is None:
        raise SpecError(f"{path} must be a Kubernetes label value")
    return parsed


def _volume_key(value: Any, path: str) -> str:
    parsed = _string(value, path)
    if len(parsed) > 253 or re.fullmatch(r"[A-Za-z0-9._-]+", parsed) is None:
        raise SpecError(f"{path} must be a valid Secret data key")
    return parsed


def _s3_prefix(value: Any, path: str) -> str:
    parsed = _string(value, path)
    uri = urlsplit(parsed)
    if (
        uri.scheme != "s3"
        or not uri.netloc
        or not uri.path.lstrip("/")
        or uri.query
        or uri.fragment
        or any(character.isspace() for character in parsed)
    ):
        raise SpecError(f"{path} must be an s3://<bucket>/<prefix> URI")
    return parsed


def _mysql(value: Any) -> MysqlSpec:
    payload = _mapping(value, "mysql")
    _keys(payload, "mysql", ("secretName",), ("secretKey",))
    return MysqlSpec(
        secret_name=_kubernetes_name(payload["secretName"], "mysql.secretName"),
        secret_key=_volume_key(
            payload.get("secretKey", "mysql.cnf"),
            "mysql.secretKey",
        ),
    )


PACK_GENERATOR_RESERVED = (
    "--name",
    "--workflow-name",
    "--shards",
    "--image",
    "--node-group",
    "--ttl-seconds-after-finished",
    "--cluster-id-in-file",
    "--cluster-id-config-map",
    "--cluster-id-dfs-object-key",
    "--output",
    "--cpu.request",
    "--mem.request",
    "--memory.request",
    "--ebs.request",
    "--ebs",
    "--ebs.storage-class",
)
PACK_COMMAND_RESERVED = (
    "--backup-name",
    "--use-id",
    "--data-dir",
    "--keyspace-select.shard-number",
    "--keyspace-select.shard-total",
    "--keyspace-select.cluster-id-in",
)
EXPORT_GENERATOR_RESERVED = (
    "--name",
    "--workflow-name",
    "--batch-task-json",
    "--output",
    "--local-gpg",
    "--runner-image",
    "--dumpling-image",
    "--cse-ctl-image",
    "--node-group",
    "--ttl-seconds-after-finished",
    "--cpu.request",
    "--mem.request",
    "--memory.request",
)
EXPORT_WORKER_RESERVED = (
    "--s3-prefix",
    "--out-prefix",
    "--flat",
    "--gpg",
    "--exclude-meta-name-regexp",
)
EXPORT_DUMPLING_RESERVED = (
    "--packed-backup",
    "--cse-ctl-path",
    "--filetype",
    "--output",
    "--s3.region",
    "--status-addr",
    "--gpg-key-file",
)
WATERMARK_GENERATOR_RESERVED = (
    "--name",
    "--workflow-name",
    "--tasks-json",
    "--mysql-secret",
    "--mysql-secret-key",
    "--out-prefix",
    "--export-instance",
    "--backup-time",
    "--aws-region",
    "--runner-image",
    "--node-group",
    "--ttl-seconds-after-finished",
    "--cpu.request",
    "--memory.request",
)
WATERMARK_WORKER_RESERVED = (
    "--watermark-instance",
    "--export-instance",
    "--tasks-json",
    "--out-prefix",
    "--backup-time",
    "--mysql-defaults-file",
    "--aws-region",
)


def _pack(value: Any) -> PackSpec:
    payload = _mapping(value, "pack")
    _keys(
        payload,
        "pack",
        ("shards", "cseCtlImage", "resources"),
        ("ebs", "generatorArgs", "commandArgs"),
    )
    generator_args = _argv(payload.get("generatorArgs"), "pack.generatorArgs")
    command_args = _argv(payload.get("commandArgs"), "pack.commandArgs")
    _reject_reserved(
        generator_args,
        "pack.generatorArgs",
        PACK_GENERATOR_RESERVED,
    )
    _reject_reserved(command_args, "pack.commandArgs", PACK_COMMAND_RESERVED)

    ebs_request = None
    ebs_storage_class = None
    if "ebs" in payload:
        ebs = _mapping(payload["ebs"], "pack.ebs")
        _keys(ebs, "pack.ebs", (), ("request", "storageClass"))
        ebs_request = _optional_string(ebs.get("request"), "pack.ebs.request")
        ebs_storage_class = _optional_string(
            ebs.get("storageClass"),
            "pack.ebs.storageClass",
        )
        if ebs_storage_class is not None and ebs_request is None:
            raise SpecError("pack.ebs.storageClass requires pack.ebs.request")

    return PackSpec(
        shards=_positive_integer(payload["shards"], "pack.shards"),
        cse_ctl_image=_string(payload["cseCtlImage"], "pack.cseCtlImage"),
        resources=_resources(payload["resources"], "pack.resources"),
        ebs_request=ebs_request,
        ebs_storage_class=ebs_storage_class,
        generator_args=generator_args,
        command_args=command_args,
    )


def _export(value: Any) -> ExportSpec:
    payload = _mapping(value, "export")
    _keys(
        payload,
        "export",
        (
            "outPrefix",
            "cseCtlImage",
            "dumplingImage",
            "resources",
        ),
        (
            "runnerImage",
            "excludeMetaNameRegexp",
            "generatorArgs",
            "workerArgs",
            "dumplingArgs",
        ),
    )
    generator_args = _argv(
        payload.get("generatorArgs"),
        "export.generatorArgs",
    )
    worker_args = _argv(payload.get("workerArgs"), "export.workerArgs")
    dumpling_args = _argv(payload.get("dumplingArgs"), "export.dumplingArgs")
    _reject_reserved(
        generator_args,
        "export.generatorArgs",
        EXPORT_GENERATOR_RESERVED,
    )
    _reject_reserved(worker_args, "export.workerArgs", EXPORT_WORKER_RESERVED)
    _reject_reserved(
        dumpling_args,
        "export.dumplingArgs",
        EXPORT_DUMPLING_RESERVED,
    )

    return ExportSpec(
        out_prefix=_s3_prefix(payload["outPrefix"], "export.outPrefix"),
        cse_ctl_image=_string(payload["cseCtlImage"], "export.cseCtlImage"),
        dumpling_image=_string(payload["dumplingImage"], "export.dumplingImage"),
        runner_image=_optional_string(
            payload.get("runnerImage"),
            "export.runnerImage",
        ),
        resources=_resources(payload["resources"], "export.resources"),
        exclude_meta_name_regexp=_optional_string(
            payload.get("excludeMetaNameRegexp"),
            "export.excludeMetaNameRegexp",
        ),
        generator_args=generator_args,
        worker_args=worker_args,
        dumpling_args=dumpling_args,
    )


def _watermark(value: Any) -> WatermarkSpec:
    payload = _mapping(value, "watermark")
    _keys(
        payload,
        "watermark",
        ("image", "awsRegion", "resources"),
        ("generatorArgs", "workerArgs"),
    )
    generator_args = _argv(
        payload.get("generatorArgs"),
        "watermark.generatorArgs",
    )
    worker_args = _argv(payload.get("workerArgs"), "watermark.workerArgs")
    _reject_reserved(
        generator_args,
        "watermark.generatorArgs",
        WATERMARK_GENERATOR_RESERVED,
    )
    _reject_reserved(
        worker_args,
        "watermark.workerArgs",
        WATERMARK_WORKER_RESERVED,
    )
    return WatermarkSpec(
        image=_string(payload["image"], "watermark.image"),
        aws_region=_string(payload["awsRegion"], "watermark.awsRegion"),
        resources=_resources(payload["resources"], "watermark.resources"),
        generator_args=generator_args,
        worker_args=worker_args,
    )


def parse_workflow_spec(payload: Any) -> WorkflowSpec:
    root = _mapping(payload, "workflow spec")
    _keys(
        root,
        "workflow spec",
        ("schemaVersion", "mysql", "pack", "export", "watermark"),
        ("nodeGroup", "ttlSecondsAfterFinished"),
    )
    version = root["schemaVersion"]
    if isinstance(version, bool) or version != SCHEMA_VERSION:
        raise SpecError(
            f"schemaVersion must be the integer {SCHEMA_VERSION}"
        )
    return WorkflowSpec(
        schema_version=version,
        node_group=_label_value(
            root.get("nodeGroup", "general"),
            "nodeGroup",
        ),
        ttl_seconds_after_finished=_non_negative_integer(
            root.get(
                "ttlSecondsAfterFinished",
                DEFAULT_TTL_SECONDS_AFTER_FINISHED,
            ),
            "ttlSecondsAfterFinished",
        ),
        mysql=_mysql(root["mysql"]),
        pack=_pack(root["pack"]),
        export=_export(root["export"]),
        watermark=_watermark(root["watermark"]),
    )


def load_workflow_spec(path: str) -> WorkflowSpec:
    spec_path = Path(path)
    try:
        raw = spec_path.read_bytes()
    except OSError as exc:
        raise SpecError(
            f"cannot read workflow spec {str(spec_path)!r} "
            f"({type(exc).__name__})"
        ) from exc
    if not raw:
        raise SpecError("workflow spec is empty")
    if len(raw) > MAX_SPEC_BYTES:
        raise SpecError(
            f"workflow spec exceeds {MAX_SPEC_BYTES} bytes"
        )
    try:
        text = raw.decode("utf-8")
    except UnicodeDecodeError as exc:
        raise SpecError("workflow spec is not valid UTF-8") from exc
    try:
        payload = yaml.load(text, Loader=UniqueKeySafeLoader)
    except yaml.YAMLError as exc:
        raise SpecError(f"workflow spec contains invalid YAML: {exc}") from exc
    return parse_workflow_spec(payload)