"""Deterministic workflow composition for cross-system Autopilot plans."""

from __future__ import annotations

from dataclasses import dataclass
from typing import TYPE_CHECKING, Iterable

from .contracts import PlanStep

if TYPE_CHECKING:
    from .tools import ToolRegistry


class WorkflowCompositionError(ValueError):
    """Raised when a workflow cannot be composed deterministically."""


@dataclass(frozen=True)
class WorkflowComposition:
    ordered_steps: list[PlanStep]
    execution_order: list[str]
    rollback_chain: list[str]
    dependency_map: dict[str, list[str]]


_INFERRED_TOOL_DEPENDENCIES: dict[str, set[str]] = {
    "slack.notify": {
        "sequence.pause",
        "sfdc.followup.enforce",
        "sfdc.record.reassign_owner",
        "sfdc.record.update",
        "sfdc.record.enrich_missing_fields",
        "tasks.create",
    },
    "sfdc.record.update": {
        "sfdc.record.enrich_missing_fields",
    },
}


def compose_steps(
    steps: list[PlanStep],
    *,
    registry: ToolRegistry | None = None,
    infer_dependencies: bool = True,
) -> WorkflowComposition:
    if not steps:
        return WorkflowComposition(
            ordered_steps=[],
            execution_order=[],
            rollback_chain=[],
            dependency_map={},
        )

    step_by_id: dict[str, PlanStep] = {}
    for step in steps:
        step_id = str(step.step_id or "").strip()
        if not step_id:
            raise WorkflowCompositionError("step_id is required for all workflow steps")
        if step_id in step_by_id:
            raise WorkflowCompositionError(f"duplicate step_id in workflow: {step_id}")
        step_by_id[step_id] = step

    dependency_sets: dict[str, set[str]] = {
        step.step_id: {str(dep).strip() for dep in step.depends_on if str(dep).strip()}
        for step in steps
    }
    if infer_dependencies:
        _apply_inferred_dependencies(steps, dependency_sets)

    for step_id, deps in dependency_sets.items():
        if step_id in deps:
            raise WorkflowCompositionError(f"step {step_id} cannot depend on itself")
        missing = sorted(dep for dep in deps if dep not in step_by_id)
        if missing:
            raise WorkflowCompositionError(
                f"step {step_id} depends on unknown step_id values: {', '.join(missing)}"
            )

    order_ids = _topological_sort(steps, dependency_sets)
    dependent_count: dict[str, int] = {step_id: 0 for step_id in dependency_sets}
    for deps in dependency_sets.values():
        for dep_id in deps:
            dependent_count[dep_id] = dependent_count.get(dep_id, 0) + 1

    ordered_steps: list[PlanStep] = []
    for step_id in order_ids:
        original = step_by_id[step_id]
        depends_on = sorted(dependency_sets.get(step_id, set()))
        rollback_on_failure = bool(original.rollback_on_failure)
        if (
            not rollback_on_failure
            and dependent_count.get(step_id, 0) > 0
            and _tool_supports_rollback(original.tool, registry)
        ):
            rollback_on_failure = True
        ordered_steps.append(
            original.model_copy(
                update={
                    "depends_on": depends_on,
                    "rollback_on_failure": rollback_on_failure,
                }
            )
        )

    rollback_chain = [
        step.step_id
        for step in reversed(ordered_steps)
        if step.rollback_on_failure
    ]
    dependency_map = {step_id: sorted(values) for step_id, values in dependency_sets.items()}
    return WorkflowComposition(
        ordered_steps=ordered_steps,
        execution_order=order_ids,
        rollback_chain=rollback_chain,
        dependency_map=dependency_map,
    )


def expand_selected_step_ids(steps: list[PlanStep], selected_step_ids: Iterable[str]) -> set[str]:
    step_ids = {step.step_id for step in steps}
    dependency_map = {
        step.step_id: {str(dep).strip() for dep in step.depends_on if str(dep).strip()}
        for step in steps
    }
    selected = {str(step_id).strip() for step_id in selected_step_ids if str(step_id).strip()}
    missing = sorted(step_id for step_id in selected if step_id not in step_ids)
    if missing:
        raise WorkflowCompositionError(f"unknown step_id values: {', '.join(missing)}")

    expanded = set(selected)
    queue = list(selected)
    while queue:
        current = queue.pop()
        deps = dependency_map.get(current, set())
        missing_deps = sorted(dep for dep in deps if dep not in step_ids)
        if missing_deps:
            raise WorkflowCompositionError(
                f"step {current} depends on unknown step_id values: {', '.join(missing_deps)}"
            )
        for dep_id in deps:
            if dep_id in expanded:
                continue
            expanded.add(dep_id)
            queue.append(dep_id)
    return expanded


def _apply_inferred_dependencies(
    steps: list[PlanStep],
    dependency_sets: dict[str, set[str]],
) -> None:
    for index, step in enumerate(steps):
        upstream_tools = _INFERRED_TOOL_DEPENDENCIES.get(step.tool, set())
        if not upstream_tools:
            continue
        for candidate in steps[:index]:
            if candidate.tool not in upstream_tools:
                continue
            if not _shares_target(step, candidate):
                continue
            dependency_sets.setdefault(step.step_id, set()).add(candidate.step_id)


def _topological_sort(steps: list[PlanStep], dependency_sets: dict[str, set[str]]) -> list[str]:
    order: list[str] = []
    unresolved = {step.step_id: set(dependency_sets.get(step.step_id, set())) for step in steps}
    stable_order = [step.step_id for step in steps]

    while unresolved:
        ready_step_id = next(
            (
                step_id
                for step_id in stable_order
                if step_id in unresolved and not unresolved[step_id]
            ),
            None,
        )
        if ready_step_id is None:
            cycle_nodes = ", ".join(sorted(unresolved.keys()))
            raise WorkflowCompositionError(f"dependency cycle detected: {cycle_nodes}")
        order.append(ready_step_id)
        unresolved.pop(ready_step_id, None)
        for deps in unresolved.values():
            deps.discard(ready_step_id)
    return order


def _shares_target(step: PlanStep, candidate: PlanStep) -> bool:
    left = _step_targets(step)
    right = _step_targets(candidate)
    if not left or not right:
        return False
    return bool(left.intersection(right))


def _step_targets(step: PlanStep) -> set[str]:
    args = step.args if isinstance(step.args, dict) else {}
    targets: list[str] = []

    for key in ("record_ids", "targets", "ids"):
        value = args.get(key)
        if isinstance(value, list):
            targets.extend(str(item).strip() for item in value if str(item).strip())

    for key in ("record_id", "event_id", "entity_id"):
        value = args.get(key)
        if isinstance(value, list):
            targets.extend(str(item).strip() for item in value if str(item).strip())
        elif value not in (None, ""):
            targets.append(str(value).strip())

    rows = args.get("records")
    if isinstance(rows, list):
        for row in rows:
            if not isinstance(row, dict):
                continue
            row_id = row.get("record_id")
            if row_id not in (None, ""):
                targets.append(str(row_id).strip())

    return {value.lower() for value in targets if value}


def _tool_supports_rollback(tool_name: str, registry: ToolRegistry | None) -> bool:
    if registry is None:
        return tool_name in {
            "sequence.pause",
            "sfdc.record.update",
            "sfdc.record.reassign_owner",
            "sfdc.record.enrich_missing_fields",
            "sfdc.followup.enforce",
            "tasks.create",
        }
    tool_def = registry.get_tool(tool_name)
    if tool_def is None:
        return False
    return bool(getattr(tool_def, "rollback_supported", False))


__all__ = [
    "WorkflowComposition",
    "WorkflowCompositionError",
    "compose_steps",
    "expand_selected_step_ids",
]
