# DOCS_REF:
# - docs/user/README.md
# - docs/user/04-services-sources-topics.md
# - docs/user/reference/api-surface.md
# DOCS_SYNC:
# - If TopicStream operators or subscribe semantics change, update the
#   referenced docs and examples/docs/{services_sources_topics_example.py,cookbook_topic_flow_topic_example.py}.

from __future__ import annotations

from flownet.api.datastream import DataStream, PartitionedStream
import flownet.api.node_control as node_control_api
from flownet.api.services import get_services
from flownet.compiler.flowtask import FlowTask
from flownet.core.exceptions import FlowDefinitionError


class TopicStream:
    _SERVICE_TOPIC_PREFIX = "service:"

    def __init__(self, topic_id: str):
        self.topic_id = topic_id
        self._flowtask = FlowTask()
        self._head = DataStream(self._flowtask, transformation=None, init=True)
        self._tail: DataStream = self._head
        self._subscribed = False
        self._subscription_id: str | None = None

    def map(self, fn_object) -> "TopicStream":
        self._tail = self._tail.map(fn_object)
        return self

    def pipe(self, fn_object) -> "TopicStream":
        self._tail = self._tail.pipe(fn_object)
        return self

    def filter(self, fn_object) -> "TopicStream":
        self._tail = self._tail.filter(fn_object)
        return self

    def flatmap(self, fn_object) -> "TopicStream":
        self._tail = self._tail.flatmap(fn_object)
        return self

    def tap(self, fn_object) -> "TopicStream":
        self._tail = self._tail.tap(fn_object)
        return self

    def write(self, sink) -> "TopicStream":
        self._tail = self._tail.write(sink)
        return self

    def publish(self, sink) -> "TopicStream":
        return self.write(sink)

    def shuffle(self, **kwargs) -> "TopicStream":
        self._tail = self._tail.shuffle(**kwargs)
        return self

    def key_by(self, **kwargs) -> "TopicStream":
        self._tail = self._tail.key_by(**kwargs)
        return self

    def repartition(self, num_partitions: int, **kwargs) -> "TopicStream":
        self._tail = self._tail.repartition(num_partitions=num_partitions, **kwargs)
        return self

    def reduce(self, **kwargs) -> "TopicStream":
        if not isinstance(self._tail, PartitionedStream):
            raise FlowDefinitionError("reduce() requires shuffle() or repartition() first.")
        self._tail = self._tail.reduce(**kwargs)
        return self

    def process_with_state(self, handler_ref, **kwargs) -> "TopicStream":
        if not isinstance(self._tail, PartitionedStream):
            raise FlowDefinitionError("process_with_state() requires shuffle(), key_by(), or repartition() first.")
        self._tail = self._tail.process_with_state(handler_ref, **kwargs)
        return self

    def set_return(self) -> "TopicStream":
        self._tail.set_return()
        return self

    # Backward-compatible alias (same as return marker).
    def set_sink(self) -> "TopicStream":
        return self.set_return()

    def _service_name(self) -> str | None:
        if self.topic_id.startswith(self._SERVICE_TOPIC_PREFIX):
            return self.topic_id[len(self._SERVICE_TOPIC_PREFIX):]
        return None

    def subscribe(self, source_address: str | None = None) -> "TopicStream":
        service_name = self._service_name()
        if source_address is None and service_name is None:
            self._subscribe_topic_registry()
            return self

        if not self._flowtask.pipeline:
            raise ValueError("TopicStream has no operators. Add map/pipe/write before subscribe().")

        hub = get_services()
        topics = hub.require_topics()
        execution_address = hub.local_address()

        if service_name is None:
            raise ValueError("source_address subscribe is only supported for service topics.")

        if source_address is None or source_address == execution_address:
            if topics.topic_registry is not None:
                self._subscribe_topic_registry()
                return self
            if not self._subscribed:
                self._subscription_id = topics.require_service_registry().subscribe_flow(
                    service_name,
                    self._flowtask,
                    execution_address,
                )
                self._subscribed = True
            return self

        self._subscription_id = node_control_api.subscribe_service_flow_on(
            source_address,
            service_name,
            self._flowtask,
            execution_address,
        )
        self._subscribed = True
        return self

    def submit(self, source_address: str | None = None) -> "TopicStream":
        return self.subscribe(source_address=source_address)

    def unsubscribe(self, source_address: str | None = None) -> bool:
        if not self._subscription_id:
            return False

        service_name = self._service_name()
        if source_address is None and service_name is None:
            return self._unsubscribe_topic_registry()

        hub = get_services()
        topics = hub.require_topics()
        execution_address = hub.local_address()

        if service_name is None:
            raise ValueError("source_address unsubscribe is only supported for service topics.")

        if source_address is None or source_address == execution_address:
            if topics.topic_registry is not None:
                return self._unsubscribe_topic_registry()
            return topics.require_service_registry().unsubscribe_flow(self._subscription_id)

        return node_control_api.unsubscribe_service_flow_on(
            source_address,
            self._subscription_id,
        )

    def _subscribe_topic_registry(self) -> None:
        if not self._flowtask.pipeline:
            raise ValueError("TopicStream has no operators. Add map/pipe/write before subscribe().")
        hub = get_services()
        registry = hub.require_topics().require_topic_registry()
        if not self._subscribed:
            self._subscription_id = registry.attach_task(
                self.topic_id,
                self._flowtask,
                hub.local_address(),
            )
            self._subscribed = True

    def _unsubscribe_topic_registry(self) -> bool:
        hub = get_services()
        registry = hub.require_topics().require_topic_registry()
        return registry.detach_task(self._subscription_id)


def topic_stream(topic_id: str) -> TopicStream:
    return TopicStream(topic_id)


def from_topic(topic_id: str) -> TopicStream:
    return topic_stream(topic_id)


def from_service(service_name: str) -> TopicStream:
    normalized_name = _normalize_service_name(service_name)
    return from_topic(f"{TopicStream._SERVICE_TOPIC_PREFIX}{normalized_name}")


def _normalize_service_name(service_name: str) -> str:
    if not isinstance(service_name, str):
        raise TypeError("service_name must be a string.")

    normalized = service_name.strip()
    if not normalized:
        raise ValueError("service_name must be a non-empty string.")
    if normalized.startswith(TopicStream._SERVICE_TOPIC_PREFIX):
        raise ValueError("from_service() expects service name without the 'service:' prefix.")
    return normalized


__all__ = [
    "TopicStream",
    "topic_stream",
    "from_topic",
    "from_service",
]
