"""Workflow container: steps, edges, validation, and ready-step selection."""

from __future__ import annotations

from collections import deque
from typing import Any

from pydantic import BaseModel, ConfigDict, Field

from agentbyte.workflow.core._structure_hash import build_workflow_structure_hash
from agentbyte.workflow.core.models import (
    Edge,
    EdgeCondition,
    JoinFailureMode,
    StepExecution,
    StepStatus,
    WorkflowCompletionMode,
    WorkflowStateConflictPolicy,
    WorkflowStateMode,
    WorkflowExecution,
    WorkflowValidationResult,
)
from agentbyte.workflow.steps.step import BaseStep


class WorkflowConfig(BaseModel):
    """Static workflow configuration metadata."""

    name: str
    description: str | None = None
    initial_state: dict[str, Any] = Field(default_factory=dict)
    metadata: dict[str, Any] = Field(default_factory=dict)
    completion_mode: WorkflowCompletionMode = WorkflowCompletionMode.ANY_END
    state_mode: WorkflowStateMode = WorkflowStateMode.IMMEDIATE
    state_conflict_policy: WorkflowStateConflictPolicy = WorkflowStateConflictPolicy.FAIL

    model_config = ConfigDict(extra="forbid")


class BaseWorkflow:
    """Base workflow that manages steps, edges, and validation."""

    def __init__(self, config: WorkflowConfig | None = None) -> None:
        self.config = config or WorkflowConfig(name="workflow")
        self.steps: dict[str, BaseStep] = {}
        self.edges: list[Edge] = []
        self.start_step_id: str | None = None
        self.end_step_ids: list[str] = []

    # Construction -----------------------------------------------------
    def add_step(self, step: BaseStep) -> "BaseWorkflow":
        self.steps[step.step_id] = step
        return self

    def add_edge(
        self,
        from_step: str | BaseStep,
        to_step: str | BaseStep,
        condition: EdgeCondition | dict[str, Any] | None = None,
    ) -> "BaseWorkflow":
        from_id = from_step.step_id if isinstance(from_step, BaseStep) else from_step
        to_id = to_step.step_id if isinstance(to_step, BaseStep) else to_step

        if isinstance(from_step, BaseStep):
            self.add_step(from_step)
        if isinstance(to_step, BaseStep):
            self.add_step(to_step)

        cond_obj = (
            condition
            if isinstance(condition, EdgeCondition)
            else EdgeCondition(**condition)  # type: ignore[arg-type]
            if condition
            else EdgeCondition()
        )
        self.edges.append(Edge(from_step=from_id, to_step=to_id, condition=cond_obj))
        return self

    def set_start_step(self, step: str | BaseStep) -> "BaseWorkflow":
        step_id = step.step_id if isinstance(step, BaseStep) else step
        if isinstance(step, BaseStep):
            self.add_step(step)
        self.start_step_id = step_id
        return self

    def add_end_step(self, step: str | BaseStep) -> "BaseWorkflow":
        step_id = step.step_id if isinstance(step, BaseStep) else step
        if isinstance(step, BaseStep):
            self.add_step(step)
        if step_id not in self.end_step_ids:
            self.end_step_ids.append(step_id)
        return self

    def configure_join(
        self,
        step: "str | BaseStep",
        *,
        failure_mode: JoinFailureMode = JoinFailureMode.FAIL_FAST,
    ) -> "BaseWorkflow":
        """Configure fan-in behaviour for a merge step.

        Call this after adding the step to customise how the runner handles
        upstream branch failures before data is passed into this join point.
        """
        step_id = step.step_id if isinstance(step, BaseStep) else step
        if step_id not in self.steps:
            msg = f"Step '{step_id}' is not part of this workflow"
            raise ValueError(msg)
        self.steps[step_id].metadata.join_failure_mode = failure_mode
        return self

    def is_step_failure_tolerated(self, step_id: str) -> bool:
        """Return True if failed upstream branches should be skipped for this step."""
        step = self.steps.get(step_id)
        if step is None:
            return False
        return step.metadata.join_failure_mode == JoinFailureMode.SKIP_FAILED

    def chain(self, *steps: BaseStep) -> "BaseWorkflow":
        if len(steps) < 2:
            msg = "chain requires at least two steps"
            raise ValueError(msg)
        for step in steps:
            self.add_step(step)
        for current, nxt in zip(steps, steps[1:]):
            self.add_edge(current, nxt)
        self.set_start_step(steps[0])
        self.add_end_step(steps[-1])
        return self

    def as_agent(
        self,
        *,
        name: str | None = None,
        description: str | None = None,
        memory_mode: str = "last_result",
        input_mapper: Any | None = None,
        output_mapper: Any | None = None,
        history_selector: Any | None = None,
        runner: Any | None = None,
    ) -> Any:
        """Wrap this workflow in a BaseAgent-compatible adapter."""
        from agentbyte.workflow.agent import WorkflowAgent

        return WorkflowAgent(
            self,
            name=name,
            description=description,
            memory_mode=memory_mode,
            input_mapper=input_mapper,
            output_mapper=output_mapper,
            history_selector=history_selector,
            runner=runner,
        )

    def get_step_dependencies(self, step_id: str) -> list[str]:
        """Get source step IDs that must feed into the given step."""
        return [edge.from_step for edge in self.edges if edge.to_step == step_id]

    # Validation -------------------------------------------------------
    def validate_workflow(self) -> WorkflowValidationResult:
        errors: list[str] = []
        warnings: list[str] = []

        if not self.steps:
            errors.append("Workflow has no steps")
            return WorkflowValidationResult(
                is_valid=False,
                errors=errors,
                warnings=warnings,
            )

        if not self.start_step_id or self.start_step_id not in self.steps:
            errors.append("Start step is not set or missing")

        for end_id in self.end_step_ids:
            if end_id not in self.steps:
                errors.append(f"End step '{end_id}' is not defined")

        for edge in self.edges:
            if edge.from_step not in self.steps:
                errors.append(f"Edge from missing step '{edge.from_step}'")
            if edge.to_step not in self.steps:
                errors.append(f"Edge to missing step '{edge.to_step}'")

        # Cycle detection via DFS
        has_cycles = self._detect_cycles()
        if has_cycles:
            errors.append("Workflow contains cycles")

        # Reachability via BFS from start
        unreachable = self._find_unreachable_steps()
        if unreachable:
            warnings.append(f"Unreachable steps: {', '.join(sorted(unreachable))}")

        # Conditional edge validation (basic contradiction detection)
        self._validate_conditional_edges(errors, warnings)

        # Graph semantics diagnostics (fan-in/fan-out/orphan terminal checks)
        self._validate_graph_semantics(warnings)

        # Type compatibility along edges
        # Fan-in join steps receive a synthetic envelope, so skip type check for their edges.
        fan_in_step_ids: set[str] = {
            step_id
            for step_id in self.steps
            if sum(
                1 for e in self.edges
                if e.to_step == step_id and e.condition.type == "always"
            ) >= 2
        }
        for edge in self.edges:
            if edge.to_step in fan_in_step_ids:
                continue  # envelope typed, not individual branch output
            if edge.from_step in self.steps and edge.to_step in self.steps:
                from_step = self.steps[edge.from_step]
                to_step = self.steps[edge.to_step]
                if not self._types_compatible(from_step, to_step):
                    errors.append(
                        f"Type mismatch from '{edge.from_step}' to '{edge.to_step}'"
                    )

        return WorkflowValidationResult(
            is_valid=not errors,
            errors=errors,
            warnings=warnings,
            has_cycles=has_cycles,
            unreachable_steps=unreachable,
        )

    def _detect_cycles(self) -> bool:
        visited: set[str] = set()
        stack: set[str] = set()

        def dfs(node: str) -> bool:
            visited.add(node)
            stack.add(node)
            for edge in self.edges:
                if edge.from_step != node:
                    continue
                nxt = edge.to_step
                if nxt not in visited:
                    if dfs(nxt):
                        return True
                elif nxt in stack:
                    return True
            stack.remove(node)
            return False

        for step_id in self.steps:
            if step_id in visited:
                continue
            if dfs(step_id):
                return True
        return False

    def _find_unreachable_steps(self) -> list[str]:
        if not self.start_step_id:
            return list(self.steps.keys())
        reachable: set[str] = set()
        queue: deque[str] = deque([self.start_step_id])
        while queue:
            node = queue.popleft()
            if node in reachable:
                continue
            reachable.add(node)
            for edge in self.edges:
                if edge.from_step == node:
                    queue.append(edge.to_step)
        return [sid for sid in self.steps if sid not in reachable]

    def _validate_conditional_edges(self, errors: list[str], warnings: list[str]) -> None:
        allowed_types = {"always", "output_based", "state_based"}
        allowed_operators = {"==", "!=", ">=", "<=", ">", "<", "in", "not_in"}

        for edge in self.edges:
            cond = edge.condition
            if cond.type not in allowed_types:
                errors.append(
                    f"Edge '{edge.from_step}' -> '{edge.to_step}' has unknown condition type '{cond.type}'"
                )
                continue

            if cond.type in {"output_based", "state_based"}:
                if not cond.field:
                    errors.append(
                        f"Edge '{edge.from_step}' -> '{edge.to_step}' with condition type '{cond.type}' must define field"
                    )
                if not cond.operator:
                    errors.append(
                        f"Edge '{edge.from_step}' -> '{edge.to_step}' with condition type '{cond.type}' must define operator"
                    )
                elif cond.operator not in allowed_operators:
                    errors.append(
                        f"Edge '{edge.from_step}' -> '{edge.to_step}' uses unsupported operator '{cond.operator}'"
                    )

        # Minimal contradiction detection: if multiple output_based edges on same field
        # with equality operators to different values, mark as potential contradiction.
        by_from: dict[str, list[Edge]] = {}
        for edge in self.edges:
            by_from.setdefault(edge.from_step, []).append(edge)

        for from_step, edges in by_from.items():
            seen_field_values: dict[str, set[Any]] = {}
            seen_branch_signatures: set[tuple[str, str, str, Any]] = set()
            equals_by_field: dict[str, set[Any]] = {}
            not_equals_by_field: dict[str, set[Any]] = {}
            for edge in edges:
                cond = edge.condition
                if cond.type == "output_based" and cond.field:
                    signature = (edge.to_step, cond.field, cond.operator or "", cond.value)
                    if signature in seen_branch_signatures:
                        warnings.append(
                            f"Duplicate conditional branch from '{from_step}' to '{edge.to_step}' on field '{cond.field}'"
                        )
                    else:
                        seen_branch_signatures.add(signature)

                if cond.type == "output_based" and cond.operator in {"==", "!="} and cond.field:
                    seen_field_values.setdefault(cond.field, set()).add(cond.value)
                    if cond.operator == "==":
                        equals_by_field.setdefault(cond.field, set()).add(cond.value)
                    elif cond.operator == "!=":
                        not_equals_by_field.setdefault(cond.field, set()).add(cond.value)

            for field, values in seen_field_values.items():
                # Suppress warning for exhaustive boolean splits (True/False) — they are
                # mutually exclusive by definition and represent a valid if/else pattern.
                is_exhaustive_bool = values == {True, False}
                if len(values) > 1 and not is_exhaustive_bool:
                    warnings.append(
                        f"Multiple conditional branches on field '{field}' with different equality values"
                    )

            for field, equal_values in equals_by_field.items():
                not_equal_values = not_equals_by_field.get(field, set())
                # Only warn when the same concrete value appears in both == and != sets,
                # meaning it could match two edges simultaneously.
                truly_overlapping = sorted(equal_values & not_equal_values)
                if truly_overlapping:
                    warnings.append(
                        (
                            f"Conditional branches from '{from_step}' on field '{field}' may overlap between "
                            "'==' and '!=' operators"
                        )
                    )

            for field, not_equal_values in not_equals_by_field.items():
                if len(not_equal_values) > 1:
                    warnings.append(
                        (
                            f"Multiple '!=' branches from '{from_step}' on field '{field}' may overlap and "
                            "activate multiple downstream paths"
                        )
                    )

    def _validate_graph_semantics(self, warnings: list[str]) -> None:
        incoming_by_step: dict[str, list[Edge]] = {step_id: [] for step_id in self.steps}
        outgoing_by_step: dict[str, list[Edge]] = {step_id: [] for step_id in self.steps}

        for edge in self.edges:
            if edge.to_step in incoming_by_step:
                incoming_by_step[edge.to_step].append(edge)
            if edge.from_step in outgoing_by_step:
                outgoing_by_step[edge.from_step].append(edge)

        for step_id in self.steps:
            incoming = incoming_by_step[step_id]
            outgoing = outgoing_by_step[step_id]

            if step_id != self.start_step_id and not incoming:
                warnings.append(
                    f"Step '{step_id}' has no incoming edges and will never run unless configured as start"
                )

            if step_id not in self.end_step_ids and not outgoing:
                warnings.append(
                    f"Step '{step_id}' has no outgoing edges but is not marked as an end step"
                )

            if len(incoming) > 1:
                always_edges = [edge for edge in incoming if edge.condition.type == "always"]
                conditional_edges = [edge for edge in incoming if edge.condition.type != "always"]
                if conditional_edges and not always_edges:
                    warnings.append(
                        (
                            f"Step '{step_id}' has only conditional fan-in edges; readiness uses OR semantics "
                            "across those branches"
                        )
                    )

            if len(outgoing) > 1:
                has_always = any(edge.condition.type == "always" for edge in outgoing)
                has_conditional = any(edge.condition.type != "always" for edge in outgoing)
                if has_always and has_conditional:
                    warnings.append(
                        (
                            f"Step '{step_id}' mixes always and conditional fan-out edges; this can create "
                            "ambiguous branch activation"
                        )
                    )

    def _types_compatible(self, from_step: BaseStep, to_step: BaseStep) -> bool:
        left_schema = from_step.output_type.model_json_schema()
        right_schema = to_step.input_type.model_json_schema()
        if left_schema == right_schema and from_step.output_type.__name__ == to_step.input_type.__name__:
            return True
        return from_step.output_type == to_step.input_type

    # Ready step selection ---------------------------------------------
    def get_ready_steps(self, execution: WorkflowExecution) -> list[BaseStep]:
        ready: list[BaseStep] = []
        _done = {StepStatus.COMPLETED, StepStatus.FAILED_TOLERATED}
        completed = {sid for sid, exec_obj in execution.step_executions.items() if exec_obj.status in _done}
        running = {
            sid
            for sid, exec_obj in execution.step_executions.items()
            if exec_obj.status in {StepStatus.RUNNING, StepStatus.SUSPENDED}
        }

        for step_id, step in self.steps.items():
            if step_id in completed or step_id in running:
                continue
            if not self.start_step_id:
                continue
            if step_id == self.start_step_id and step_id not in execution.step_executions:
                ready.append(step)
                continue

            incoming = [edge for edge in self.edges if edge.to_step == step_id]
            if not incoming:
                continue

            always_edges = [edge for edge in incoming if edge.condition.type == "always"]
            conditional_edges = [edge for edge in incoming if edge.condition.type != "always"]

            if always_edges:
                always_dependencies = {edge.from_step for edge in always_edges}
                if not always_dependencies.issubset(completed):
                    continue

            if not conditional_edges:
                if always_edges:
                    ready.append(step)
                continue

            # Conditional OR: any completed conditional dependency that passes gates readiness,
            # after all unconditional dependencies are satisfied.
            for edge in conditional_edges:
                if edge.from_step not in completed:
                    continue
                if self._evaluate_edge_condition(edge, execution):
                    ready.append(step)
                    break

        return ready

    def _evaluate_edge_condition(self, edge: Edge, execution: WorkflowExecution) -> bool:
        condition = edge.condition
        if condition.type == "always":
            return True
        if condition.type == "output_based":
            prev_exec: StepExecution | None = execution.step_executions.get(edge.from_step)
            if not prev_exec or not prev_exec.output_data:
                return False
            field_value = prev_exec.output_data.get(condition.field)
            return self._compare_values(field_value, condition.operator, condition.value)
        if condition.type == "state_based":
            field_value = execution.state.get(condition.field)
            return self._compare_values(field_value, condition.operator, condition.value)
        return False

    @staticmethod
    def _compare_values(left: Any, operator: str | None, right: Any) -> bool:
        try:
            if operator == "==":
                return left == right
            if operator == "!=":
                return left != right
            if operator == ">=":
                return left >= right
            if operator == "<=":
                return left <= right
            if operator == ">":
                return left > right
            if operator == "<":
                return left < right
            if operator == "in":
                return left in right
            if operator == "not_in":
                return left not in right
            return False
        except Exception:
            return False

    # Hashing ----------------------------------------------------------
    def compute_structure_hash(self) -> str:
        return build_workflow_structure_hash(
            steps=self.steps,
            edges=self.edges,
            start_step_id=self.start_step_id,
            end_step_ids=self.end_step_ids,
            completion_mode=self.config.completion_mode,
            state_mode=self.config.state_mode,
            state_conflict_policy=self.config.state_conflict_policy,
        )

    @property
    def id(self) -> str:
        return self.config.name


class Workflow(BaseWorkflow):
    """Concrete workflow container for runtime graph execution."""


__all__ = ["WorkflowConfig", "BaseWorkflow", "Workflow"]
