# DOCS_REF:
# - docs/user/05-datastream-distributed-ops.md
# - docs/user/reference/api-surface.md
# DOCS_SYNC:
# - If SchedulePlan/build_plan behavior changes, update distributed-ops docs.

from __future__ import annotations

import hashlib
import uuid
from dataclasses import dataclass
from typing import Iterable


def _stable_hash(value: str) -> int:
    digest = hashlib.sha256(value.encode("utf-8")).hexdigest()
    return int(digest, 16)


def consistent_hash_address(stage_id: str, request_id: str, addresses: Iterable[str]) -> str:
    address_list = list(addresses)
    if not address_list:
        raise ValueError("addresses must not be empty")
    idx = _stable_hash(f"{stage_id}:{request_id}") % len(address_list)
    return address_list[idx]


@dataclass(frozen=True)
class SchedulePlan:
    plan_id: str
    stage_id: str
    num_partitions: int
    partitions: dict[str, str]
    merge_address: str | None
    merge_partitions: dict[str, str] | None = None
    epoch: int = 1

    def as_dict(self) -> dict[str, object]:
        payload = {
            "plan_id": self.plan_id,
            "stage_id": self.stage_id,
            "num_partitions": self.num_partitions,
            "partitions": dict(self.partitions),
            "epoch": self.epoch,
        }
        if self.merge_address is not None:
            payload["merge_address"] = self.merge_address
        if self.merge_partitions is not None:
            payload["merge_partitions"] = dict(self.merge_partitions)
        return payload


class SimpleScheduler:
    """A minimal scheduler that produces deterministic partition/merge plans."""

    def __init__(self, addresses: Iterable[str] | None = None):
        self._addresses = list(addresses or [])

    def update_addresses(self, addresses: Iterable[str]) -> None:
        self._addresses = list(addresses)

    def make_plan(
        self,
        *,
        stage_id: str,
        request_id: str,
        num_partitions: int,
        addresses: Iterable[str] | None = None,
        merge_addresses: Iterable[str] | None = None,
        merge_partitions: dict[str, str] | None = None,
        include_merge_address: bool = True,
        merge_sharded: bool = False,
        epoch: int = 1,
        plan_id: str | None = None,
    ) -> SchedulePlan:
        if num_partitions <= 0:
            raise ValueError("num_partitions must be > 0")
        address_list = list(addresses or self._addresses)
        if not address_list:
            raise ValueError("addresses must not be empty")
        plan_id = plan_id or uuid.uuid4().hex
        partitions = {
            str(pid): address_list[pid % len(address_list)]
            for pid in range(num_partitions)
        }
        merge_address = None
        merge_partition_map = None
        if merge_partitions is not None:
            merge_partition_map = {
                str(pid): address for pid, address in merge_partitions.items()
            }
        elif merge_sharded:
            merge_candidates = list(merge_addresses or address_list)
            if not merge_candidates:
                raise ValueError("merge_addresses must not be empty when merge_sharded=True")
            merge_partition_map = {
                str(pid): merge_candidates[
                    _stable_hash(f"{stage_id}:{request_id}:{pid}") % len(merge_candidates)
                ]
                for pid in range(num_partitions)
            }
        if include_merge_address:
            if merge_partition_map is None:
                merge_candidates = list(merge_addresses or address_list)
                merge_address = consistent_hash_address(stage_id, request_id, merge_candidates)
        return SchedulePlan(
            plan_id=plan_id,
            stage_id=stage_id,
            num_partitions=num_partitions,
            partitions=partitions,
            merge_address=merge_address,
            merge_partitions=merge_partition_map,
            epoch=epoch,
        )


def build_plan(
    *,
    stage_id: str,
    request_id: str,
    num_partitions: int,
    addresses: Iterable[str],
    merge_addresses: Iterable[str] | None = None,
    merge_partitions: dict[str, str] | None = None,
    include_merge_address: bool = True,
    merge_sharded: bool = False,
    epoch: int = 1,
) -> SchedulePlan:
    scheduler = SimpleScheduler(addresses)
    return scheduler.make_plan(
        stage_id=stage_id,
        request_id=request_id,
        num_partitions=num_partitions,
        merge_addresses=merge_addresses,
        merge_partitions=merge_partitions,
        include_merge_address=include_merge_address,
        merge_sharded=merge_sharded,
        epoch=epoch,
    )


__all__ = [
    "SchedulePlan",
    "SimpleScheduler",
    "build_plan",
    "consistent_hash_address",
]
