from __future__ import annotations

from collections import defaultdict
from typing import Any

from flownet.core.exceptions import FlowDefinitionError
from flownet.core.request import FlowRequest, TransformationFrame
from flownet.core.types import PayloadData
from flownet.engine.context import OperatorContext
from flownet.engine.meta_validation import (
    require_dict,
    require_positive_int,
    validate_shuffle_meta_contract,
)

from .base import OperatorAction
from .reducers.common import _OperatorStats
from .router import RouterOperatorHandler


class ShuffleRouterOperatorHandler(RouterOperatorHandler):
    """Internal shuffle router that partitions records to reducer shards."""

    def __init__(self) -> None:
        self._stats = _OperatorStats()
        self._partition_load: defaultdict[tuple[str, int], int] = defaultdict(int)

    def stats_snapshot(self) -> dict[str, int]:
        snapshot = self._stats.snapshot()
        if not self._partition_load:
            return snapshot
        loads = list(self._partition_load.values())
        snapshot["partition_load_entries"] = len(loads)
        snapshot["partition_load_max"] = max(loads)
        snapshot["partition_load_min"] = min(loads)
        snapshot["partition_load_skew"] = max(loads) - min(loads)
        return snapshot

    async def handle_request(
        self,
        ctx: OperatorContext,
        request: FlowRequest,
        transformation_frame: TransformationFrame,
    ) -> OperatorAction:
        meta = ctx.current_frame_meta or {}
        validate_shuffle_meta_contract(meta)
        shuffle = require_dict(meta.get("shuffle"), "shuffle")

        sys_meta = meta.get("sys") or {}
        plan = sys_meta.get("plan") if isinstance(sys_meta, dict) else None
        plan_partitions = plan.get("partitions") if isinstance(plan, dict) else None
        plan_merge_partitions = plan.get("merge_partitions") if isinstance(plan, dict) else None

        num_partitions = require_positive_int(
            shuffle.get("num_partitions", 1), "shuffle.num_partitions"
        )
        if isinstance(plan, dict) and plan.get("num_partitions") is not None:
            num_partitions = require_positive_int(
                plan.get("num_partitions"), "sys.plan.num_partitions"
            )

        next_meta = meta.get("next_meta")
        if not isinstance(next_meta, dict):
            next_meta = {}
        control = str(next_meta.get("control", "data"))

        if control == "done":
            partition_id = self._extract_done_partition_id(
                transformation_frame.payload,
                num_partitions=num_partitions,
            )
            partition_key = partition_id
            self._stats.inc("route_done_total")
        else:
            try:
                key_index = int(shuffle.get("key_index", 0))
            except (TypeError, ValueError) as exc:
                raise FlowDefinitionError("shuffle.key_index must be an integer >= 0.") from exc
            if key_index < 0:
                raise FlowDefinitionError("shuffle.key_index must be an integer >= 0.")
            key_fn = shuffle.get("key_fn")
            if key_fn is None:
                key = self._extract_key(transformation_frame.payload, key_index, shuffle.get("key_field"))
            else:
                resolved, key = await self._extract_key_from_key_fn(
                    ctx,
                    request,
                    transformation_frame,
                    key_fn=key_fn,
                )
                if not resolved:
                    self._stats.inc("key_fn_failed_total")
                    return OperatorAction.RETURN
            partition_key = self._normalize_partition_key(key)
            partition_id = hash(partition_key) % num_partitions
        stage_id = str(shuffle.get("stage_id") or transformation_frame.transformation.trans_id)
        self._stats.inc("route_total")
        self._partition_load[(stage_id, partition_id)] += 1

        # TODO: resolve reducer addresses from cluster scheduler/state service.
        addresses = shuffle.get("addresses") or []
        if not isinstance(addresses, (list, tuple)):
            raise FlowDefinitionError("shuffle.addresses must be a list when provided.")
        address = None
        if isinstance(plan_partitions, dict):
            address = plan_partitions.get(str(partition_id)) or plan_partitions.get(partition_id)
        if address is None:
            address = addresses[partition_id] if partition_id < len(addresses) else ctx.local_address()
        if address == ctx.local_address():
            self._stats.inc("route_local_total")
        else:
            self._stats.inc("route_remote_total")

        route_meta = dict(next_meta)
        route_meta.setdefault("partition_id", partition_id)
        route_meta.setdefault("state_partition_key", partition_key)
        if "reducer" in meta:
            route_meta.setdefault("reducer", meta["reducer"])
            merge_address = None
            if isinstance(plan_merge_partitions, dict):
                merge_address = plan_merge_partitions.get(str(partition_id)) or plan_merge_partitions.get(
                    partition_id
                )
            if merge_address is None and isinstance(plan, dict) and plan.get("merge_address") is not None:
                merge_address = plan.get("merge_address")
            if merge_address is not None:
                reducer_meta = dict(route_meta["reducer"])
                reducer_meta.setdefault("merge_address", merge_address)
                if isinstance(plan_merge_partitions, dict):
                    shard_partitions = [
                        int(pid)
                        for pid, addr in plan_merge_partitions.items()
                        if addr == merge_address
                    ]
                    shard_partitions.sort()
                    reducer_meta.setdefault("merge_partitions", shard_partitions)
                route_meta["reducer"] = reducer_meta
            reducer_meta = route_meta["reducer"]
            if (
                isinstance(reducer_meta, dict)
                and reducer_meta.get("merge_transformation") is not None
                and reducer_meta.get("merge_address") is None
                and not isinstance(plan_merge_partitions, dict)
                and not (isinstance(plan, dict) and plan.get("merge_address") is not None)
            ):
                raise FlowDefinitionError(
                    "reducer.merge_transformation requires merge target metadata "
                    "(reducer.merge_address or sys.plan.merge_address/merge_partitions)."
                )
        if "join" in meta:
            route_meta.setdefault("join", meta["join"])
        if "side" in meta:
            route_meta.setdefault("side", meta["side"])
        if sys_meta:
            route_meta.setdefault("sys", sys_meta)

        # NOTE: This is where a reducer "gets launched":
        # we emit a continuation FlowRequest to the target container address.
        # The remote runtime receives the flowrequest via its network interface,
        # then runs FlowEngine -> operator handler (`sys/reducer`) locally.
        route = {
            "address": address,
            "transformation": shuffle.get("downstream"),
            "payload": transformation_frame.payload,
            "meta": route_meta,
        }

        target_address, next_request = self._build_route_request(
            ctx, request, transformation_frame, route
        )
        await ctx.emit_flow_request(target_address, next_request)
        return OperatorAction.RETURN

    @staticmethod
    def _extract_done_partition_id(
        payload: Any,
        *,
        num_partitions: int,
    ) -> int:
        candidate: Any = payload
        if isinstance(payload, PayloadData):
            if payload.args:
                candidate = payload.args[0]
            elif "partition_id" in payload.kwargs:
                candidate = payload.kwargs["partition_id"]
            else:
                raise FlowDefinitionError(
                    "shuffle done control payload must provide partition_id."
                )
        try:
            partition_id = int(candidate)
        except (TypeError, ValueError) as exc:
            raise FlowDefinitionError(
                "shuffle done control payload partition_id must be an integer."
            ) from exc
        if partition_id < 0 or partition_id >= int(num_partitions):
            raise FlowDefinitionError(
                "shuffle done control payload partition_id out of range."
            )
        return partition_id

    @staticmethod
    async def _extract_key_from_key_fn(
        ctx: OperatorContext,
        request: FlowRequest,
        transformation_frame: TransformationFrame,
        *,
        key_fn: Any,
    ) -> tuple[bool, Any]:
        from flownet.api.actor_method_ref import ActorMethodRef
        from flownet.core.stateless import StatelessOpRef

        if isinstance(key_fn, ActorMethodRef):
            ok, value = await ctx.call_operator(
                key_fn,
                transformation_frame.payload,
                request,
                transformation_frame.transformation,
            )
            if not ok:
                return False, None
            return True, ShuffleRouterOperatorHandler._coerce_key_fn_result(value)

        if isinstance(key_fn, StatelessOpRef):
            ok, value = await ctx.call_stateless(
                key_fn,
                transformation_frame.payload,
                request,
                transformation_frame.transformation,
            )
            if not ok:
                return False, None
            return True, ShuffleRouterOperatorHandler._coerce_key_fn_result(value)

        raise FlowDefinitionError(
            "shuffle.key_fn must be ActorMethodRef or StatelessOpRef when provided."
        )

    @staticmethod
    def _coerce_key_fn_result(value: Any) -> Any:
        if isinstance(value, PayloadData):
            if len(value.args) == 1 and not value.kwargs:
                return value.args[0]
            if not value.args and "key" in value.kwargs:
                return value.kwargs["key"]
        return value

    @staticmethod
    def _extract_key(payload: PayloadData, key_index: int, key_field: str | None) -> Any:
        if isinstance(payload, PayloadData):
            if key_field is not None and key_field in payload.kwargs:
                return payload.kwargs[key_field]
            if payload.args:
                # Common stream path carries one logical record in args[0].
                if len(payload.args) == 1:
                    record = payload.args[0]
                    if key_field is not None and isinstance(record, dict):
                        if key_field not in record:
                            raise FlowDefinitionError(
                                f"shuffle.key_field '{key_field}' is missing in payload."
                            )
                        return record[key_field]
                    if isinstance(record, (list, tuple)):
                        if key_index >= len(record):
                            raise FlowDefinitionError("shuffle.key_index out of range.")
                        return record[key_index]
                    if key_index == 0:
                        return record
                    raise FlowDefinitionError("shuffle.key_index out of range.")

                if key_index >= len(payload.args):
                    raise FlowDefinitionError("shuffle.key_index out of range.")
                return payload.args[key_index]
            if payload.kwargs:
                # fallback: first value
                return next(iter(payload.kwargs.values()))
            return None

        if isinstance(payload, (list, tuple)):
            if key_index >= len(payload):
                raise FlowDefinitionError("shuffle.key_index out of range.")
            return payload[key_index]

        return payload

    @classmethod
    def _normalize_partition_key(cls, key: Any) -> Any:
        """Convert composite/unhashable keys into a deterministic hashable form."""
        if isinstance(key, dict):
            normalized_items = []
            for raw_key, raw_value in key.items():
                normalized_items.append(
                    (
                        cls._normalize_partition_key(raw_key),
                        cls._normalize_partition_key(raw_value),
                    )
                )
            normalized_items.sort(key=lambda item: repr(item[0]))
            return ("dict", tuple(normalized_items))

        if isinstance(key, (list, tuple)):
            return ("seq", tuple(cls._normalize_partition_key(item) for item in key))

        if isinstance(key, set):
            values = [cls._normalize_partition_key(item) for item in key]
            values.sort(key=repr)
            return ("set", tuple(values))

        if isinstance(key, bytearray):
            return bytes(key)

        try:
            hash(key)
        except TypeError as exc:
            raise FlowDefinitionError(
                "shuffle key is not hashable after normalization; "
                f"type={type(key).__name__}. "
                "Use scalar, tuple/list, set, or dict composite keys."
            ) from exc
        return key
