from __future__ import annotations

from dataclasses import dataclass
import hashlib
import json
from typing import Any, Callable

from flownet.api.actor_method_ref import ActorMethodRef
from flownet.core.exceptions import FlowDefinitionError
from flownet.core.request import FlowRequest, TransformationFrame
from flownet.core.stateless import StatelessOpRef
from flownet.core.types import PayloadData
from flownet.engine.context import OperatorContext
from flownet.engine.meta_validation import require_dict, require_positive_int, validate_partition_id

from .base import OperatorAction, OperatorHandler


def _payload_value(payload: PayloadData) -> Any:
    if isinstance(payload, PayloadData):
        if len(payload.args) == 1 and not payload.kwargs:
            return payload.args[0]
        return payload
    return payload


def _stable_key_fragment(value: Any) -> str:
    encoded = json.dumps(value, sort_keys=True, separators=(",", ":"), default=str)
    return hashlib.sha1(encoded.encode("utf-8")).hexdigest()


@dataclass
class StateContext:
    state_services: Any
    namespace: str
    key: str
    stage_id: str
    partition_id: int
    request_id: str
    default_ttl_seconds: float | None = None

    def get(self, default: Any = None) -> Any:
        return self.state_services.get(self.namespace, self.key, default=default)

    def set(self, value: Any, ttl_seconds: float | None = None) -> bool:
        ttl = self.default_ttl_seconds if ttl_seconds is None else ttl_seconds
        return bool(
            self.state_services.put(
                self.namespace,
                self.key,
                value,
                ttl_seconds=ttl,
            )
        )

    def update(
        self,
        mutator: Callable[[Any], Any],
        *,
        default: Any = None,
        ttl_seconds: float | None = None,
    ) -> Any:
        current = self.get(default=default)
        next_value = mutator(current)
        self.set(next_value, ttl_seconds=ttl_seconds)
        return next_value

    def clear(self) -> bool:
        return bool(self.state_services.delete(self.namespace, self.key))


class StatefulProcessOperatorHandler(OperatorHandler):
    async def handle_request(
        self,
        ctx: OperatorContext,
        request: FlowRequest,
        transformation_frame: TransformationFrame,
    ) -> OperatorAction:
        meta = ctx.current_frame_meta or {}
        process_meta = require_dict(meta.get("stateful_process"), "stateful_process")
        control = str(meta.get("control", "data"))
        if control == "done":
            await ctx.post_process_result(request, None, transformation_frame.payload)
            return OperatorAction.RETURN
        if control != "data":
            raise FlowDefinitionError(
                "stateful_process control must be one of: data, done."
            )

        if ctx.state_services is None:
            raise FlowDefinitionError("Runtime state services are not configured.")

        stage_id = str(process_meta.get("stage_id") or transformation_frame.transformation.trans_id)
        num_partitions = require_positive_int(
            process_meta.get("num_partitions", 1),
            "stateful_process.num_partitions",
        )
        partition_id = validate_partition_id(
            meta.get("partition_id", 0),
            num_partitions,
            field="partition_id",
        )
        namespace = str(
            process_meta.get("state_namespace") or f"stateful_process:{stage_id}"
        )
        default_ttl = process_meta.get("state_ttl_s")
        ttl_seconds: float | None = None
        if default_ttl is not None:
            ttl_seconds = float(default_ttl)
            if ttl_seconds <= 0:
                raise FlowDefinitionError("stateful_process.state_ttl_s must be > 0 when provided.")

        state_key = self._resolve_state_key(
            payload=transformation_frame.payload,
            frame_meta=meta,
            process_meta=process_meta,
            partition_id=partition_id,
        )
        state_context = StateContext(
            state_services=ctx.state_services,
            namespace=namespace,
            key=state_key,
            stage_id=stage_id,
            partition_id=partition_id,
            request_id=request.request_id,
            default_ttl_seconds=ttl_seconds,
        )

        payload_value = _payload_value(transformation_frame.payload)
        invocation_payload = PayloadData(args=[payload_value, state_context], kwargs={})
        target = transformation_frame.transformation.target
        if isinstance(target, ActorMethodRef):
            success, result = await ctx.call_operator(
                target,
                invocation_payload,
                request,
                transformation_frame.transformation,
            )
        elif isinstance(target, StatelessOpRef):
            success, result = await ctx.call_stateless(
                target,
                invocation_payload,
                request,
                transformation_frame.transformation,
            )
        else:
            raise FlowDefinitionError(
                "process_with_state() target must be ActorMethodRef or StatelessOpRef."
            )
        if not success:
            return OperatorAction.RETURN

        await ctx.post_process_result(request, result, transformation_frame.payload)
        return OperatorAction.RETURN

    @staticmethod
    def _resolve_state_key(
        *,
        payload: PayloadData,
        frame_meta: dict[str, Any],
        process_meta: dict[str, Any],
        partition_id: int,
    ) -> str:
        state_spec = process_meta.get("state_spec")
        key_source: Any = None
        if isinstance(state_spec, dict):
            payload_value = _payload_value(payload)
            key_field = state_spec.get("key_field")
            if key_field is not None and isinstance(payload_value, dict):
                if key_field not in payload_value:
                    raise FlowDefinitionError(
                        f"stateful_process.state_spec.key_field '{key_field}' is missing in payload."
                    )
                key_source = payload_value[key_field]
            key_index = state_spec.get("key_index")
            if key_source is None and key_index is not None and isinstance(payload_value, (list, tuple)):
                idx = int(key_index)
                if idx < 0 or idx >= len(payload_value):
                    raise FlowDefinitionError("stateful_process.state_spec.key_index out of range.")
                key_source = payload_value[idx]

        if key_source is None:
            key_source = frame_meta.get("state_partition_key", partition_id)
        return f"p{partition_id}:{_stable_key_fragment(key_source)}"


__all__ = ["StateContext", "StatefulProcessOperatorHandler"]
