from __future__ import annotations

import os
from typing import Callable

from flownet.api import topic_stream, TopicStream
from flownet.api.runtime import create_actor
from flownet.platform.user.startup import DEFAULT_NODE_ADDRESS
from flownet.proactive.registry import service_topic_id


class PlatformClientActor:
    """
    Platform user-level client that subscribes to node runtime services
    and consumes service events via FlowNet's dataflow pipeline.
    """

    def __init__(
        self,
        node_address: str | None = None,
        auto_subscribe: bool = False,
    ):
        self.node_address = (
            node_address
            or os.environ.get("FLOWNET_NODE_ADDR")
            or DEFAULT_NODE_ADDRESS
        )
        self._streams: dict[str, TopicStream] = {}
        if auto_subscribe:
            self.start()

    def add_service_topic(
        self,
        service_name: str,
        build: Callable[[TopicStream], TopicStream] | None = None,
    ) -> TopicStream:
        stream = topic_stream(service_topic_id(service_name))
        if build is not None:
            stream = build(stream)
        self._streams[service_name] = stream
        return stream

    def add_handler(
        self,
        service_name: str,
        handler: Callable,
        payload_only: bool = True,
    ) -> TopicStream:
        stream = topic_stream(service_topic_id(service_name))
        if payload_only:
            stream = stream.map(_payload_ref())
        stream = stream.write(handler)
        self._streams[service_name] = stream
        return stream

    def start(self) -> None:
        for service_name, stream in self._streams.items():
            stream.subscribe(source_address=self.node_address)

    def subscribe(self, service_name: str) -> bool:
        if service_name not in self._streams:
            raise ValueError(f"Service topic not registered: {service_name}")
        self._streams[service_name].subscribe(source_address=self.node_address)
        return True


__all__ = [
    "PlatformClientActor",
]


_payload_actor_ref = None


def _payload_ref():
    global _payload_actor_ref
    if _payload_actor_ref is None:
        class PayloadExtractor:
            def extract(self, ev):
                return ev.payload

        _payload_actor_ref = create_actor(PayloadExtractor)
    return _payload_actor_ref.extract
