from __future__ import annotations

from collections.abc import Iterator, Mapping, Sequence
from copy import deepcopy
from dataclasses import dataclass
from typing import Any

from wf_core.errors import WorkflowExecutionError
from wf_core.run_state import ExecutionFrame, LineageState, RunState, StateWrite
from wf_core.runtime.ops.state import (
    StatePatch,
    commit_state_patch,
    safe_set_nested_value,
)


@dataclass(slots=True)
class LineageStateView:
    """Committed state plus writes visible inside one child lineage.

    Concurrent foreach item writes live in ``RunState.lineages``; the barrier
    keeps only each item's lineage identity. Future native subgraphs and
    fork/gather should reuse this primitive instead of rebuilding lineage
    overlay logic.
    """

    base_state: Mapping[str, Any]
    writes: Sequence[StateWrite]

    def to_state_dict(self) -> dict[str, Any]:
        """Materialize the lineage-visible state as an isolated mutable dict."""
        # Correctness first: this full copy isolates sibling reads. If state grows
        # large, replace this with a lazy/copy-on-write overlay.
        state_view = deepcopy(dict(self.base_state))
        for write in self.writes:
            safe_set_nested_value(
                state_view,
                list(write.path.parts),
                write.visible_value,
            )
        return state_view


def lineage_writes_for_frame(
    run: RunState, frame: ExecutionFrame
) -> Sequence[StateWrite]:
    """Return ancestor and current-lineage writes visible to this frame.

    An empty child lineage still inherits writes buffered by its ancestors, as
    happens when an outer concurrent foreach writes before entering an inner
    foreach. Lineage existence and scope therefore control traversal; the
    current lineage having its own writes does not.
    """
    lineage = run.lineages.get(frame.lineage_id)
    if lineage is not None and lineage.scope_id == frame.scope_id:
        return tuple(
            lineage_state_writes(
                run, scope_id=frame.scope_id, lineage_id=frame.lineage_id
            )
        )
    return ()


def is_scope_root_lineage_frame(run: RunState, frame: ExecutionFrame) -> bool:
    """Return whether writes from this frame commit to its scope state root."""
    lineage = run.lineages.get(frame.lineage_id)
    return (
        lineage is not None
        and lineage.scope_id == frame.scope_id
        and lineage.parent_id is None
    )


def commit_patch_for_frame(
    run: RunState, frame: ExecutionFrame, patch: StatePatch
) -> dict[str, Any]:
    """Commit at a scope root or buffer writes in the frame lineage.

    Child workflow root frames own a committed child-state root just like the
    top-level root frame owns `RunState.state`. Descendant branch/item frames
    remain isolated until an explicit barrier or future gather commits them.
    """
    if is_scope_root_lineage_frame(run, frame):
        return commit_state_patch(scope_state_for_frame(run, frame), patch)
    append_lineage_writes(
        run,
        scope_id=frame.scope_id,
        lineage_id=frame.lineage_id,
        writes=patch.writes,
    )
    return {}


def commit_foreach_aware_patch(
    run: RunState, frame: ExecutionFrame, patch: StatePatch
) -> dict[str, Any]:
    """Commit one write patch with foreach-aware routing.

    Ordinary frames commit (or buffer) through their own lineage. The walk
    climbs through every serial item owner until it reaches either the
    workflow/subgraph scope root, where it commits, or a concurrent item
    boundary, where it buffers in that item lineage for the barrier to
    merge. The whole ancestry is validated first: the write lands only
    after the chain reaches an acyclic non-item ancestor, so a parent
    cycle fails closed even when it passes through a concurrent
    boundary. Malformed ownership, missing parents, parent cycles, and
    closed or superseded activations fail closed.
    """
    from wf_core.runtime.foreach_state import (
        item_frame_owner,
        require_foreach_activation,
    )

    current = frame
    seen: set[str] = set()
    buffer_in: ExecutionFrame | None = None
    while True:
        owner = item_frame_owner(current)
        if owner is None:
            break
        if current.id in seen:
            raise WorkflowExecutionError(
                f"cycle detected in foreach parent chain at frame {current.id!r}"
            )
        seen.add(current.id)
        parent_frame = run.frames.get(owner.parent_frame_id)
        if parent_frame is None:
            raise WorkflowExecutionError(
                "foreach item state references missing parent frame "
                f"{owner.parent_frame_id!r} for child frame {current.id!r}"
            )
        activation = require_foreach_activation(
            parent_frame, owner.foreach_node_id, owner.activation_id
        )
        if buffer_in is None and activation.barrier.mode == "concurrent":
            buffer_in = current
        current = parent_frame
    if buffer_in is not None:
        append_lineage_writes(
            run,
            scope_id=buffer_in.scope_id,
            lineage_id=buffer_in.lineage_id,
            writes=patch.writes,
        )
        return {}
    return commit_patch_for_frame(run, current, patch)


def scope_state_for_frame(run: RunState, frame: ExecutionFrame) -> dict[str, Any]:
    """Return the committed state root for the frame's runtime scope."""
    scope = run.scopes.get(frame.scope_id)
    if scope is None:
        raise ValueError(f"unknown scope {frame.scope_id!r}")
    return scope.committed_state


def scope_input_for_frame(run: RunState, frame: ExecutionFrame) -> dict[str, Any]:
    """Return the invocation input associated with the frame's workflow scope."""
    scope = run.scopes.get(frame.scope_id)
    if scope is None:
        raise ValueError(f"unknown scope {frame.scope_id!r}")
    return scope.workflow_input


def add_lineage(
    run: RunState,
    *,
    scope_id: str,
    lineage_id: str,
    parent_id: str | None,
) -> None:
    """Create one lineage record inside an existing runtime scope."""
    if scope_id not in run.scopes:
        raise ValueError(f"unknown scope {scope_id!r}")
    if lineage_id in run.lineages:
        raise ValueError(f"duplicate lineage {lineage_id!r}")
    if parent_id is not None and parent_id not in run.lineages:
        raise ValueError(f"unknown parent lineage {parent_id!r}")
    run.lineages[lineage_id] = LineageState(
        id=lineage_id,
        scope_id=scope_id,
        parent_id=parent_id,
    )


def append_lineage_writes(
    run: RunState,
    *,
    scope_id: str,
    lineage_id: str,
    writes: Sequence[StateWrite],
) -> None:
    """Append ordered writes to an existing lineage without committing state."""
    lineage = _lineage(run, scope_id=scope_id, lineage_id=lineage_id)
    lineage.writes.extend(writes)


def lineage_patch(
    run: RunState,
    *,
    scope_id: str,
    lineage_id: str,
) -> StatePatch:
    """Return a replayable patch for one lineage's pending writes.

    Barrier/gather code should consume this instead of reconstructing a patch
    from visible state. Incoming values are the replay source of truth.
    """
    lineage = _lineage(run, scope_id=scope_id, lineage_id=lineage_id)
    return StatePatch(writes=list(lineage.writes))


def lineage_state_view(
    run: RunState,
    *,
    scope_id: str,
    lineage_id: str,
) -> dict[str, Any]:
    """Materialize scope committed state plus ancestor/current lineage writes."""
    scope = run.scopes.get(scope_id)
    if scope is None:
        raise ValueError(f"unknown scope {scope_id!r}")
    writes: list[StateWrite] = []
    for lineage in _lineage_chain(run, scope_id=scope_id, lineage_id=lineage_id):
        writes.extend(lineage.writes)
    return LineageStateView(scope.committed_state, writes).to_state_dict()


def lineage_state_writes(
    run: RunState,
    *,
    scope_id: str,
    lineage_id: str,
) -> Iterator[StateWrite]:
    """Yield ancestor and current lineage writes in read-visibility order."""
    for lineage in _lineage_chain(run, scope_id=scope_id, lineage_id=lineage_id):
        yield from lineage.writes


def _lineage(run: RunState, *, scope_id: str, lineage_id: str) -> LineageState:
    lineage = run.lineages.get(lineage_id)
    if lineage is None:
        raise ValueError(f"unknown lineage {lineage_id!r}")
    if lineage.scope_id != scope_id:
        raise ValueError(
            f"lineage {lineage_id!r} belongs to scope {lineage.scope_id!r}, "
            f"not {scope_id!r}"
        )
    return lineage


def _lineage_chain(
    run: RunState, *, scope_id: str, lineage_id: str
) -> Iterator[LineageState]:
    lineage = _lineage(run, scope_id=scope_id, lineage_id=lineage_id)
    reverse_chain: list[LineageState] = []
    seen: set[str] = set()
    while True:
        if lineage.id in seen:
            raise WorkflowExecutionError(
                f"cycle detected in lineage chain at {lineage.id!r}"
            )
        seen.add(lineage.id)
        reverse_chain.append(lineage)
        if lineage.parent_id is None:
            break
        lineage = _lineage(run, scope_id=scope_id, lineage_id=lineage.parent_id)
    yield from reversed(reverse_chain)