from __future__ import annotations

from pathlib import Path

from pydantic import BaseModel

from agentbyte.tools import FunctionTool
from agentbyte.workflow import StepMetadata, Workflow, WorkflowConfig, WorkflowVisualizer
from agentbyte.workflow.core.models import WorkflowContext
from agentbyte.workflow.steps.step import BaseStep


class InputModel(BaseModel):
    value: str


class OutputModel(BaseModel):
    result: str


class _DummyAgent:
    def __init__(self) -> None:
        self.name = "dummy_agent"
        self.tools = [FunctionTool(str.upper), FunctionTool(str.lower)]


class _AgentStep(BaseStep[InputModel, OutputModel]):
    def __init__(self, step_id: str) -> None:
        super().__init__(step_id, StepMetadata(name=step_id), InputModel, OutputModel)
        self.agent = _DummyAgent()

    async def execute(self, input_data: InputModel, context: WorkflowContext) -> OutputModel:  # noqa: ARG002
        return OutputModel(result=input_data.value)


class _SimpleStep(BaseStep[OutputModel, OutputModel]):
    def __init__(self, step_id: str) -> None:
        super().__init__(step_id, StepMetadata(name=step_id), OutputModel, OutputModel)

    async def execute(self, input_data: OutputModel, context: WorkflowContext) -> OutputModel:  # noqa: ARG002
        return input_data


def _build_workflow() -> Workflow:
    workflow = Workflow(WorkflowConfig(name="viz_test"))
    workflow.add_step(_AgentStep("planner"))
    workflow.add_step(_SimpleStep("coordinator"))
    workflow.set_start_step("planner")
    workflow.add_end_step("coordinator")
    workflow.add_edge(
        "planner",
        "coordinator",
        {"type": "output_based", "field": "result", "operator": "==", "value": "ok"},
    )
    return workflow


def test_mermaid_topology_contains_structure() -> None:
    mermaid = WorkflowVisualizer.to_mermaid(_build_workflow(), view="topology")

    assert "flowchart TD" in mermaid
    assert "planner [start]" in mermaid
    assert "coordinator [end]" in mermaid
    assert 'output_based: result == ok' in mermaid


def test_mermaid_detailed_contains_agent_tools() -> None:
    mermaid = WorkflowVisualizer.to_mermaid(_build_workflow(), view="detailed")

    assert "classDef agent" in mermaid
    assert "agent: dummy_agent" in mermaid
    assert "tool:" in mermaid
    assert "START" in mermaid
    assert "END" in mermaid


def test_save_writes_graph_file(tmp_path: Path) -> None:
    content = "flowchart TD\n  A-->B\n"
    out = WorkflowVisualizer.save(content, tmp_path / "graph.mmd")

    assert out.exists()
    assert out.read_text(encoding="utf-8") == content


def test_dot_generation() -> None:
    dot = WorkflowVisualizer.to_dot(_build_workflow(), view="detailed")

    assert "digraph" in dot
    assert "planner" in dot
    assert "coordinator" in dot
    assert "agent: dummy_agent" in dot


def test_mermaid_can_include_validation_warnings_as_comments() -> None:
    workflow = Workflow(WorkflowConfig(name="viz_warning_test"))
    workflow.add_step(_AgentStep("planner"))
    workflow.add_step(_SimpleStep("unused"))
    workflow.set_start_step("planner")

    mermaid = WorkflowVisualizer.to_mermaid(
        workflow,
        view="topology",
        include_validation=True,
    )

    assert "%% Validation warnings" in mermaid
    assert "%% Unreachable steps: unused" in mermaid


def test_mermaid_detailed_attaches_warning_nodes() -> None:
    workflow = Workflow(WorkflowConfig(name="viz_warning_nodes"))
    workflow.add_step(_AgentStep("planner"))
    workflow.add_step(_SimpleStep("coordinator"))
    workflow.set_start_step("planner")
    workflow.add_edge("planner", "coordinator")

    mermaid = WorkflowVisualizer.to_mermaid(
        workflow,
        view="detailed",
        include_validation=True,
    )

    assert "classDef warning" in mermaid
    assert "coordinator_warning_" in mermaid
    assert "-.warning.->" in mermaid
