
from __future__ import annotations

from typing import Any
import queue
import weakref
import uuid
import threading
import warnings

from flownet.api.actor_method_ref import ActorMethodRef
from flownet.compiler.flow_context import in_flow_build
from flownet.compiler.flowtask import FlowTask
from flownet.core.request import FlowAckRoute
from flownet.core.types import PayloadData
from flownet.runtime.flows.request_registry import (
    register_run_request,
    ack_run_request,
)


class BoundFlow:
    def __init__(self, flow_def: "FlowDef", instance: Any):
        self._flow_def = flow_def
        self._instance = instance

    def compile(self, *args, **kwargs) -> FlowTask:
        return self._flow_def._get_task(self._instance)

    def to_ref(self) -> ActorMethodRef:
        return self._flow_def.to_ref(self._instance)

    def put(self, *args, **kwargs):
        self._flow_def._fire_and_forget(self._instance, *args, **kwargs)

    def submit(self, *args, **kwargs):
        return self._flow_def._submit(self._instance, *args, **kwargs)

    def open(self):
        return self._flow_def._open(self._instance)

    def call(self, *args, **kwargs):
        return self._flow_def._call(self._instance, *args, **kwargs)

    def __call__(self, *args, **kwargs):
        return self.call(*args, **kwargs)

    def put_request(
        self,
        request_id,
        *args,
        ack_route=None,
        _next_meta: dict[str, Any] | None = None,
        **kwargs,
    ):
        return self._flow_def._dispatch(
            self._instance,
            *args,
            request_id=request_id,
            ack_route=ack_route,
            _next_meta=_next_meta,
            **kwargs,
        )


class FlowDef:
    # Development notes:
    # - FlowDef currently bundles compile/caching + request tracking + result sink.
    # - Potential next step: extract request tracking + result sink into a separate
    #   lightweight component (e.g., RequestTracker/ResultSink) and let FlowDef
    #   hold a reference instead of owning the state directly.
    # - Keep call/submit/put semantics unchanged; refactor should be internal-only.
    # ARCH_NOTE:
    # - For current M0 behavior, FlowDef is the owner of compile-time topology
    #   materialization and in-process task cache.
    # - "submit" does not publish a global flow catalog entry; it creates a run
    #   handle in runtime FlowRunService using the compiled task.
    def __init__(
        self,
        flowfunction,
        static=None,
    ):
        self.flowfunction = flowfunction
        self._method_name = flowfunction.__name__
        self._default_task = None
        self._object_tasks = weakref.WeakKeyDictionary()
        self._strong_object_tasks: dict[int, tuple[Any, FlowTask]] = {}
        self._default_output = None
        self._object_outputs = weakref.WeakKeyDictionary()
        self._strong_object_outputs: dict[int, tuple[Any, Any]] = {}

    _TO_REF_DEPRECATION_MESSAGE = (
        "FlowDef.to_ref()/BoundFlow.to_ref() is deprecated because it implicitly "
        "registers an object actor with hidden lifecycle side effects. "
        "Prefer flow_program_ref(flow_uri, version) for composition, or explicit "
        "actor registration via create_actor(...) + ActorMethodRef."
    )

    def __get__(self, instance, owner): 
        """
        当作为类成员访问时，会触发此方法。
        将绑定的实例记录下来，转为“bound method”。
        因为它本质上是class直辖的对象，而不是class object的成员对象
        """
        if in_flow_build():
            return self._to_ref_internal(instance)
        if instance is None:
            return self
        return BoundFlow(self, instance)
        
    def __call__(self, *args, **kwargs):
        return self.call(*args, **kwargs)

    def compile(self, *args, **kwargs) -> FlowTask:
        return self._get_task(None)

    def put(self, *args, **kwargs):
        self._fire_and_forget(None, *args, **kwargs)

    def submit(self, *args, **kwargs):
        return self._submit(None, *args, **kwargs)

    def open(self):
        return self._open(None)

    def call(self, *args, **kwargs):
        return self._call(None, *args, **kwargs)
    
    def to_ref(self, instance=None) -> ActorMethodRef:
        warnings.warn(
            self._TO_REF_DEPRECATION_MESSAGE,
            DeprecationWarning,
            stacklevel=2,
        )
        return self._to_ref_internal(instance)

    def _to_ref_internal(self, instance=None) -> ActorMethodRef:
        actor_object = instance if instance is not None else self
        method_name = self._method_name if instance is not None else "compile"
        from flownet.api.runtime import make_object_actor
        actor_ref = make_object_actor(actor_object)
        return ActorMethodRef(
            address=actor_ref.address,
            actor_id=actor_ref.actor_id,
            method=method_name,
        )

    def _compile_task(self, instance) -> FlowTask:
        from flownet.compiler.v1_compile_facade import compile_v0_flow_task

        return compile_v0_flow_task(
            self.flowfunction,
            instance=instance,
        )

    def _set_strong_task(self, instance, task: FlowTask) -> FlowTask:
        key = id(instance)
        self._strong_object_tasks[key] = (instance, task)
        return task

    def _get_strong_task(self, instance) -> FlowTask:
        key = id(instance)
        cached = self._strong_object_tasks.get(key)
        if cached is not None:
            cached_instance, cached_task = cached
            if cached_instance is instance:
                return cached_task
        task = self._compile_task(instance)
        self._strong_object_tasks[key] = (instance, task)
        return task

    def _get_task(self, instance) -> FlowTask:
        if instance is None:
            if self._default_task is None:
                self._default_task = self._compile_task(None)
            return self._default_task

        try:
            cached = self._object_tasks.get(instance)
            if cached is not None:
                return cached
        except TypeError:
            return self._get_strong_task(instance)

        task = self._compile_task(instance)
        try:
            self._object_tasks[instance] = task
        except TypeError:
            return self._set_strong_task(instance, task)
        return task

    def _make_output_channel(self, instance):
        from flownet.runtime.flows.output_channel import RunOutputChannel

        target = self._to_ref_internal(instance)
        return RunOutputChannel(target, runtime_managed=True)

    def _set_strong_output(self, instance, output):
        key = id(instance)
        self._strong_object_outputs[key] = (instance, output)
        return output

    def _get_strong_output(self, instance):
        key = id(instance)
        cached = self._strong_object_outputs.get(key)
        if cached is not None:
            cached_instance, cached_output = cached
            if cached_instance is instance:
                return cached_output
        output = self._make_output_channel(instance)
        self._strong_object_outputs[key] = (instance, output)
        return output

    def _get_output_channel(self, instance):
        if instance is None:
            if self._default_output is None:
                self._default_output = self._make_output_channel(None)
            return self._default_output

        try:
            cached = self._object_outputs.get(instance)
            if cached is not None:
                return cached
        except TypeError:
            return self._get_strong_output(instance)

        output = self._make_output_channel(instance)
        try:
            self._object_outputs[instance] = output
        except TypeError:
            return self._set_strong_output(instance, output)
        return output

    def _make_tracked_request_id(self, tracker) -> str:
        from flownet.api.runtime import make_object_actor
        actor_ref = make_object_actor(tracker)
        return f"{actor_ref.actor_id}:{uuid.uuid4().hex}"

    def _default_ack_route(self) -> FlowAckRoute:
        from flownet.api.runtime import make_object_actor

        actor_ref = make_object_actor(self)
        return FlowAckRoute(
            coordinator_target=actor_ref.actor_id,
            coordinator_address=str(actor_ref.address or "").strip() or None,
        )

    def _dispatch(
        self,
        instance,
        *args,
        request_id: str,
        ack_route=None,
        _next_meta: dict[str, Any] | None = None,
        **kwargs,
    ):
        task = self._get_task(instance)
        payload = PayloadData(
            args=list(args),
            kwargs=kwargs,
        )
        from flownet.runtime.kernel import get_runtime_kernel

        hub = get_runtime_kernel()
        request = task.make_request(
            payload,
            request_id=request_id,
            origin_address=hub.local_address(),
            ack_route=ack_route,
        )
        if isinstance(_next_meta, dict) and _next_meta:
            top = request.flow_stack.top()
            if top is not None:
                top.meta = {"next_meta": dict(_next_meta)}
        hub.require_flows().submit(request)
        return request_id

    def _fire_and_forget(self, instance, *args, **kwargs):
        run = self._create_flow_run(
            instance,
            egress_kind=None,
            with_output_channel=False,
        )

        new_request_id = getattr(run, "new_request_id", None)
        if not callable(new_request_id):
            raise RuntimeError(
                "FlowDef.put() requires runtime run support. "
                "Run must implement new_request_id(tracker)."
            )

        dispatch = getattr(run, "dispatch", None)
        if not callable(dispatch):
            raise RuntimeError(
                "FlowDef.put() requires runtime run support. "
                "Run must implement dispatch(*args, request_id=..., **kwargs)."
            )

        request_id = new_request_id(tracker=self)
        dispatch(*args, request_id=request_id, **kwargs)
        close_ingress = getattr(run, "close_ingress", None)
        if callable(close_ingress):
            close_ingress()
        return request_id

    def _make_augmented_task_id(self, egress_kind: str | None) -> str:
        output_kind = egress_kind or "default"
        return f"{self._method_name}.wrapper.{output_kind}"

    def _create_flow_run(
        self,
        instance,
        *,
        egress_kind: str | None,
        with_output_channel: bool,
    ):
        # ARCH_NOTE:
        # - The compiled FlowTask is attached via run_config for this submission.
        # - Runtime side indexes the result by run_id (FlowRunRef), so this path
        #   is run-oriented even though the topology comes from cached FlowDef.
        from flownet.api import runtime as runtime_api

        ack_route_holder = {"value": None}
        run_config = {
            "instance": instance,
            "flow_task": self._get_task(instance),
            "flow_id": self._method_name,
            "native_flow_id": self._method_name,
            "augmented_task_id": self._make_augmented_task_id(egress_kind)
            if with_output_channel
            else None,
            "dispatch": lambda *args, request_id, _next_meta=None, **kwargs: self._dispatch(
                instance,
                *args,
                request_id=request_id,
                ack_route=ack_route_holder["value"],
                _next_meta=_next_meta,
                **kwargs,
            ),
            "tracked_request_id_factory": self._make_tracked_request_id,
            # FlowDef.submit() performs explicit request registration before
            # dispatch and therefore must not auto-increment on dispatch().
            "dispatch_tracks_requests": False,
        }
        if not with_output_channel:
            run_config["ack_route"] = self._default_ack_route()
        if with_output_channel:
            run_config["output_factory"] = lambda: self._get_output_channel(instance)

        egress = {"kind": egress_kind} if egress_kind is not None else None

        run = runtime_api.submit_flow(
            self,
            ingress={"kind": "manual_push"},
            egress=egress,
            run_config=run_config,
        )
        output_fn = getattr(run, "output", None)
        if callable(output_fn):
            ack_route_holder["value"] = output_fn("ack_route")
        return run

    def _open(self, instance):
        run = self._create_flow_run(
            instance,
            egress_kind="buffered",
            with_output_channel=True,
        )
        return run.require_output("output")

    def _submit(self, instance, *args, **kwargs):
        run = self._create_flow_run(
            instance,
            egress_kind=None,
            with_output_channel=False,
        )
        request_id = run.new_request_id(tracker=self)
        event = register_run_request(run, request_id, initial_delta=1)

        run.dispatch(*args, request_id=request_id, **kwargs)
        close_ingress = getattr(run, "close_ingress", None)
        if callable(close_ingress):
            close_ingress()

        def _status() -> str:
            request_done = getattr(run, "request_done", None)
            if callable(request_done):
                try:
                    if request_done(request_id):
                        return "completed"
                except Exception:
                    pass
            run_status = getattr(run, "status", None)
            if callable(run_status):
                try:
                    return str(run_status())
                except Exception:
                    pass
            return "completed" if event.is_set() else "running"

        def _cancel() -> bool:
            cancel_run = getattr(run, "cancel", None)
            if not callable(cancel_run):
                return False
            return bool(cancel_run())

        return RequestHandle(
            request_id,
            event,
            status_fn=_status,
            cancel_fn=_cancel,
        )

    def _call(self, instance, *args, **kwargs):
        from flownet.runtime.flows.output_channel import RunOutputDone
        from flownet.core.exceptions import FlowDefinitionError
        from flownet.runtime.flows.output_channel import RunOutputError
        from flownet.runtime.flows.output_channel import RunOutputGap
        from flownet.core.exceptions import FlowException, ExceptionEvent, ExceptionContext

        task = self._get_task(instance)
        if task.no_return:
            raise FlowDefinitionError("Flow has no return stream; call() is not supported.")

        run = self._create_flow_run(
            instance,
            egress_kind="return_once",
            with_output_channel=True,
        )
        request_id = run.write(*args, **kwargs)
        close_ingress = getattr(run, "close_ingress", None)
        if callable(close_ingress):
            close_ingress()
        results = []
        done_seen = False
        supports_late_drain = callable(getattr(run, "output", None))
        while True:
            try:
                item = run.read(
                    timeout=1.0 if done_seen and supports_late_drain else 5.0,
                    request_id=request_id,
                )
            except queue.Empty:
                if done_seen and supports_late_drain:
                    break
                raise
            if isinstance(item, RunOutputDone):
                if not supports_late_drain:
                    break
                done_seen = True
                continue
            if isinstance(item, RunOutputError):
                if item.error_type == "FlowDefinitionError":
                    raise FlowDefinitionError(item.message)
                if isinstance(item.event, ExceptionEvent):
                    raise FlowException(item.event)
                raise FlowException(
                    ExceptionEvent(
                        error_type=item.error_type,
                        message=item.message,
                        traceback=item.traceback_text,
                        context=ExceptionContext(
                            request_id=request_id,
                            phase="flow_call",
                        ),
                        error=None,
                    )
                )
            if isinstance(item, RunOutputGap):
                continue
            results.append(item)
        return results

    def event_chain_done(self, request_id: str, delta: int):
        ack_run_request(request_id, delta)


class RequestHandle:
    def __init__(
        self,
        request_id: str,
        event: threading.Event,
        status_fn=None,
        cancel_fn=None,
    ):
        self.request_id = request_id
        self._event = event
        self._status_fn = status_fn
        self._cancel_fn = cancel_fn

    def id(self) -> str:
        return self.request_id

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

    def status(self) -> str:
        if callable(self._status_fn):
            return str(self._status_fn())
        return "completed" if self._event.is_set() else "running"

    def cancel(self) -> bool:
        if not callable(self._cancel_fn):
            return False
        return bool(self._cancel_fn())

    def wait(self, timeout: float | None = None) -> bool:
        return self._event.wait(timeout=timeout)

    def done(self) -> bool:
        return self._event.is_set()
