from __future__ import annotations

from typing import TYPE_CHECKING, Any, Mapping

from flownet.compiler.flow_def import BoundFlow, FlowDef
from flownet.compiler.flowtask import FlowTask
from flownet.core.flow_program import (
    FlowProgramRef,
    ResolvedFlowProgram,
    normalize_flow_program_version,
    normalize_flow_uri,
)
from flownet.core.exceptions import FlowDefinitionError
from flownet.core.request import FlowRequest, TransformationFrame
from flownet.core.types import PayloadData
from flownet.engine.context import OperatorContext
from flownet.engine.utils.flow_stack import append_frame

from .base import OperatorAction, OperatorHandler


if TYPE_CHECKING:
    from flownet.compiler.transformation import Transformation


class PipeOperatorHandler(OperatorHandler):
    """Operator behavior for pipe(): accepts subgraphs and continues execution."""

    @staticmethod
    def _next_frame_meta(frame_meta: Any) -> dict[str, Any] | None:
        if not isinstance(frame_meta, dict):
            return None
        propagated: dict[str, Any] = {}
        next_meta = frame_meta.get("next_meta")
        if isinstance(next_meta, dict) and next_meta:
            propagated["next_meta"] = dict(next_meta)
        sys_meta = frame_meta.get("sys")
        if isinstance(sys_meta, dict) and sys_meta:
            propagated["sys"] = dict(sys_meta)
        return propagated or None

    @staticmethod
    def _normalize_lineage_entry(
        entry: Mapping[str, Any],
        *,
        fallback_ref: FlowProgramRef,
    ) -> dict[str, str]:
        flow_uri = entry.get("flow_uri")
        version = entry.get("version")
        try:
            normalized_uri = normalize_flow_uri(str(flow_uri))
        except Exception:
            normalized_uri = fallback_ref.flow_uri
        try:
            normalized_version = normalize_flow_program_version(str(version))
        except Exception:
            normalized_version = fallback_ref.version
        normalized: dict[str, str] = {
            "flow_uri": normalized_uri,
            "version": normalized_version,
        }
        program_id = entry.get("program_id")
        if isinstance(program_id, str) and program_id.strip():
            normalized["program_id"] = program_id
        program_hash = entry.get("program_hash")
        if isinstance(program_hash, str) and program_hash.strip():
            normalized["program_hash"] = program_hash
        return normalized

    @classmethod
    def _normalize_lineage_list(
        cls,
        value: Any,
        *,
        fallback_ref: FlowProgramRef,
    ) -> list[dict[str, str]]:
        if not isinstance(value, list):
            return []
        lineage: list[dict[str, str]] = []
        for item in value:
            if not isinstance(item, Mapping):
                continue
            try:
                lineage.append(cls._normalize_lineage_entry(item, fallback_ref=fallback_ref))
            except Exception:
                continue
        return lineage

    @classmethod
    def _build_flow_program_frame_meta(
        cls,
        frame_meta: Any,
        *,
        program_ref: FlowProgramRef,
        lineage_entry: Mapping[str, Any],
    ) -> dict[str, Any]:
        propagated = cls._next_frame_meta(frame_meta) or {}
        next_meta_obj = propagated.get("next_meta")
        next_meta = dict(next_meta_obj) if isinstance(next_meta_obj, dict) else {}

        lineage = cls._normalize_lineage_list(
            next_meta.get("flow_program_lineage"),
            fallback_ref=program_ref,
        )
        current_entry = cls._normalize_lineage_entry(lineage_entry, fallback_ref=program_ref)
        current_flow_id = f"{current_entry['flow_uri']}@{current_entry['version']}"
        current_path = [f"{item['flow_uri']}@{item['version']}" for item in lineage]
        if current_flow_id in set(current_path):
            path_text = " -> ".join(current_path + [current_flow_id])
            raise FlowDefinitionError(
                "pipe() FlowProgramRef cycle detected at runtime: "
                f"{path_text}"
            )

        lineage.append(current_entry)
        next_meta["flow_program_ref"] = dict(current_entry)
        next_meta["flow_program_lineage"] = lineage
        propagated["next_meta"] = next_meta
        return propagated

    async def handle_result(
        self,
        ctx: OperatorContext,
        request: FlowRequest,
        transformation: Transformation,
        payload: PayloadData,
        result: Any,
        *,
        frame_meta: dict[str, Any] | None = None,
    ) -> OperatorAction:
        propagated_meta = frame_meta
        if propagated_meta is None:
            propagated_meta = self._next_frame_meta(ctx.current_frame_meta)
        if isinstance(result, FlowTask):
            if result.no_return and transformation.downstreams:
                raise FlowDefinitionError(
                    "pipe() subflow has no return but downstream operators exist."
                )
            entry = result.resolve_entry_transformation()
            if entry is None:
                raise FlowDefinitionError("pipe() subflow has no entry transformation.")
            new_frame = TransformationFrame(
                entry,
                payload=payload,
                meta=propagated_meta,
            )
            request.flow_stack = append_frame(request.flow_stack, new_frame)
            return OperatorAction.CONTINUE

        if isinstance(result, (FlowDef, BoundFlow)):
            task = result.compile(*payload.args, **payload.kwargs)
            if task.no_return and transformation.downstreams:
                raise FlowDefinitionError(
                    "pipe() subflow has no return but downstream operators exist."
                )
            entry = task.resolve_entry_transformation()
            if entry is None:
                raise FlowDefinitionError("pipe() subflow has no entry transformation.")
            new_frame = TransformationFrame(
                entry,
                payload=payload,
                meta=propagated_meta,
            )
            request.flow_stack = append_frame(request.flow_stack, new_frame)
            return OperatorAction.CONTINUE

        raise TypeError(
            "pipe() expects a FlowDef/BoundFlow or FlowTask to build a subgraph."
        )

    async def handle_request(
        self,
        ctx: OperatorContext,
        request: FlowRequest,
        transformation_frame: TransformationFrame,
    ) -> OperatorAction:
        embedded_pipe_flow = transformation_frame.transformation.pipe_flow
        if embedded_pipe_flow is not None:
            return await self.handle_result(
                ctx,
                request,
                transformation_frame.transformation,
                transformation_frame.payload,
                embedded_pipe_flow,
                frame_meta=self._next_frame_meta(ctx.current_frame_meta),
            )

        target = transformation_frame.transformation.target
        if isinstance(target, FlowProgramRef):
            # Pre-check lineage cycle before any catalog/network resolution.
            self._build_flow_program_frame_meta(
                ctx.current_frame_meta,
                program_ref=target,
                lineage_entry={
                    "flow_uri": target.flow_uri,
                    "version": target.version,
                },
            )
            resolver = getattr(ctx, "resolve_flow_program", None)
            if resolver is None:
                raise FlowDefinitionError(
                    "pipe() FlowProgramRef resolution is unavailable in current runtime context."
                )
            try:
                resolved = await resolver(target)
            except Exception as exc:
                raise FlowDefinitionError(
                    f"pipe() failed to resolve FlowProgramRef {target.flow_id()}."
                ) from exc
            resolved_payload = resolved
            lineage_entry: dict[str, Any] = {
                "flow_uri": target.flow_uri,
                "version": target.version,
            }
            if isinstance(resolved, ResolvedFlowProgram):
                resolved_payload = resolved.payload
                lineage_entry = resolved.lineage_entry()
            propagated_meta = self._build_flow_program_frame_meta(
                ctx.current_frame_meta,
                program_ref=target,
                lineage_entry=lineage_entry,
            )
            return await self.handle_result(
                ctx,
                request,
                transformation_frame.transformation,
                transformation_frame.payload,
                resolved_payload,
                frame_meta=propagated_meta,
            )
        return await super().handle_request(ctx, request, transformation_frame)



class FilterOperatorHandler(OperatorHandler):
    """Operator behavior for filter(): drops payloads when predicate is falsy."""

    async def handle_result(
        self,
        ctx: OperatorContext,
        request: FlowRequest,
        transformation: Transformation,
        payload: PayloadData,
        result: Any,
    ) -> OperatorAction:
        if isinstance(result, (FlowTask, FlowDef, BoundFlow)):
            raise TypeError(
                "Subgraph results are only allowed for pipe(); use pipe() instead of filter."
            )
        if isinstance(result, PayloadData):
            await ctx.post_process_result(request, result, payload)
            return OperatorAction.RETURN
        if result:
            await ctx.post_process_result(request, payload, payload)
            return OperatorAction.RETURN
        await ctx.post_process_result(request, None, payload)
        return OperatorAction.RETURN
