from __future__ import annotations

import threading
import uuid
from typing import TYPE_CHECKING, Any, Callable, Mapping

from flownet.core.request import FlowAckRoute
from flownet.runtime.contexts.flow import current_flow_request

if TYPE_CHECKING:
    from flownet.api.actor_method_ref import ActorMethodRef
    from flownet.api.datastream import DataStream


class RunTopicEgressChannel:
    """
    Runtime-managed topic egress adapter.

    It wraps a flow target and publishes emitted payloads to a configured topic.
    Request lifecycle events are also surfaced as topic events (`done`/`error`).
    """

    def __init__(
        self,
        target: "ActorMethodRef",
        *,
        topic_id: str,
        event_type: str = "data",
        done_event_type: str = "done",
        error_event_type: str = "error",
        tags: Mapping[str, str] | None = None,
        runtime_managed: bool = True,
    ):
        self.target = target
        self.topic_id = str(topic_id)
        self.event_type = str(event_type or "data")
        self.done_event_type = str(done_event_type or "done")
        self.error_event_type = str(error_event_type or "error")
        self.runtime_managed = runtime_managed
        self._base_tags = {str(k): str(v) for k, v in dict(tags or {}).items()}
        self._pending_requests: dict[str, int] = {}
        self._completed_requests: set[str] = set()
        self._request_seq: dict[str, int] = {}
        self._request_tracking_origin: dict[str, str] = {}
        self._pending_lock = threading.Lock()
        self._actor_id: str | None = None
        self._actor_ref = None
        self._feed_ref = None
        self._runtime_start_request: Callable[[str, int], Any] | None = None
        self._runtime_ack_request_delta: Callable[[str, int], Any] | None = None
        self._runtime_next_request_meta: Callable[[str], dict[str, Any]] | None = None
        self._runtime_finish_request: Callable[[str], Any] | None = None
        self._runtime_close_ingress: Callable[[], bool] | None = None
        self._runtime_request_done: Callable[[str], bool] | None = None
        self.output_task = self._build_output_task()

    def _build_output_task(self):
        from flownet.compiler.flow_def import FlowDef

        return FlowDef(type(self)._output_task_impl, static=True).__get__(self, type(self))

    def _output_task_impl(self, init_stream: "DataStream"):
        return init_stream.pipe(self.target).map(self._get_feed_ref())

    def _get_actor_id(self) -> str:
        if self._actor_id is None:
            from flownet.api.runtime import make_object_actor

            self._actor_ref = make_object_actor(self)
            self._actor_id = self._actor_ref.actor_id
        return self._actor_id

    def _get_actor_address(self) -> str | None:
        actor_ref = self._actor_ref
        if actor_ref is None:
            return None
        normalized = str(getattr(actor_ref, "address", "") or "").strip()
        return normalized or None

    def _get_feed_ref(self):
        if self._feed_ref is None:
            from flownet.api.actor_method_ref import ActorMethodRef

            actor_ref = self._actor_ref
            if actor_ref is None:
                from flownet.api.runtime import make_object_actor

                actor_ref = make_object_actor(self)
                self._actor_ref = actor_ref
                self._actor_id = actor_ref.actor_id
            self._feed_ref = ActorMethodRef(
                address=actor_ref.address,
                actor_id=actor_ref.actor_id,
                method="feed",
            )
        return self._feed_ref

    def ack_route(self) -> FlowAckRoute:
        return FlowAckRoute(
            coordinator_target=self._get_actor_id(),
            coordinator_address=self._get_actor_address(),
        )

    def feed(self, *args, **kwargs):
        request = current_flow_request()
        if request is None:
            return
        self._remember_request_origin(request.request_id, request.origin_address)
        with self._pending_lock:
            next_seq = self._request_seq.get(request.request_id, 0) + 1
            self._request_seq[request.request_id] = next_seq
        payload = args[0] if len(args) == 1 and not kwargs else (args, kwargs)
        self._publish(
            request_id=request.request_id,
            event_type=self.event_type,
            payload=payload,
            seq=next_seq,
        )

    def write(
        self,
        *args,
        request_id: str | None = None,
        autocommit: bool | None = None,
        **kwargs,
    ):
        new_request = request_id is None
        if new_request:
            request_id = f"{self._get_actor_id()}:{uuid.uuid4().hex}"
        elif not isinstance(request_id, str) or not request_id:
            raise ValueError("request_id must be a non-empty string when provided.")
        elif not self._has_pending_request(request_id):
            raise ValueError(
                "Custom runtime request_id injection is not allowed. "
                "Call write() without request_id to start a request and reuse the returned id.",
            )

        assert request_id is not None
        with self._pending_lock:
            self._pending_requests[request_id] = self._pending_requests.get(request_id, 0) + 1
        if self._runtime_start_request is not None:
            self._runtime_start_request(request_id, 1)
        next_meta: dict[str, Any] | None = None
        if self._runtime_next_request_meta is not None:
            next_meta = self._runtime_next_request_meta(request_id)
        dispatched_request_id = self.output_task.put_request(
            request_id,
            *args,
            ack_route=self.ack_route(),
            _next_meta=next_meta,
            **kwargs,
        )
        if autocommit is None:
            autocommit = new_request
        if autocommit:
            self.finish(request_id)
        return dispatched_request_id

    def finish(self, request_id: str) -> bool:
        if not request_id:
            return False
        if self._runtime_finish_request is not None:
            finished = bool(self._runtime_finish_request(request_id))
            if finished:
                self.event_chain_done(request_id, 0)
            return finished
        return False

    def close(self) -> bool:
        closed = False
        if self._runtime_close_ingress is not None:
            closed = bool(self._runtime_close_ingress())
        if closed:
            with self._pending_lock:
                pending_ids = tuple(self._pending_requests.keys())
            for request_id in pending_ids:
                self.event_chain_done(request_id, 0)
        return closed

    def close_ingress(self) -> bool:
        return self.close()

    def event_chain_done(self, request_id: str, delta: int):
        runtime_pending: int | None = None
        if self._runtime_ack_request_delta is not None:
            runtime_pending = self._runtime_ack_request_delta(request_id, delta)
        done = False
        final_seq = 0
        with self._pending_lock:
            if request_id in self._completed_requests:
                return
            pending = int(self._pending_requests.get(request_id, 0))
            if pending == 0 and delta > 0:
                pending = int(delta)
            else:
                pending += int(delta)

            if runtime_pending is not None:
                runtime_pending_int = int(runtime_pending)
                if runtime_pending_int > 0:
                    pending = runtime_pending_int
                elif runtime_pending_int == 0 and pending <= 0:
                    pending = 0

            if pending <= 0:
                request_done = True
                if self._runtime_request_done is not None:
                    request_done = bool(self._runtime_request_done(request_id))
                if request_done:
                    self._pending_requests.pop(request_id, None)
                    self._completed_requests.add(request_id)
                    final_seq = int(self._request_seq.pop(request_id, 0))
                    done = True
                else:
                    # Keep request open until producer_done arrives.
                    self._pending_requests[request_id] = 0
            else:
                self._pending_requests[request_id] = pending
        if done:
            self._publish(
                request_id=request_id,
                event_type=self.done_event_type,
                payload={
                    "request_id": request_id,
                    "request_ref_id": request_id,
                    "reason": "runtime_done",
                    "final_seq": final_seq,
                    "expected_total_events": final_seq,
                    "request_done": True,
                    "status": "completed",
                    "subscriber_done": True,
                },
                final_seq=final_seq,
            )
            self._clear_request_tracking_origin(request_id)

    def flow_error(
        self,
        request_id: str,
        error_type: str,
        message: str,
        traceback_text: str,
        event: Any | None = None,
    ):
        final_seq = 0
        with self._pending_lock:
            pending = self._pending_requests.pop(request_id, None)
            if request_id in self._completed_requests:
                return
            self._completed_requests.add(request_id)
            final_seq = int(self._request_seq.pop(request_id, 0))
        if pending is not None and self._runtime_ack_request_delta is not None:
            self._runtime_ack_request_delta(request_id, -pending)

        self._publish(
            request_id=request_id,
            event_type=self.error_event_type,
            payload={
                "request_id": request_id,
                "request_ref_id": request_id,
                "error_type": error_type,
                "message": message,
                "traceback_text": traceback_text,
                "event": event,
            },
            final_seq=final_seq,
        )
        self._publish(
            request_id=request_id,
            event_type=self.done_event_type,
            payload={
                "request_id": request_id,
                "request_ref_id": request_id,
                "reason": "runtime_error",
                "final_seq": final_seq,
                "expected_total_events": final_seq,
                "request_done": True,
                "status": "failed",
                "subscriber_done": True,
            },
            final_seq=final_seq,
        )
        self._clear_request_tracking_origin(request_id)

    def _has_pending_request(self, request_id: str) -> bool:
        with self._pending_lock:
            return request_id in self._pending_requests

    def _publish(
        self,
        *,
        request_id: str,
        event_type: str,
        payload: Any,
        seq: int | None = None,
        final_seq: int | None = None,
    ) -> None:
        from flownet.runtime.kernel import get_runtime_kernel

        tags = dict(self._base_tags)
        tags["request_id"] = request_id
        tags.setdefault("request_ref_id", request_id)
        tags.setdefault("__track_request_root__", "1")
        tracking_origin = self._resolve_request_tracking_origin(request_id)
        if tracking_origin is not None:
            tags.setdefault("__request_tracker_origin__", tracking_origin)
        if seq is not None and seq >= 0:
            tags["__seq__"] = str(int(seq))
        if final_seq is not None and final_seq >= 0:
            tags["__final_seq__"] = str(int(final_seq))
        get_runtime_kernel().require_topics().emit_topic_event(
            self.topic_id,
            event_type,
            payload,
            tags,
        )

    def _remember_request_origin(self, request_id: str, origin_address: str | None) -> None:
        normalized_request_id = str(request_id or "").strip()
        normalized_origin = str(origin_address or "").strip()
        if not normalized_request_id or not normalized_origin:
            return
        with self._pending_lock:
            self._request_tracking_origin.setdefault(normalized_request_id, normalized_origin)

    def _resolve_request_tracking_origin(self, request_id: str) -> str | None:
        normalized_request_id = str(request_id or "").strip()
        if not normalized_request_id:
            return None
        with self._pending_lock:
            return self._request_tracking_origin.get(normalized_request_id)

    def _clear_request_tracking_origin(self, request_id: str) -> None:
        normalized_request_id = str(request_id or "").strip()
        if not normalized_request_id:
            return
        with self._pending_lock:
            self._request_tracking_origin.pop(normalized_request_id, None)

    def bind_runtime_tracker(
        self,
        *,
        start_request: Callable[[str, int], Any] | None = None,
        ack_request_delta: Callable[[str, int], Any] | None = None,
        next_request_meta: Callable[[str], dict[str, Any]] | None = None,
        finish_request: Callable[[str], Any] | None = None,
        close_ingress: Callable[[], bool] | None = None,
        request_done: Callable[[str], bool] | None = None,
    ) -> None:
        self._runtime_start_request = start_request
        self._runtime_ack_request_delta = ack_request_delta
        self._runtime_next_request_meta = next_request_meta
        self._runtime_finish_request = finish_request
        self._runtime_close_ingress = close_ingress
        self._runtime_request_done = request_done


__all__ = ["RunTopicEgressChannel"]
