from __future__ import annotations
from typing import TYPE_CHECKING, Generic, TypeVar

from flownet.api.datastream import (
    DataStream,
    SinkStream,
    _ensure_flow_target,
    _normalize_shuffle_key,
    _plan_meta,
    _system_target,
)
from flownet.compiler.transformation import Transformation
from flownet.core.exceptions import FlowDefinitionError
from flownet.compiler.flowtask import FlowTask



if TYPE_CHECKING:
    from flownet.api.actor_method_ref import ActorMethodRef
    from flownet.api.scheduler import SchedulePlan
    from flownet.core.stateless import StatelessOpRef

# TODO:注意，我们要让ConnectedStreams也能set_sink。等于说流返回一个流相加结果
class ConnectedStreams(DataStream):
    def __init__(self, flowtask: FlowTask, transformations: list[Transformation]):
        self.flowtask = flowtask
        self.transformations = transformations
        
        assert len(transformations) >= 2, "ConnectedStreams requires at least 2 transformations"
        for trans in transformations:
            assert trans.flowtask == flowtask, "All transformations must belong to the same flowtask"
    
    def connect(self, other: DataStream | ConnectedStreams) -> ConnectedStreams:
        if isinstance(other, DataStream):
            # ConnectedStreams + DataStream -> ConnectedStreams
            new_transformations = self.transformations + [other.transformation]
            return ConnectedStreams(self.flowtask, new_transformations)
        else:  # ConnectedStreams
            # ConnectedStreams + ConnectedStreams -> ConnectedStreams
            new_transformations = self.transformations + other.transformations
            return ConnectedStreams(self.flowtask, new_transformations)
        
    def apply(self, transformation: Transformation) -> DataStream:
        transformation.exception_handler_stack = list(self.flowtask._exception_handler_stack)
        for upstream_trans in self.transformations:
            transformation.add_upstream(upstream_trans)
        self.flowtask.pipeline.append(transformation)
        self.flowtask.transformations[transformation.trans_id] = transformation
        return DataStream(self.flowtask, transformation)

    def set_return(self):
        for transformation in self.transformations:
            self.flowtask.mark_return(transformation)

    # Backward-compatible alias (previously used as return marker).
    def set_sink(self):
        self.set_return()

    def write(self, sink: "ActorMethodRef | StatelessOpRef") -> SinkStream:
        target = sink
        if not callable(sink) and hasattr(sink, "write"):
            target = getattr(sink, "write")
        transformation = Transformation(
            self.flowtask,
            _ensure_flow_target(target, "write"),
            type="sink",
        )
        result = self.apply(transformation)
        self.flowtask.mark_sink(result.transformation)
        return SinkStream(self.flowtask, result.transformation)

    def publish(self, sink: "ActorMethodRef | StatelessOpRef") -> SinkStream:
        return self.write(sink)

    def join(
        self,
        *,
        key: int | str | None = None,
        key_index: int | None = None,
        key_field: str | None = None,
        num_partitions: int = 1,
        plan: "SchedulePlan | None" = None,
        addresses: list[str] | None = None,
        stage_id: str | None = None,
    ) -> DataStream:
        if len(self.transformations) != 2:
            raise FlowDefinitionError("join() requires exactly two streams.")
        if plan is not None and stage_id is not None and plan.stage_id != stage_id:
            raise FlowDefinitionError("join stage_id must match plan.stage_id when plan is provided.")
        if plan is None and not addresses:
            raise FlowDefinitionError("join() requires addresses when plan is not provided.")

        key_index, key_field = _normalize_shuffle_key(key, key_index, key_field)
        if stage_id is None:
            stage_id = plan.stage_id if plan is not None else None

        shuffle_meta: dict[str, object] = {
            "num_partitions": num_partitions,
            "key_index": key_index,
        }
        if key_field is not None:
            shuffle_meta["key_field"] = key_field
        if addresses:
            shuffle_meta["addresses"] = list(addresses)
        if stage_id is not None:
            shuffle_meta["stage_id"] = stage_id

        join_meta = {
            "stage_id": stage_id,
            "num_partitions": num_partitions,
        }

        base_config: dict[str, object] = {
            "shuffle": shuffle_meta,
            "join": join_meta,
        }
        if plan is not None:
            base_config["sys"] = _plan_meta(plan)

        left_config = dict(base_config)
        left_config["side"] = "left"
        right_config = dict(base_config)
        right_config["side"] = "right"

        left_router = Transformation(
            self.flowtask,
            _system_target(),
            type="sys/join",
            operator_config=left_config,
        )
        right_router = Transformation(
            self.flowtask,
            _system_target(),
            type="sys/join",
            operator_config=right_config,
        )
        join_reducer = Transformation(
            self.flowtask,
            _system_target(),
            type="sys/join-reducer",
        )

        left_router.add_upstream(self.transformations[0])
        right_router.add_upstream(self.transformations[1])
        join_reducer.add_upstream(left_router)
        join_reducer.add_upstream(right_router)

        for trans in (left_router, right_router, join_reducer):
            trans.exception_handler_stack = list(self.flowtask._exception_handler_stack)
            self.flowtask.pipeline.append(trans)
            self.flowtask.transformations[trans.trans_id] = trans

        return DataStream(self.flowtask, join_reducer)
