"""Tests for the Spec 08.7 declarative workflow schema."""

from __future__ import annotations

import json
from typing import Any

import pytest
from pydantic import BaseModel
import yaml

from agentbyte.workflow import (
    EdgeConditionSchema,
    EdgeSchema,
    EchoStep,
    FunctionStep,
    StepKind,
    StepMetadata,
    StepSchema,
    Workflow,
    WorkflowBuilder,
    WorkflowLoader,
    WorkflowSchemaBuildError,
    WorkflowConfig,
    WorkflowSchema,
    WorkflowStepRegistry,
    workflow_to_schema,
)
from agentbyte.workflow.core.models import (
    JoinFailureMode,
    WorkflowCompletionMode,
    WorkflowStateConflictPolicy,
    WorkflowStateMode,
)
from agentbyte.workflow.steps.http import HttpStep


# ---------------------------------------------------------------------------
# Helpers
# ---------------------------------------------------------------------------

class Msg(BaseModel):
    text: str


class Out(BaseModel):
    result: str


def _identity(input_data: Msg, _ctx: Any) -> Out:
    return Out(result=input_data.text)


def _build_linear_function_workflow(name: str = "test_wf") -> Workflow:
    step_a = FunctionStep("step_a", StepMetadata(name="step_a"), Msg, Out, _identity)
    step_b = FunctionStep("step_b", StepMetadata(name="step_b"), Out, Out, lambda d, _: d)
    return Workflow(WorkflowConfig(name=name)).chain(step_a, step_b)


# ---------------------------------------------------------------------------
# EdgeConditionSchema
# ---------------------------------------------------------------------------

class TestEdgeConditionSchema:
    def test_defaults(self) -> None:
        schema = EdgeConditionSchema()
        assert schema.type == "always"
        assert schema.field is None
        assert schema.operator is None
        assert schema.value is None

    def test_round_trip_via_edge_condition(self) -> None:
        schema = EdgeConditionSchema(type="output_based", field="score", operator=">=", value=5)
        edge_cond = schema.to_edge_condition()
        restored = EdgeConditionSchema.from_edge_condition(edge_cond)
        assert restored == schema

    def test_serialises_to_json_safe_dict(self) -> None:
        schema = EdgeConditionSchema(type="state_based", field="flag", operator="==", value=True)
        d = schema.model_dump(mode="json")
        assert d["type"] == "state_based"
        assert d["value"] is True


# ---------------------------------------------------------------------------
# StepSchema
# ---------------------------------------------------------------------------

class TestStepSchema:
    def test_defaults(self) -> None:
        step = StepSchema(id="s1", name="my_step")
        assert step.kind == StepKind.UNKNOWN
        assert step.max_retries == 0
        assert step.tags == []
        assert step.config == {}

    def test_extra_fields_forbidden(self) -> None:
        with pytest.raises(Exception):
            StepSchema(id="s1", name="my_step", unknown_field="x")  # type: ignore[call-arg]


# ---------------------------------------------------------------------------
# WorkflowSchema
# ---------------------------------------------------------------------------

class TestWorkflowSchema:
    def _minimal_schema(self) -> WorkflowSchema:
        return WorkflowSchema(
            name="wf",
            start_step="a",
            end_steps=["b"],
            steps=[
                StepSchema(id="a", name="a"),
                StepSchema(id="b", name="b"),
            ],
            edges=[EdgeSchema(from_step="a", to_step="b")],
        )

    def test_defaults(self) -> None:
        schema = self._minimal_schema()
        assert schema.schema_version == "1"
        assert schema.state_mode == WorkflowStateMode.IMMEDIATE
        assert schema.state_conflict_policy == WorkflowStateConflictPolicy.FAIL
        assert schema.completion_mode == WorkflowCompletionMode.ANY_END
        assert schema.initial_state == {}
        assert schema.metadata == {}

    def test_to_json_round_trip(self) -> None:
        schema = self._minimal_schema()
        json_str = schema.to_json()
        restored = WorkflowSchema.from_json(json_str)
        assert restored == schema

    def test_to_dict_round_trip(self) -> None:
        schema = self._minimal_schema()
        d = schema.to_dict()
        restored = WorkflowSchema.from_dict(d)
        assert restored == schema

    def test_json_is_valid_json(self) -> None:
        schema = self._minimal_schema()
        parsed = json.loads(schema.to_json())
        assert parsed["name"] == "wf"
        assert parsed["schema_version"] == "1"

    def test_extra_fields_forbidden(self) -> None:
        with pytest.raises(Exception):
            WorkflowSchema(  # type: ignore[call-arg]
                name="wf",
                start_step="a",
                end_steps=["a"],
                steps=[],
                surprise="field",
            )


# ---------------------------------------------------------------------------
# workflow_to_schema (extractor)
# ---------------------------------------------------------------------------

class TestWorkflowToSchema:
    def test_linear_function_workflow(self) -> None:
        wf = _build_linear_function_workflow("pipe")
        schema = workflow_to_schema(wf)

        assert schema.name == "pipe"
        assert schema.start_step == "step_a"
        assert schema.end_steps == ["step_b"]
        assert len(schema.steps) == 2
        assert len(schema.edges) == 1

    def test_step_kinds_detected(self) -> None:
        wf = _build_linear_function_workflow()
        schema = workflow_to_schema(wf)
        for step in schema.steps:
            assert step.kind == StepKind.FUNCTION

    def test_function_ref_captured(self) -> None:
        wf = _build_linear_function_workflow()
        schema = workflow_to_schema(wf)
        step_a = next(s for s in schema.steps if s.id == "step_a")
        # function_ref is a dotted module.qualname string
        assert "function_ref" in step_a.config
        assert step_a.config["function_ref"] is not None
        assert "_identity" in step_a.config["function_ref"]

    def test_echo_step_config_captured(self) -> None:
        echo = EchoStep(
            "greet",
            StepMetadata(name="greet"),
            prefix="Hello, ",
            suffix="!",
            delay_seconds=0.5,
        )
        wf = Workflow(WorkflowConfig(name="echo_wf"))
        wf.add_step(echo)
        wf.set_start_step(echo)
        wf.add_end_step(echo)

        schema = workflow_to_schema(wf)
        step_schema = schema.steps[0]
        assert step_schema.kind == StepKind.ECHO
        assert step_schema.config["prefix"] == "Hello, "
        assert step_schema.config["suffix"] == "!"
        assert step_schema.config["delay_seconds"] == 0.5

    def test_http_step_kind_detected(self) -> None:
        http = HttpStep("fetch", StepMetadata(name="fetch"))
        wf = Workflow(WorkflowConfig(name="http_wf"))
        wf.add_step(http)
        wf.set_start_step(http)
        wf.add_end_step(http)

        schema = workflow_to_schema(wf)
        assert schema.steps[0].kind == StepKind.HTTP

    def test_edge_condition_serialised(self) -> None:
        step_a = FunctionStep("a", StepMetadata(name="a"), Msg, Out, _identity)
        step_b = FunctionStep("b", StepMetadata(name="b"), Msg, Out, _identity)
        wf = Workflow(WorkflowConfig(name="cond_wf"))
        wf.add_step(step_a)
        wf.add_step(step_b)
        wf.add_edge(
            "a", "b",
            condition={"type": "output_based", "field": "result", "operator": "==", "value": "ok"},
        )
        wf.set_start_step("a")
        wf.add_end_step("b")

        schema = workflow_to_schema(wf)
        edge = schema.edges[0]
        assert edge.condition.type == "output_based"
        assert edge.condition.field == "result"
        assert edge.condition.operator == "=="
        assert edge.condition.value == "ok"

    def test_workflow_config_fields_preserved(self) -> None:
        cfg = WorkflowConfig(
            name="cfg_wf",
            description="test workflow",
            state_mode=WorkflowStateMode.STAGED,
            state_conflict_policy=WorkflowStateConflictPolicy.LAST_WRITE_WINS,
            completion_mode=WorkflowCompletionMode.ALL_END,
            initial_state={"seed": 42},
            metadata={"owner": "team-a"},
        )
        step = FunctionStep("s", StepMetadata(name="s"), Msg, Out, _identity)
        wf = Workflow(cfg)
        wf.add_step(step)
        wf.set_start_step(step)
        wf.add_end_step(step)

        schema = workflow_to_schema(wf)
        assert schema.state_mode == WorkflowStateMode.STAGED
        assert schema.state_conflict_policy == WorkflowStateConflictPolicy.LAST_WRITE_WINS
        assert schema.completion_mode == WorkflowCompletionMode.ALL_END
        assert schema.initial_state == {"seed": 42}
        assert schema.metadata == {"owner": "team-a"}
        assert schema.description == "test workflow"

    def test_step_metadata_fields_preserved(self) -> None:
        meta = StepMetadata(
            name="careful_step",
            description="runs with retries",
            tags=["critical", "prod"],
            max_retries=3,
            timeout_seconds=30.0,
            join_failure_mode=JoinFailureMode.SKIP_FAILED,
        )
        step = FunctionStep("s", meta, Msg, Out, _identity)
        wf = Workflow(WorkflowConfig(name="meta_wf"))
        wf.add_step(step)
        wf.set_start_step(step)
        wf.add_end_step(step)

        schema = workflow_to_schema(wf)
        s = schema.steps[0]
        assert s.name == "careful_step"
        assert s.description == "runs with retries"
        assert s.tags == ["critical", "prod"]
        assert s.max_retries == 3
        assert s.timeout_seconds == 30.0
        assert s.join_failure_mode == JoinFailureMode.SKIP_FAILED

    def test_schema_round_trips_through_json(self) -> None:
        wf = _build_linear_function_workflow("json_rt")
        schema = workflow_to_schema(wf)
        json_str = schema.to_json()
        restored = WorkflowSchema.from_json(json_str)
        assert restored.name == "json_rt"
        assert restored.start_step == "step_a"
        assert len(restored.steps) == 2
        assert len(restored.edges) == 1


# ---------------------------------------------------------------------------
# Public export check
# ---------------------------------------------------------------------------

class TestPublicExports:
    def test_schema_symbols_importable_from_workflow_package(self) -> None:
        from agentbyte.workflow import (  # noqa: F401
            EdgeConditionSchema,
            EdgeSchema,
            StepKind,
            StepSchema,
            WorkflowBuilder,
            WorkflowLoader,
            WorkflowSchemaBuildError,
            WorkflowSchema,
            WorkflowStepRegistry,
            workflow_to_schema,
        )


class TestWorkflowBuilder:
    def test_builds_echo_workflow_from_schema(self) -> None:
        wf = Workflow(WorkflowConfig(name="echo_roundtrip"))
        echo = EchoStep(
            "echo",
            StepMetadata(name="echo"),
            prefix="Hello, ",
            suffix="!",
        )
        wf.add_step(echo)
        wf.set_start_step(echo)
        wf.add_end_step(echo)

        schema = workflow_to_schema(wf)
        rebuilt = WorkflowBuilder().from_schema(schema)

        assert rebuilt.config.name == "echo_roundtrip"
        assert rebuilt.start_step_id == "echo"
        assert rebuilt.end_step_ids == ["echo"]
        rebuilt_step = rebuilt.steps["echo"]
        assert isinstance(rebuilt_step, EchoStep)
        assert rebuilt_step.prefix == "Hello, "
        assert rebuilt_step.suffix == "!"

    def test_builds_http_workflow_from_schema(self) -> None:
        wf = Workflow(WorkflowConfig(name="http_roundtrip"))
        step = HttpStep("fetch", StepMetadata(name="fetch"))
        wf.add_step(step)
        wf.set_start_step(step)
        wf.add_end_step(step)

        rebuilt = WorkflowBuilder().from_schema(workflow_to_schema(wf))

        assert isinstance(rebuilt.steps["fetch"], HttpStep)

    def test_builds_transform_step_from_schema(self) -> None:
        class TransformInput(BaseModel):
            message: str
            count: int

        class TransformOutput(BaseModel):
            result: str
            total: int

        from agentbyte.workflow import TransformStep

        wf = Workflow(WorkflowConfig(name="transform_roundtrip"))
        step = TransformStep(
            "convert",
            StepMetadata(name="convert"),
            input_type=TransformInput,
            output_type=TransformOutput,
            mappings={"result": "message", "total": "count"},
        )
        wf.add_step(step)
        wf.set_start_step(step)
        wf.add_end_step(step)

        rebuilt = WorkflowBuilder().from_schema(workflow_to_schema(wf))

        rebuilt_step = rebuilt.steps["convert"]
        assert rebuilt_step.__class__.__name__ == "TransformStep"
        assert rebuilt_step.mappings == {"result": "message", "total": "count"}

    def test_requires_registry_for_function_steps(self) -> None:
        schema = workflow_to_schema(_build_linear_function_workflow("func_registry_needed"))

        with pytest.raises(WorkflowSchemaBuildError, match="Function step 'step_a'"):
            WorkflowBuilder().from_schema(schema)

    def test_uses_registry_to_build_function_steps(self) -> None:
        schema = workflow_to_schema(_build_linear_function_workflow("func_registry"))

        registry = WorkflowStepRegistry(
            function_by_id={
                "step_a": lambda step: FunctionStep(
                    step.id,
                    StepMetadata(name=step.name),
                    Msg,
                    Out,
                    _identity,
                ),
                "step_b": lambda step: FunctionStep(
                    step.id,
                    StepMetadata(name=step.name),
                    Out,
                    Out,
                    lambda data, _ctx: data,
                ),
            }
        )

        rebuilt = WorkflowBuilder(registry=registry).from_schema(schema)

        assert rebuilt.start_step_id == "step_a"
        assert list(rebuilt.steps) == ["step_a", "step_b"]

    def test_registry_rejects_wrong_step_id(self) -> None:
        schema = workflow_to_schema(_build_linear_function_workflow("func_bad_registry"))
        registry = WorkflowStepRegistry(
            function_by_id={
                "step_a": lambda step: FunctionStep(
                    "wrong_id",
                    StepMetadata(name=step.name),
                    Msg,
                    Out,
                    _identity,
                ),
            }
        )

        with pytest.raises(WorkflowSchemaBuildError, match="must match"):
            WorkflowBuilder(registry=registry).from_schema(schema)


class TestWorkflowLoader:
    def test_loads_workflow_from_json(self) -> None:
        wf = Workflow(WorkflowConfig(name="echo_json"))
        step = EchoStep("echo", StepMetadata(name="echo"), prefix=">", suffix="<")
        wf.add_step(step)
        wf.set_start_step(step)
        wf.add_end_step(step)

        schema = workflow_to_schema(wf)
        rebuilt = WorkflowLoader().from_json(schema.to_json())

        assert rebuilt.config.name == "echo_json"
        assert isinstance(rebuilt.steps["echo"], EchoStep)

    def test_loads_workflow_from_dict(self) -> None:
        wf = Workflow(WorkflowConfig(name="echo_dict"))
        step = EchoStep("echo", StepMetadata(name="echo"))
        wf.add_step(step)
        wf.set_start_step(step)
        wf.add_end_step(step)

        rebuilt = WorkflowLoader().from_dict(workflow_to_schema(wf).to_dict())
        assert rebuilt.start_step_id == "echo"

    def test_loads_workflow_from_yaml(self) -> None:
        wf = Workflow(WorkflowConfig(name="echo_yaml"))
        step = EchoStep("echo", StepMetadata(name="echo"), prefix="[", suffix="]")
        wf.add_step(step)
        wf.set_start_step(step)
        wf.add_end_step(step)

        yaml_text = yaml.safe_dump(workflow_to_schema(wf).to_dict(), sort_keys=False)
        rebuilt = WorkflowLoader().from_yaml(yaml_text)

        rebuilt_step = rebuilt.steps["echo"]
        assert isinstance(rebuilt_step, EchoStep)
        assert rebuilt_step.prefix == "["
        assert rebuilt_step.suffix == "]"
