"""Pipeline — composable workflow orchestrator."""
from __future__ import annotations

import time
from datetime import UTC, datetime
from typing import Any, AsyncIterator, Callable
from uuid import uuid4

from logos_adk.runnable import Runnable
from logos_adk.types import Response, StreamChunk, StreamDone, TextChunk, Usage
from logos_adk.workflow.context import PipelineContext, PipelineResult, PipelineState, StepResult
from logos_adk.workflow.step import Step
from logos_adk.workflow.steps import (
    AgentStep,
    ApprovalStep,
    ConditionalStep,
    FunctionStep,
    LoopStep,
    ParallelStep,
)


class _StepConfig:
    """Internal config for a step in the pipeline."""

    def __init__(
        self,
        step: Step,
        retry: Any = None,
        timeout: float | None = None,
        guardrails: list | None = None,
    ) -> None:
        self.step = step
        self.retry = retry
        self.timeout = timeout
        self.guardrails = guardrails or []


class Pipeline(Step, Runnable):
    """Composable workflow pipeline."""

    def __init__(
        self,
        name: str,
        *,
        retry: Any = None,
        timeout: float | None = None,
        state_store: Any = None,
        learning_engine: Any = None,
    ) -> None:
        super().__init__(name)
        self._default_retry = retry
        self._default_timeout = timeout
        self._state_store = state_store
        self._learning_engine = learning_engine
        self._steps_config: list[_StepConfig] = []
        self._step_counter = 0

    def _auto_name(self, prefix: str = "step") -> str:
        name = f"{prefix}_{self._step_counter}"
        self._step_counter += 1
        return name

    # ── Builder API ──────────────────────────────────────────────

    def step(self, target: Any, *, name: str | None = None, retry: Any = None,
             timeout: float | None = None, guardrails: list | None = None) -> Pipeline:
        if isinstance(target, Step):
            s = target
            if name and s.name != name:
                s._name = name
        elif isinstance(target, Runnable):
            s = AgentStep(name or self._auto_name("agent"), target)
        else:
            raise TypeError(f"step() target must be Step or Runnable, got {type(target)}")
        self._steps_config.append(_StepConfig(s, retry=retry, timeout=timeout, guardrails=guardrails))
        return self

    def function(self, fn: Callable, *, name: str | None = None, retry: Any = None,
                 timeout: float | None = None) -> Pipeline:
        step_name = name or fn.__name__
        s = FunctionStep(step_name, fn)
        self._steps_config.append(_StepConfig(s, retry=retry, timeout=timeout))
        return self

    def graph(self, graph: Any, *, name: str | None = None, retry: Any = None,
              timeout: float | None = None, guardrails: list | None = None) -> Pipeline:
        if not isinstance(graph, Step):
            raise TypeError(f"graph() target must be a Step-compatible graph, got {type(graph)}")
        if name and graph.name != name:
            graph._name = name
        self._steps_config.append(_StepConfig(graph, retry=retry, timeout=timeout, guardrails=guardrails))
        return self

    def parallel(self, *targets: Any, name: str | None = None) -> Pipeline:
        sub_steps: list[Step] = []
        for t in targets:
            if isinstance(t, Step):
                sub_steps.append(t)
            elif callable(t):
                sub_steps.append(FunctionStep(t.__name__, t))
            elif isinstance(t, Runnable):
                sub_steps.append(AgentStep(self._auto_name("agent"), t))
            else:
                raise TypeError(f"parallel() target type not supported: {type(t)}")
        s = ParallelStep(name or self._auto_name("parallel"), sub_steps)
        self._steps_config.append(_StepConfig(s))
        return self

    def conditional(self, condition: Callable, *, then: Any, else_: Any = None,
                    name: str | None = None) -> Pipeline:
        then_step = self._wrap_target(then, "then")
        else_step = self._wrap_target(else_, "else") if else_ is not None else None
        s = ConditionalStep(name or self._auto_name("cond"), condition, then_step, else_step)
        self._steps_config.append(_StepConfig(s))
        return self

    def loop(self, target: Any, *, until: Callable, max_iterations: int = 10,
             name: str | None = None) -> Pipeline:
        inner = self._wrap_target(target, "loop_body")
        s = LoopStep(name or self._auto_name("loop"), inner, until, max_iterations)
        self._steps_config.append(_StepConfig(s))
        return self

    def approval(self, *, name: str | None = None, message: str | None = None) -> Pipeline:
        s = ApprovalStep(name or self._auto_name("approval"), message)
        self._steps_config.append(_StepConfig(s))
        return self

    def tool_workflow_learning(self, engine: Any) -> Pipeline:
        """Attach tool/workflow learning used during execution."""
        self._learning_engine = engine
        return self

    def _wrap_target(self, target: Any, prefix: str) -> Step:
        if isinstance(target, Step):
            return target
        if isinstance(target, Runnable):
            return AgentStep(self._auto_name(prefix), target)
        if callable(target):
            return FunctionStep(getattr(target, "__name__", self._auto_name(prefix)), target)
        raise TypeError(f"Cannot wrap {type(target)} as Step")

    # ── Step interface (for nesting) ─────────────────────────────

    async def execute(self, ctx: PipelineContext) -> StepResult:
        """Execute pipeline as a step (for nesting)."""
        result = await self._execute_pipeline(ctx)
        last_output = ""
        if result.outputs:
            last_key = list(result.outputs.keys())[-1]
            last_output = result.outputs[last_key].output
        return StepResult(
            output=last_output,
            status=result.status,
            duration_ms=result.duration_ms,
        )

    # ── Runnable interface ───────────────────────────────────────

    async def run(self, prompt: str, *, session_id: str | None = None, **kwargs: Any) -> Response:
        ctx = PipelineContext(
            pipeline_id=uuid4().hex,
            metadata={
                "trace_id": kwargs.get("_trace_id", uuid4().hex),
                "workspace_id": kwargs.get("_workspace_id"),
                "task_type": kwargs.get("_task_type"),
                "user_segment": kwargs.get("_user_segment"),
                "session_id": session_id,
            },
        )
        ctx["_last"] = prompt
        result = await self._execute_pipeline(ctx)

        last_output = prompt
        if result.outputs:
            last_key = list(result.outputs.keys())[-1]
            out = result.outputs[last_key].output
            last_output = str(out) if out is not None else ""

        return Response(
            content=last_output,
            tool_calls=[],
            usage=Usage(input_tokens=0, output_tokens=0),
            raw={"pipeline_result": _serialize_result(result)},
        )

    async def stream(self, prompt: str, **kwargs: Any) -> AsyncIterator[StreamChunk]:
        ctx = PipelineContext(
            pipeline_id=uuid4().hex,
            metadata={
                "trace_id": kwargs.get("_trace_id", uuid4().hex),
                "workspace_id": kwargs.get("_workspace_id"),
                "task_type": kwargs.get("_task_type"),
                "user_segment": kwargs.get("_user_segment"),
            },
        )
        ctx["_last"] = prompt

        start = time.monotonic()
        outputs: dict[str, StepResult] = {}
        status = "completed"

        for cfg in self._steps_config:
            if isinstance(cfg.step, LoopStep) and self._learning_engine is not None:
                ctx.metadata[f"adk_loop_max_iterations:{cfg.step.name}"] = (
                    self._learning_engine.recommend_loop_limit(
                        self.name,
                        cfg.step.name,
                        base_limit=cfg.step._max_iterations,
                        workspace_id=ctx.metadata.get("workspace_id"),
                    )
                )
            result = await self._execute_step(cfg, ctx)
            outputs[cfg.step.name] = result
            self._record_learning_signals(cfg.step, result, ctx)

            if result.status == "paused":
                status = "paused"
                break
            if result.status == "failed":
                status = "failed"
                break

            ctx[cfg.step.name] = result.output
            ctx["_last"] = result.output
            if result.output is not None:
                yield TextChunk(text=str(result.output))

        duration = (time.monotonic() - start) * 1000
        pipeline_result = PipelineResult(
            status=status, outputs=outputs, context=ctx, duration_ms=duration
        )
        last_output = ""
        if outputs:
            last_key = list(outputs.keys())[-1]
            out = outputs[last_key].output
            last_output = str(out) if out is not None else ""

        yield StreamDone(
            response=Response(
                content=last_output,
                tool_calls=[],
                usage=Usage(input_tokens=0, output_tokens=0),
                raw={"pipeline_result": _serialize_result(pipeline_result)},
            )
        )

    # ── Resume ────────────────────────────────────────────────────

    async def resume(self, pipeline_id: str, *, approval: bool = True) -> PipelineResult:
        """Resume a paused pipeline."""
        from logos_adk.errors import ConfigError

        if self._state_store is None:
            raise ConfigError("Cannot resume: pipeline has no state_store configured")

        agent_state = self._state_store.load(self.name, pipeline_id)
        if agent_state is None:
            raise ConfigError(f"No saved state for pipeline_id={pipeline_id}")

        ps = PipelineState.from_agent_state(agent_state)
        paused_cfg = self._steps_config[ps.current_step_index]

        if isinstance(paused_cfg.step, ApprovalStep) and self._learning_engine is not None:
            self._learning_engine.record_approval_outcome(
                pipeline_name=self.name,
                step_name=paused_cfg.step.name,
                approved=approval,
                corrected=not approval,
                workspace_id=ps.metadata_data.get("workspace_id"),
                trace_id=ps.metadata_data.get("trace_id"),
            )

        ctx = PipelineContext(
            pipeline_id=pipeline_id,
            data=ps.context_data,
            metadata={
                **ps.metadata_data,
                "_blackboard_versions": dict(ps.blackboard_versions_data),
            },
            blackboard=ps.blackboard_data,
        )
        outputs: dict[str, StepResult] = {}

        # Restore completed step results
        for step_name, sr_dict in ps.step_results.items():
            outputs[step_name] = StepResult(**sr_dict)

        if not approval:
            ps.status = "failed"
            ps.updated_at = datetime.now(UTC)
            self._state_store.save(self.name, ps.to_agent_state())
            return PipelineResult(status="failed", outputs=outputs, context=ctx, duration_ms=0)

        # Continue from the paused step so approval and future resumable steps
        # can finalize their own output into the restored context.
        start = time.monotonic()
        status = "completed"
        resume_index = ps.current_step_index
        ctx.metadata["_resume_node"] = paused_cfg.step.name

        for i in range(resume_index, len(self._steps_config)):
            cfg = self._steps_config[i]
            if isinstance(cfg.step, LoopStep) and self._learning_engine is not None:
                ctx.metadata[f"adk_loop_max_iterations:{cfg.step.name}"] = (
                    self._learning_engine.recommend_loop_limit(
                        self.name,
                        cfg.step.name,
                        base_limit=cfg.step._max_iterations,
                        workspace_id=ctx.metadata.get("workspace_id"),
                    )
                )
            result = await self._execute_step(cfg, ctx)
            outputs[cfg.step.name] = result
            self._record_learning_signals(cfg.step, result, ctx)

            if result.status == "paused":
                status = "paused"
                self._save_state(ctx, i, outputs, status)
                break
            if result.status == "failed":
                status = "failed"
                break
            ctx[cfg.step.name] = result.output
            ctx["_last"] = result.output

        duration = (time.monotonic() - start) * 1000
        ctx.metadata.pop("_resume_node", None)
        if self._state_store is not None and self._steps_config:
            self._save_state(
                ctx,
                min(max(resume_index, 0), len(self._steps_config) - 1),
                outputs,
                status,
            )
        return PipelineResult(status=status, outputs=outputs, context=ctx, duration_ms=duration)

    # ── Internal execution ───────────────────────────────────────

    def _save_state(self, ctx: PipelineContext, step_index: int,
                    outputs: dict[str, StepResult], status: str) -> None:
        """Save pipeline state to state_store for resume/checkpoint."""
        ps = PipelineState(
            pipeline_id=ctx.pipeline_id,
            pipeline_name=self.name,
            current_step_index=step_index,
            status=status,
            context_data=dict(ctx.data),
            metadata_data=dict(ctx.metadata),
            blackboard_data=dict(ctx.blackboard),
            blackboard_versions_data=ctx.blackboard.versions_snapshot(),
            step_results={
                name: {"output": sr.output, "status": sr.status,
                       "duration_ms": sr.duration_ms, "metadata": sr.metadata}
                for name, sr in outputs.items()
            },
        )
        self._state_store.save(self.name, ps.to_agent_state())

    async def _execute_pipeline(self, ctx: PipelineContext) -> PipelineResult:
        start = time.monotonic()
        outputs: dict[str, StepResult] = {}
        status = "completed"

        for i, cfg in enumerate(self._steps_config):
            if isinstance(cfg.step, LoopStep) and self._learning_engine is not None:
                ctx.metadata[f"adk_loop_max_iterations:{cfg.step.name}"] = (
                    self._learning_engine.recommend_loop_limit(
                        self.name,
                        cfg.step.name,
                        base_limit=cfg.step._max_iterations,
                        workspace_id=ctx.metadata.get("workspace_id"),
                    )
                )
            result = await self._execute_step(cfg, ctx)
            outputs[cfg.step.name] = result
            self._record_learning_signals(cfg.step, result, ctx)

            if result.status == "paused":
                status = "paused"
                if self._state_store is not None:
                    self._save_state(ctx, i, outputs, status)
                break
            if result.status == "failed":
                status = "failed"
                break

            ctx[cfg.step.name] = result.output
            ctx["_last"] = result.output

            # Checkpoint after each step if state_store configured
            if self._state_store is not None:
                self._save_state(ctx, i, outputs, "running")

        duration = (time.monotonic() - start) * 1000
        return PipelineResult(status=status, outputs=outputs, context=ctx, duration_ms=duration)

    def _record_learning_signals(self, step: Step, result: StepResult, ctx: PipelineContext) -> None:
        if self._learning_engine is None:
            return
        workspace_id = ctx.metadata.get("workspace_id")
        trace_id = ctx.metadata.get("trace_id")
        if isinstance(step, ConditionalStep) and "branch" in result.metadata:
            self._learning_engine.record_workflow_branch(
                pipeline_name=self.name,
                step_name=step.name,
                branch_name=str(result.metadata.get("branch")),
                success=result.status == "completed",
                duration_ms=result.duration_ms,
                workspace_id=workspace_id,
                trace_id=trace_id,
            )
        if isinstance(step, LoopStep) and "iterations" in result.metadata:
            self._learning_engine.record_loop_execution(
                pipeline_name=self.name,
                step_name=step.name,
                iterations=int(result.metadata.get("iterations", 0) or 0),
                max_iterations=int(result.metadata.get("max_iterations", step._max_iterations) or 1),
                termination_reason=str(result.metadata.get("termination_reason", "unknown")),
                success=result.status == "completed",
                workspace_id=workspace_id,
                trace_id=trace_id,
            )
        if isinstance(step, ApprovalStep) and result.status == "paused":
            self._learning_engine.record_approval_request(
                pipeline_name=self.name,
                step_name=step.name,
                workspace_id=workspace_id,
                trace_id=trace_id,
            )

    async def _execute_step(self, cfg: _StepConfig, ctx: PipelineContext) -> StepResult:
        import asyncio as _asyncio

        retry = cfg.retry or self._default_retry
        timeout = cfg.timeout or self._default_timeout

        async def _run() -> StepResult:
            if timeout is not None:
                return await _asyncio.wait_for(cfg.step.execute(ctx), timeout=timeout)
            return await cfg.step.execute(ctx)

        try:
            if retry is not None:
                result = await retry.call(_run)
            else:
                result = await _run()
        except Exception as exc:
            result = StepResult(
                output=None,
                status="failed",
                duration_ms=0,
                metadata={"error": f"{type(exc).__name__}: {exc}"},
            )

        # Apply guardrails
        if result.status == "completed" and cfg.guardrails:
            for guard in cfg.guardrails:
                check = await guard.check(str(result.output), role="assistant")
                if not check.passed and check.action == "block":
                    result = StepResult(
                        output=None,
                        status="failed",
                        duration_ms=result.duration_ms,
                        metadata={"error": f"Guardrail blocked: {check.message}"},
                    )
                    break

        return result


def _serialize_result(result: PipelineResult) -> dict:
    """Serialize PipelineResult for Response.raw."""
    payload = {
        "pipeline_id": result.context.pipeline_id,
        "status": result.status,
        "outputs": {
            name: {
                "output": r.output if isinstance(r.output, (str, int, float, bool, type(None), dict, list)) else str(r.output),
                "status": r.status,
                "duration_ms": r.duration_ms,
                "metadata": r.metadata,
            }
            for name, r in result.outputs.items()
        },
        "duration_ms": result.duration_ms,
    }
    graph_results = result.context.metadata.get("adk_graph_results")
    if isinstance(graph_results, dict) and graph_results:
        payload["graph_results"] = graph_results
    return payload
