from __future__ import annotations

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

from flownet.api import flow
from flownet.core.lockfree import lockfree
from flownet.core.request import FlowAckRoute
from flownet.core.types import PayloadData
from flownet.runtime.contexts.flow import current_flow_request
from flownet.runtime.flows.ordered_event_channel import (
    EventEnvelope,
    GapNotice,
    OrderedEventChannel,
    TerminalEnvelope,
)

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


class RunOutputRequest:
    """Request-scoped helper for grouped writes on a shared request_id."""

    def __init__(self, channel: "RunIOChannel", request_id: str | None = None):
        self._channel = channel
        self.request_id = request_id
        self._finished = False

    def write(self, *args, **kwargs) -> str:
        request_id = self.request_id
        if request_id is None:
            request_id = self._channel.write(*args, autocommit=False, **kwargs)
            self.request_id = request_id
            return request_id
        self._channel.write(*args, request_id=request_id, autocommit=False, **kwargs)
        return request_id

    def finish(self) -> bool:
        request_id = self.request_id
        if not request_id:
            return False
        self._finished = self._channel.finish(request_id)
        return self._finished

    def collect(
        self,
        *,
        timeout: float = 5.0,
        late_drain_timeout: float = 1.0,
    ) -> list[Any]:
        request_id = self.request_id
        if not request_id:
            return []
        return self._channel.collect(
            request_id=request_id,
            timeout=timeout,
            late_drain_timeout=late_drain_timeout,
        )

    def __enter__(self) -> "RunOutputRequest":
        return self

    def __exit__(self, exc_type, exc, tb) -> bool:
        if exc_type is None:
            self.finish()
        return False


class RunIOChannel:
    """
    Runtime-managed compatibility output adapter used by buffered/return_once runs.
    This class is intentionally transport-focused: it should not own run lifecycle
    semantics, which belong to FlowRunRef.
    """

    def __init__(
        self,
        target: "ActorMethodRef",
        *,
        runtime_managed: bool = True,
        ordered_capacity: int = 256,
        ordered_backpressure: str = "block",
        ordered_terminal_hold_s: float = 0.05,
    ):
        self.target = target
        self.runtime_managed = runtime_managed

        self.return_buffer = queue.Queue()
        self._ordered_capacity = int(ordered_capacity)
        self._ordered_backpressure = str(ordered_backpressure)
        self._ordered_terminal_hold_s = float(ordered_terminal_hold_s)
        self._request_channels: dict[str, OrderedEventChannel] = {}
        self._buffers_lock = threading.Lock()
        self._pending_requests: dict[str, int] = {}
        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

    @flow
    def output_task(self, init_stream: "DataStream"):
        return init_stream.pipe(self.target).map(self._get_feed_ref())

    def feed(self, *args, **kwargs):
        request = current_flow_request()
        request_id = request.request_id if request is not None else None
        data = PayloadData(args=args, kwargs=kwargs)
        if request_id:
            channel = self._get_request_channel(request_id, create=False)
            if channel is None and self._has_pending_request(request_id):
                channel = self._get_request_channel(request_id, create=True)
            if channel is not None:
                channel.publish_data(data, kind="data")
        self.return_buffer.put(data)

    def _get_request_channel(
        self,
        request_id: str,
        create: bool = True,
    ) -> OrderedEventChannel | None:
        if not request_id:
            return None
        with self._buffers_lock:
            channel = self._request_channels.get(request_id)
            if channel is None and create:
                channel = OrderedEventChannel(
                    request_id=request_id,
                    capacity=self._ordered_capacity,
                    backpressure=self._ordered_backpressure,  # type: ignore[arg-type]
                    terminal_hold_s=self._ordered_terminal_hold_s,
                )
                self._request_channels[request_id] = channel
            return channel

    def _release_request_channel(self, request_id: str):
        if not request_id:
            return
        channel: OrderedEventChannel | None
        with self._buffers_lock:
            channel = self._request_channels.pop(request_id, None)
        if channel is not None:
            channel.close()

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

    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(),
        )

    @lockfree
    def read(self, timeout: float = 5.0, request_id: str | None = None):
        if request_id is None:
            result = self.return_buffer.get(timeout=timeout)
            if isinstance(result, (RunOutputDone, RunOutputError)):
                return result
            return result.unpack()

        channel = self._get_request_channel(request_id, create=True)
        assert channel is not None
        result = channel.read(timeout=timeout)
        if result is None:
            if not self._has_pending_request(request_id):
                self._release_request_channel(request_id)
            raise queue.Empty()
        decoded = self._decode_channel_item(request_id, result)
        if isinstance(decoded, RunOutputError):
            self._release_request_channel(request_id)
        return decoded

    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
        self._get_request_channel(request_id, create=True)
        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 run(
        self,
        *args,
        collect_timeout: float = 5.0,
        late_drain_timeout: float = 1.0,
        **kwargs,
    ) -> list[Any]:
        request_id = self.write(*args, autocommit=True, **kwargs)
        return self.collect(
            request_id=request_id,
            timeout=collect_timeout,
            late_drain_timeout=late_drain_timeout,
        )

    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 collect(
        self,
        *,
        request_id: str,
        timeout: float = 5.0,
        late_drain_timeout: float = 1.0,
    ) -> list[Any]:
        if not request_id:
            raise ValueError("request_id is required for collect().")
        outputs: list[Any] = []
        done_seen = False
        while True:
            try:
                item = self.read(
                    timeout=late_drain_timeout if done_seen else timeout,
                    request_id=request_id,
                )
            except queue.Empty:
                if done_seen:
                    break
                raise
            if isinstance(item, RunOutputDone):
                done_seen = True
                continue
            if isinstance(item, RunOutputError):
                raise RuntimeError(
                    f"{item.error_type}: {item.message}\n{item.traceback_text}"
                )
            if isinstance(item, RunOutputGap):
                continue
            outputs.append(item)
        return outputs

    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)
        channels: list[OrderedEventChannel]
        with self._buffers_lock:
            channels = list(self._request_channels.values())
            self._request_channels.clear()
        for channel in channels:
            channel.close()
        return closed

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

    def request(self, request_id: str | None = None) -> RunOutputRequest:
        return RunOutputRequest(self, request_id=request_id)

    def __enter__(self) -> "RunIOChannel":
        return self

    def __exit__(self, exc_type, exc, tb) -> bool:
        self.close()
        return False

    def subscribe(
        self,
        handler: Callable[[Any], Any],
        *,
        request_id: str,
        from_seq: int | None = None,
        poll_timeout: float = 0.2,
    ):
        if not request_id:
            raise ValueError("request_id is required for subscribe().")
        channel = self._get_request_channel(request_id, create=True)
        assert channel is not None

        def _on_item(item: EventEnvelope | TerminalEnvelope | GapNotice):
            handler(self._decode_channel_item(request_id, item))

        return channel.subscribe(
            _on_item,
            from_seq=from_seq,
            poll_timeout=poll_timeout,
        )

    def request_delivery_done(self, request_id: str) -> bool:
        channel = self._get_request_channel(request_id, create=False)
        if channel is None:
            return not self._has_pending_request(request_id)
        return self._channel_delivery_done(channel)

    def request_delivery_status(self, request_id: str) -> dict[str, Any]:
        channel = self._get_request_channel(request_id, create=False)
        if channel is None:
            return {
                "request_id": request_id,
                "request_ref_id": request_id,
                "delivery_done": not self._has_pending_request(request_id),
                "watermark_done": not self._has_pending_request(request_id),
                "terminal_emitted": False,
                "drained_seq": 0,
                "final_seq": None,
                "stats": {},
            }
        stats = channel.stats()
        terminal_final_seq = channel.terminal_final_seq()
        next_read_seq = int(stats.get("next_read_seq", 1))
        terminal_emitted = channel.terminal_emitted()
        drained_seq = max(0, next_read_seq - 1)
        watermark_done = (
            terminal_final_seq is not None
            and terminal_emitted
            and next_read_seq > terminal_final_seq
        )
        return {
            "request_id": request_id,
            "request_ref_id": request_id,
            "delivery_done": watermark_done,
            "watermark_done": watermark_done,
            "terminal_emitted": terminal_emitted,
            "drained_seq": drained_seq,
            "final_seq": terminal_final_seq,
            "stats": stats,
        }

    @staticmethod
    def _channel_delivery_done(channel: OrderedEventChannel) -> bool:
        terminal_final_seq = channel.terminal_final_seq()
        if terminal_final_seq is None:
            return False
        if not channel.terminal_emitted():
            return False
        stats = channel.stats()
        try:
            next_read_seq = int(stats.get("next_read_seq", 1))
        except (TypeError, ValueError):
            next_read_seq = 1
        return next_read_seq > terminal_final_seq

    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
        with self._pending_lock:
            pending = self._pending_requests.get(request_id)
            if pending is None and runtime_pending is None:
                return
            pending = (pending or 0) + delta
            if runtime_pending is not None:
                pending = int(runtime_pending)
            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)
                    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:
            channel = self._get_request_channel(request_id, create=True)
            assert channel is not None
            final_seq = channel.last_published_seq()
            channel.publish_terminal(final_seq=final_seq, reason="runtime_done")
            done_event = RunOutputDone(
                request_id,
                final_seq=final_seq,
                expected_total_events=final_seq,
                reason="runtime_done",
                request_done=True,
                status="completed",
                subscriber_done=True,
            )
            self.return_buffer.put(done_event)
            return

    def flow_error(
        self,
        request_id: str,
        error_type: str,
        message: str,
        traceback_text: str,
        event: Any | None = None,
    ):
        error = RunOutputError(
            request_id=request_id,
            error_type=error_type,
            message=message,
            traceback_text=traceback_text,
            event=event,
        )
        if request_id:
            with self._pending_lock:
                pending = self._pending_requests.pop(request_id, None)
            if pending is not None and self._runtime_ack_request_delta is not None:
                self._runtime_ack_request_delta(request_id, -pending)
            channel = self._get_request_channel(request_id, create=True)
            if channel is not None:
                channel.publish_data(error, kind="system")
                final_seq = channel.last_published_seq()
                channel.publish_terminal(final_seq=final_seq, reason="runtime_error")
        self.return_buffer.put(error)

    def _decode_channel_item(
        self,
        request_id: str,
        item: EventEnvelope | TerminalEnvelope | GapNotice,
    ) -> Any:
        if isinstance(item, GapNotice):
            return RunOutputGap(
                request_id=request_id,
                seq=item.seq,
                reason=item.reason,
            )
        if isinstance(item, TerminalEnvelope):
            status = "failed" if str(item.reason or "").strip().lower() == "runtime_error" else "completed"
            return RunOutputDone(
                request_id=request_id,
                final_seq=item.final_seq,
                expected_total_events=item.final_seq,
                reason=item.reason,
                request_done=True,
                status=status,
                subscriber_done=True,
            )
        payload = item.payload
        if isinstance(payload, PayloadData):
            return payload.unpack()
        return payload

    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


class RunOutputError:
    def __init__(
        self,
        request_id: str,
        error_type: str,
        message: str,
        traceback_text: str,
        event: Any | None = None,
    ):
        self.request_id = request_id
        self.error_type = error_type
        self.message = message
        self.traceback_text = traceback_text
        self.event = event

    @property
    def request_ref_id(self) -> str:
        """Compatibility alias of request_id for request-ref terminology."""
        return self.request_id


class RunOutputDone:
    def __init__(
        self,
        request_id: str,
        *,
        final_seq: int | None = None,
        expected_total_events: int | None = None,
        reason: str | None = None,
        request_done: bool | None = None,
        status: str | None = None,
        subscriber_done: bool | None = None,
    ):
        self.request_id = request_id
        self.final_seq = final_seq
        self.expected_total_events = expected_total_events
        self.reason = reason
        self.request_done = request_done
        self.status = status
        self.subscriber_done = subscriber_done

    @property
    def request_ref_id(self) -> str:
        """Compatibility alias of request_id for request-ref terminology."""
        return self.request_id


class RunOutputGap:
    def __init__(self, request_id: str, seq: int, reason: str):
        self.request_id = request_id
        self.seq = seq
        self.reason = reason

    @property
    def request_ref_id(self) -> str:
        """Compatibility alias of request_id for request-ref terminology."""
        return self.request_id


# Backward-compatible alias during naming migration from output-only wording
# to explicit run ingress+egress IO semantics.
RunOutputChannel = RunIOChannel


__all__ = [
    "RunIOChannel",
    "RunOutputChannel",
    "RunOutputRequest",
    "RunOutputDone",
    "RunOutputGap",
    "RunOutputError",
]
