"""Concrete step implementations for workflow pipelines."""
from __future__ import annotations

import asyncio
import time
import json
from typing import Any, Awaitable, Callable

from logos_adk.checkpoint_json import JSONObject, ensure_json_object
from logos_adk.event_bus import EventBus
from logos_adk.runnable import Runnable
from logos_adk.workflow.context import BlackboardConflictError, PipelineContext, StepResult
from logos_adk.workflow.step import Step


class AgentStep(Step):
    """Step that runs an Agent or Runnable."""

    def __init__(self, name: str, agent: Runnable) -> None:
        super().__init__(name)
        self._agent = agent

    async def execute(self, ctx: PipelineContext) -> StepResult:
        last_val = ctx.get("_last", "")
        # Check if we are resuming THIS specific node
        is_this_resume = ctx.metadata.get("_resume_node") == self.name
        
        if isinstance(last_val, (dict, list)):
            prompt = _safe_serialize(last_val)
        else:
            prompt = str(last_val)
            
        start = time.monotonic()
        try:
            from logos_adk.context import RunContext, use_run_context
            
            run_ctx = RunContext(
                trace_id=ctx.metadata.get("trace_id"),
                session_id=ctx.metadata.get("session_id"),
                pipeline_ctx=ctx
            )

            with use_run_context(run_ctx):
                # If resuming, we only run if this is the target node OR if it's a normal run.
                # If we are resuming but this is NOT the target node, we might skip or run with empty.
                # To avoid InputValidationError on empty prompts for helpers:
                actual_prompt = prompt
                if ctx.metadata.get("_resume_node") and not is_this_resume:
                    # We are in a resume flow, but this node is just a bystander (e.g. in a ParallelStep)
                    # If it's already finished, we could skip it. For now, we just ensure it doesn't crash.
                    if not actual_prompt:
                        actual_prompt = "Continuing parallel execution..."
                elif not actual_prompt:
                    actual_prompt = "Continue."

                response = await self._agent.run(
                    actual_prompt,
                    session_id=ctx.metadata.get("session_id"),
                )
            duration = (time.monotonic() - start) * 1000
            
            # --- Internal Reflection Learning ---
            if response.content and "REJECTED:" in response.content:
                # If this result is a rejection from a critic, we should find the PREVIOUS agent 
                # and teach it. But a simpler way: if the CURRENT agent sees it's rejected, 
                # it should learn. 
                # Actually, in a loop A -> Critic -> A, the Critic produces REJECTED.
                # When A starts again, it should have learned.
                pass
            
            return StepResult(
                output=response.content,
                status=response.status,
                duration_ms=duration,
            )
        except Exception as exc:
            duration = (time.monotonic() - start) * 1000
            return StepResult(
                output=None,
                status="failed",
                duration_ms=duration,
                metadata={"error": f"{type(exc).__name__}: {exc}"},
            )


class FunctionStep(Step):
    """Step that runs an async function."""

    def __init__(self, name: str, fn: Callable[[PipelineContext], Awaitable[Any]]) -> None:
        super().__init__(name)
        self._fn = fn

    async def execute(self, ctx: PipelineContext) -> StepResult:
        start = time.monotonic()
        try:
            output = await self._fn(ctx)
            duration = (time.monotonic() - start) * 1000
            return StepResult(output=output, status="completed", duration_ms=duration)
        except Exception as exc:
            duration = (time.monotonic() - start) * 1000
            return StepResult(
                output=None,
                status="failed",
                duration_ms=duration,
                metadata={"error": f"{type(exc).__name__}: {exc}"},
            )


def _safe_serialize(obj: Any) -> str:
    """Safe JSON serialization that handles non-serializable objects."""
    def default(o):
        if hasattr(o, 'isoformat'):
            return o.isoformat()
        if hasattr(o, 'to_dict'):
            return o.to_dict()
        return str(o)
    return json.dumps(obj, ensure_ascii=False, indent=2, default=default)


class PublishEventStep(Step):
    """Step that publishes a JSON payload into the event bus."""

    def __init__(
        self,
        name: str,
        event_bus: EventBus,
        *,
        topic: str | Callable[[PipelineContext], str],
        payload: JSONObject | Callable[[PipelineContext], JSONObject] | None = None,
        correlation_id: str | Callable[[PipelineContext], str | None] | None = None,
        run_id: str | Callable[[PipelineContext], str | None] | None = None,
        publisher: str | Callable[[PipelineContext], str | None] | None = None,
        available_in_seconds: float = 0.0,
    ) -> None:
        super().__init__(name)
        self._event_bus = event_bus
        self._topic = topic
        self._payload = payload
        self._correlation_id = correlation_id
        self._run_id = run_id
        self._publisher = publisher
        self._available_in_seconds = available_in_seconds

    async def execute(self, ctx: PipelineContext) -> StepResult:
        start = time.monotonic()
        try:
            topic = _resolve_value(self._topic, ctx)
            if not isinstance(topic, str) or not topic.strip():
                raise ValueError("PublishEventStep topic must resolve to a non-empty string")
            payload_value = self._payload(ctx) if callable(self._payload) else self._payload
            if payload_value is None:
                last_value = ctx.get("_last")
                payload_value = last_value if isinstance(last_value, dict) else {"input": last_value}
            payload = ensure_json_object(payload_value, path=f"publish_event[{self.name}].payload")
            correlation_id = _resolve_optional_string(self._correlation_id, ctx, default=ctx.metadata.get("correlation_id"))
            run_id = _resolve_optional_string(self._run_id, ctx, default=ctx.pipeline_id)
            publisher = _resolve_optional_string(self._publisher, ctx, default=self.name)
            record = self._event_bus.publish(
                topic,
                payload,
                correlation_id=correlation_id,
                run_id=run_id,
                publisher=publisher,
                available_in_seconds=self._available_in_seconds,
            )
            duration = (time.monotonic() - start) * 1000
            return StepResult(
                output=dict(payload),
                status="completed",
                duration_ms=duration,
                metadata={
                    "event_id": record.id,
                    "topic": record.topic,
                    "correlation_id": record.correlation_id,
                    "run_id": record.run_id,
                },
            )
        except Exception as exc:
            duration = (time.monotonic() - start) * 1000
            return StepResult(
                output=None,
                status="failed",
                duration_ms=duration,
                metadata={"error": f"{type(exc).__name__}: {exc}"},
            )


class WaitForEventStep(Step):
    """Step that waits for an event matching the current run context."""

    def __init__(
        self,
        name: str,
        event_bus: EventBus,
        *,
        topic: str | Callable[[PipelineContext], str],
        correlation_id: str | Callable[[PipelineContext], str | None] | None = None,
        run_id: str | Callable[[PipelineContext], str | None] | None = None,
        lease_seconds: float = 60.0,
        message: str | None = None,
    ) -> None:
        super().__init__(name)
        self._event_bus = event_bus
        self._topic = topic
        self._correlation_id = correlation_id
        self._run_id = run_id
        self._lease_seconds = lease_seconds
        self._message = message

    async def execute(self, ctx: PipelineContext) -> StepResult:
        start = time.monotonic()
        try:
            topic = _resolve_value(self._topic, ctx)
            if not isinstance(topic, str) or not topic.strip():
                raise ValueError("WaitForEventStep topic must resolve to a non-empty string")
            correlation_id = _resolve_optional_string(self._correlation_id, ctx, default=ctx.metadata.get("correlation_id"))
            run_id = _resolve_optional_string(self._run_id, ctx, default=ctx.pipeline_id)
            consumer = f"{ctx.pipeline_id}:{self.name}"
            record = self._event_bus.claim(
                consumer,
                [topic],
                lease_seconds=self._lease_seconds,
                correlation_id=correlation_id,
                run_id=run_id,
            )
            duration = (time.monotonic() - start) * 1000
            if record is None:
                return StepResult(
                    output=None,
                    status="paused",
                    duration_ms=duration,
                    metadata={
                        "message": self._message or f"Waiting for event on topic '{topic}'",
                        "topic": topic,
                        "correlation_id": correlation_id,
                        "run_id": run_id,
                    },
                )
            self._event_bus.ack(record.id, consumer)
            return StepResult(
                output=dict(record.payload),
                status="completed",
                duration_ms=duration,
                metadata={
                    "event_id": record.id,
                    "topic": record.topic,
                    "correlation_id": record.correlation_id,
                    "run_id": record.run_id,
                },
            )
        except Exception as exc:
            duration = (time.monotonic() - start) * 1000
            return StepResult(
                output=None,
                status="failed",
                duration_ms=duration,
                metadata={"error": f"{type(exc).__name__}: {exc}"},
            )


def _resolve_value(value: Any, ctx: PipelineContext) -> Any:
    return value(ctx) if callable(value) else value


def _resolve_optional_string(
    value: str | Callable[[PipelineContext], str | None] | None,
    ctx: PipelineContext,
    *,
    default: str | None = None,
) -> str | None:
    resolved = _resolve_value(value, ctx)
    if resolved is None:
        return default
    if not isinstance(resolved, str):
        raise ValueError(f"Expected string-compatible value, got {type(resolved).__name__}")
    return resolved

class ParallelStep(Step):
    """Step that runs multiple steps in parallel."""

    def __init__(self, name: str, steps: list[Step]) -> None:
        super().__init__(name)
        self._steps = steps

    async def execute(self, ctx: PipelineContext) -> StepResult:
        if not self._steps:
            return StepResult(output={}, status="completed", duration_ms=0)

        start = time.monotonic()
        
        # Helper to dispatch events for sub-steps if a callback exists
        callback = ctx.metadata.get("_graph_event_callback")
        parent_versions = ctx.blackboard.versions_snapshot()
        
        async def execute_and_report(step):
            child_ctx = PipelineContext(
                pipeline_id=ctx.pipeline_id,
                data=dict(ctx.data),
                metadata=dict(ctx.metadata),
                blackboard=ctx.blackboard.clone(),
            )
            if callback:
                from datetime import datetime, UTC
                # Simple event reporting for parallel sub-steps
                try:
                    await callback({
                        "event": "node_started",
                        "event_type": "node",
                        "node_name": step.name,
                        "timestamp": datetime.now(UTC).isoformat()
                    })
                except Exception:
                    pass
                
            res = await step.execute(child_ctx)
            
            if callback:
                try:
                    await callback({
                        "event": f"node_{res.status}",
                        "event_type": "node",
                        "node_name": step.name,
                        "status": res.status,
                        "output": str(res.output)[:1000],
                        "timestamp": datetime.now(UTC).isoformat()
                    })
                except Exception:
                    pass
            return res, child_ctx

        results = await asyncio.gather(
            *(execute_and_report(step) for step in self._steps),
            return_exceptions=True,
        )
        duration = (time.monotonic() - start) * 1000

        output: dict[str, Any] = {}
        status = "completed"
        metadata: dict[str, Any] = {"errors": {}}
        paused_steps = []
        failure_count = 0

        for step, result in zip(self._steps, results):
            if isinstance(result, Exception):
                output[step.name] = None
                metadata["errors"][step.name] = f"{type(result).__name__}: {result}"
                failure_count += 1
            else:
                step_result, child_ctx = result
                if step_result.status == "failed":
                    output[step.name] = None
                    metadata["errors"][step.name] = step_result.metadata.get(
                        "error",
                        "Unknown error",
                    )
                    failure_count += 1
                else:
                    output[step.name] = step_result.output
                    if step_result.status == "paused":
                        status = "paused"
                        # Collect all paused metadata into a list to avoid collisions
                        paused_steps.append({
                            "node_name": step.name,
                            "metadata": step_result.metadata
                        })
                    if child_ctx.blackboard:
                        metadata.setdefault("step_blackboards", {})[step.name] = dict(child_ctx.blackboard)
                    try:
                        for key, value in child_ctx.blackboard.dirty_writes().items():
                            expected_version = parent_versions.get(key, 0)
                            ctx.blackboard.compare_and_set(
                                key,
                                value,
                                expected_version=expected_version,
                            )
                        for key in child_ctx.blackboard.dirty_deletes():
                            expected_version = parent_versions.get(key, 0)
                            if key in ctx.blackboard:
                                ctx.blackboard.delete_if_version(
                                    key,
                                    expected_version=expected_version,
                                )
                    except BlackboardConflictError as exc:
                        return StepResult(
                            output=output,
                            status="failed",
                            duration_ms=duration,
                            metadata={
                                "error": f"{type(exc).__name__}: {exc}",
                                "conflict_key": key,
                                "step_name": step.name,
                            },
                        )
        
        if paused_steps:
            metadata["paused_steps"] = paused_steps
            # For backward compatibility with UI, also keep the last one in root
            metadata.update(paused_steps[-1]["metadata"])
        if not metadata["errors"]:
            metadata.pop("errors")
        elif failure_count == len(self._steps) and status != "paused":
            status = "failed"

        return StepResult(output=output, status=status, duration_ms=duration, metadata=metadata)


class ConditionalStep(Step):
    """Step that branches based on a condition."""

    def __init__(
        self,
        name: str,
        condition: Callable[[PipelineContext], bool],
        then: Step,
        else_: Step | None = None,
    ) -> None:
        super().__init__(name)
        self._condition = condition
        self._then = then
        self._else = else_

    async def execute(self, ctx: PipelineContext) -> StepResult:
        start = time.monotonic()
        if self._condition(ctx):
            branch_name = self._then.name or "then"
            result = await self._then.execute(ctx)
        elif self._else is not None:
            branch_name = self._else.name or "else"
            result = await self._else.execute(ctx)
        else:
            duration = (time.monotonic() - start) * 1000
            return StepResult(
                output=None,
                status="skipped",
                duration_ms=duration,
                metadata={"branch": "skipped"},
            )
        metadata = dict(result.metadata)
        metadata["branch"] = branch_name
        return StepResult(
            output=result.output,
            status=result.status,
            duration_ms=result.duration_ms,
            metadata=metadata,
        )


class LoopStep(Step):
    """Step that repeats until a condition is met or max iterations."""

    def __init__(
        self,
        name: str,
        step: Step,
        until: Callable[[PipelineContext], bool],
        max_iterations: int = 10,
    ) -> None:
        super().__init__(name)
        self._step = step
        self._until = until
        self._max_iterations = max_iterations

    async def execute(self, ctx: PipelineContext) -> StepResult:
        start = time.monotonic()
        last_output: Any = None
        iterations = 0
        termination_reason = "max_iterations"
        loop_limit = int(ctx.metadata.get(f"adk_loop_max_iterations:{self.name}", self._max_iterations))

        for i in range(loop_limit):
            iterations = i + 1
            result = await self._step.execute(ctx)
            if result.status == "failed":
                metadata = dict(result.metadata)
                metadata.update(
                    {
                        "iterations": iterations,
                        "max_iterations": loop_limit,
                        "termination_reason": "failed",
                    }
                )
                return StepResult(
                    output=result.output,
                    status=result.status,
                    duration_ms=result.duration_ms,
                    metadata=metadata,
                )
            if result.status == "paused":
                metadata = dict(result.metadata)
                metadata.update(
                    {
                        "iterations": iterations,
                        "max_iterations": loop_limit,
                        "termination_reason": "paused",
                    }
                )
                return StepResult(
                    output=result.output,
                    status="paused",
                    duration_ms=result.duration_ms,
                    metadata=metadata,
                )
            last_output = result.output
            ctx["_last"] = last_output
            if self._until(ctx):
                termination_reason = "condition_met"
                break

        duration = (time.monotonic() - start) * 1000
        return StepResult(
            output=last_output,
            status="completed",
            duration_ms=duration,
            metadata={
                "iterations": iterations,
                "max_iterations": loop_limit,
                "termination_reason": termination_reason,
            },
        )


class ApprovalStep(Step):
    """Step that pauses the pipeline for human approval or choice."""

    def __init__(self, name: str, message: str | None = None, options: list[str] | None = None) -> None:
        super().__init__(name)
        self._message = message
        self._options = options

    async def execute(self, ctx: PipelineContext) -> StepResult:
        if ctx.metadata.get("_resume_node") == self.name:
            # If we are resuming THIS node, it means the user has made a choice
            # We return the choice as output and complete the step
            return StepResult(
                output=ctx.get("_last"),
                status="completed",
                duration_ms=0
            )

        return StepResult(
            output=None,
            status="paused",
            duration_ms=0,
            metadata={
                "message": self._message or "Awaiting action",
                "options": self._options
            },
        )
