from __future__ import annotations

from flownet.core.exceptions import FlowDefinitionError


def require_dict(value: object, field: str) -> dict:
    if not isinstance(value, dict):
        raise FlowDefinitionError(f"{field} must be a dict.")
    return value


def require_non_empty_string(value: object, field: str) -> str:
    if not isinstance(value, str) or not value.strip():
        raise FlowDefinitionError(f"{field} must be a non-empty string.")
    return value


def require_positive_int(value: object, field: str) -> int:
    try:
        parsed = int(value)
    except (TypeError, ValueError) as exc:
        raise FlowDefinitionError(f"{field} must be an integer > 0.") from exc
    if parsed <= 0:
        raise FlowDefinitionError(f"{field} must be an integer > 0.")
    return parsed


def require_non_negative_int(value: object, field: str) -> int:
    try:
        parsed = int(value)
    except (TypeError, ValueError) as exc:
        raise FlowDefinitionError(f"{field} must be an integer >= 0.") from exc
    if parsed < 0:
        raise FlowDefinitionError(f"{field} must be an integer >= 0.")
    return parsed


def validate_partition_id(partition_id: object, num_partitions: int, *, field: str = "partition_id") -> int:
    parsed = require_non_negative_int(partition_id, field)
    if parsed >= num_partitions:
        raise FlowDefinitionError(
            f"{field}={parsed} out of range for num_partitions={num_partitions}."
        )
    return parsed


def validate_partition_map(partitions: object, *, num_partitions: int, field: str) -> None:
    mapping = require_dict(partitions, field)
    seen: set[int] = set()
    for raw_pid, address in mapping.items():
        pid = require_non_negative_int(raw_pid, f"{field} key")
        if pid >= num_partitions:
            raise FlowDefinitionError(
                f"{field} key {pid} out of range for num_partitions={num_partitions}."
            )
        require_non_empty_string(address, f"{field}[{pid}]")
        seen.add(pid)

    missing = [str(pid) for pid in range(num_partitions) if pid not in seen]
    if missing:
        raise FlowDefinitionError(
            f"{field} is missing partition ids: {', '.join(missing)}."
        )


def validate_shuffle_meta_contract(meta: dict) -> None:
    shuffle = require_dict(meta.get("shuffle"), "shuffle")
    shuffle_num_partitions = require_positive_int(
        shuffle.get("num_partitions", 1), "shuffle.num_partitions"
    )
    shuffle_stage_id = shuffle.get("stage_id")
    if shuffle_stage_id is not None:
        shuffle_stage_id = require_non_empty_string(
            shuffle_stage_id, "shuffle.stage_id"
        )

    sys_meta_obj = meta.get("sys")
    if sys_meta_obj is None:
        return
    sys_meta = require_dict(sys_meta_obj, "sys")

    plan_obj = sys_meta.get("plan")
    if plan_obj is None:
        return
    plan = require_dict(plan_obj, "sys.plan")

    plan_num_partitions = require_positive_int(
        plan.get("num_partitions", shuffle_num_partitions), "sys.plan.num_partitions"
    )
    validate_partition_map(
        plan.get("partitions"),
        num_partitions=plan_num_partitions,
        field="sys.plan.partitions",
    )

    require_non_empty_string(sys_meta.get("plan_id"), "sys.plan_id")
    sys_stage_id = require_non_empty_string(sys_meta.get("stage_id"), "sys.stage_id")

    plan_stage_id = plan.get("stage_id")
    if plan_stage_id is not None:
        plan_stage_id = require_non_empty_string(plan_stage_id, "sys.plan.stage_id")
        if plan_stage_id != sys_stage_id:
            raise FlowDefinitionError("sys.stage_id must match sys.plan.stage_id.")

    if shuffle_stage_id is not None and shuffle_stage_id != sys_stage_id:
        raise FlowDefinitionError("shuffle.stage_id must match sys.stage_id when sys.plan is provided.")

    merge_address = plan.get("merge_address")
    if merge_address is not None:
        require_non_empty_string(merge_address, "sys.plan.merge_address")

    merge_partitions = plan.get("merge_partitions")
    if merge_partitions is not None:
        validate_partition_map(
            merge_partitions,
            num_partitions=plan_num_partitions,
            field="sys.plan.merge_partitions",
        )
