"""
Tests for the Workflow Executor -- automated workflow step dispatch.

Covers:
  - Executor picks up ready steps and dispatches them
  - Result routing triggers next step
  - Full 3-step chain end-to-end with mock LLM
  - Error handling (step failure does not crash executor)
  - Task-mode dispatch creates broker tasks with workflow metadata
  - Parallel step dispatch
"""

import json
from pathlib import Path

import pytest


@pytest.fixture(autouse=True)
def _isolate(tmp_path, monkeypatch):
    """Redirect all workflow and broker storage to a temp directory."""
    channels = tmp_path / "channels"
    sessions = tmp_path / "sessions"
    cursors = tmp_path / "cursors"
    tasks = tmp_path / "tasks"
    workflows = tmp_path / "workflows"
    definitions = workflows / "definitions"
    runs = workflows / "runs"

    for d in (channels, sessions, cursors, tasks, workflows, definitions, runs):
        d.mkdir(parents=True, exist_ok=True)

    import brynq.server as srv
    monkeypatch.setattr(srv, "BROKER_DIR", tmp_path)
    monkeypatch.setattr(srv, "CHANNELS_DIR", channels)
    monkeypatch.setattr(srv, "SESSIONS_DIR", sessions)
    monkeypatch.setattr(srv, "CURSORS_DIR", cursors)

    import brynq.workflow as wf
    monkeypatch.setattr(wf, "WORKFLOWS_DIR", workflows)
    monkeypatch.setattr(wf, "DEFINITIONS_DIR", definitions)
    monkeypatch.setattr(wf, "RUNS_DIR", runs)

    import brynq.workflow_executor as we
    monkeypatch.setattr(we, "TASKS_DIR", tasks)
    monkeypatch.setattr(we, "RUNS_DIR", runs)


# ── Helpers ───────────────────────────────────────────────────────────────


def _linear_steps():
    """A -> B -> C linear chain with prompt templates."""
    return [
        {
            "id": "analyze",
            "type": "task",
            "title": "Analyze the problem",
            "depends_on": [],
            "config": {"prompt": "Analyze the following problem."},
        },
        {
            "id": "implement",
            "type": "task",
            "title": "Implement solution",
            "depends_on": ["analyze"],
            "inputs": {"analysis": "analyze"},
            "config": {"prompt": "Implement a solution based on this analysis:\n{analysis}"},
        },
        {
            "id": "review",
            "type": "task",
            "title": "Review implementation",
            "depends_on": ["implement"],
            "inputs": {"code": "implement"},
            "config": {"prompt": "Review this implementation:\n{code}"},
        },
    ]


def _parallel_steps():
    """A and B run in parallel, then C depends on both."""
    return [
        {
            "id": "research",
            "type": "task",
            "title": "Research topic",
            "depends_on": [],
            "config": {"prompt": "Research the topic."},
        },
        {
            "id": "gather_data",
            "type": "task",
            "title": "Gather data",
            "depends_on": [],
            "config": {"prompt": "Gather relevant data."},
        },
        {
            "id": "synthesize",
            "type": "task",
            "title": "Synthesize findings",
            "depends_on": ["research", "gather_data"],
            "inputs": {"research": "research", "data": "gather_data"},
            "config": {"prompt": "Synthesize:\nResearch: {research}\nData: {data}"},
        },
    ]


def _make_mock_llm(responses=None, default="Mock LLM output"):
    """Create a mock LLM function that tracks calls.

    Args:
        responses: dict mapping prompt substrings to responses.
        default: fallback response.

    Returns:
        (mock_fn, call_log) tuple.
    """
    responses = responses or {}
    call_log = []

    def mock_fn(messages):
        prompt = messages[-1]["content"] if messages else ""
        call_log.append(prompt)
        for trigger, resp in responses.items():
            if trigger in prompt:
                return resp
        return default

    return mock_fn, call_log


# ── Test: Basic Dispatch ──────────────────────────────────────────────────


class TestBasicDispatch:
    def test_executor_picks_up_ready_step(self):
        """Executor detects a ready step and dispatches it."""
        from brynq.workflow import create_workflow, start_workflow
        from brynq.workflow_executor import WorkflowExecutor

        mock_fn, calls = _make_mock_llm(default="Analysis complete")
        wf = create_workflow("Test", steps=_linear_steps())
        start_workflow(wf["workflow_id"])

        executor = WorkflowExecutor(llm_fn=mock_fn)
        dispatched = executor.tick_once()

        assert dispatched == 1
        assert len(calls) == 1
        assert "Analyze" in calls[0]

    def test_executor_skips_already_dispatched_step(self):
        """Executor does not re-dispatch a step that is already running."""
        from brynq.workflow import create_workflow, start_workflow
        from brynq.workflow_executor import WorkflowExecutor

        mock_fn, calls = _make_mock_llm(default="Done")
        wf = create_workflow("Test", steps=_linear_steps())
        run = start_workflow(wf["workflow_id"])

        # Executor dispatches step "analyze" -- since it's direct mode,
        # it completes immediately and advances, so "implement" becomes ready.
        executor = WorkflowExecutor(llm_fn=mock_fn)
        executor.tick_once()

        # Second tick should dispatch "implement", not "analyze" again
        executor.tick_once()
        assert len(calls) == 2

    def test_executor_does_nothing_when_no_runs(self):
        """Executor returns 0 when there are no running workflow runs."""
        from brynq.workflow_executor import WorkflowExecutor

        mock_fn, calls = _make_mock_llm()
        executor = WorkflowExecutor(llm_fn=mock_fn)
        dispatched = executor.tick_once()

        assert dispatched == 0
        assert len(calls) == 0


# ── Test: Result Routing ──────────────────────────────────────────────────


class TestResultRouting:
    def test_result_routes_to_next_step(self):
        """Completing a step feeds its result into the next step's prompt."""
        from brynq.workflow import create_workflow, start_workflow, get_run
        from brynq.workflow_executor import WorkflowExecutor

        mock_fn, calls = _make_mock_llm(
            responses={
                "Analyze": "Found a bug in module X",
                "analysis": "Fix applied to module X",
            },
            default="Looks good",
        )

        wf = create_workflow("Chain", steps=_linear_steps())
        run = start_workflow(wf["workflow_id"])
        run_id = run["run_id"]

        executor = WorkflowExecutor(llm_fn=mock_fn)

        # Tick 1: dispatch "analyze"
        executor.tick_once()
        # Tick 2: dispatch "implement" (with analysis result in prompt)
        executor.tick_once()

        assert len(calls) >= 2
        # The second call should contain the result from the first step
        assert "Found a bug in module X" in calls[1]

    def test_completed_step_unlocks_downstream(self):
        """After a step completes, downstream steps become ready."""
        from brynq.workflow import create_workflow, start_workflow, get_run
        from brynq.workflow_executor import WorkflowExecutor

        mock_fn, _ = _make_mock_llm(default="Done")
        wf = create_workflow("Unlock", steps=_linear_steps())
        run = start_workflow(wf["workflow_id"])
        run_id = run["run_id"]

        executor = WorkflowExecutor(llm_fn=mock_fn)
        executor.tick_once()  # Completes "analyze"

        # Check that "implement" is now ready
        fresh_run = get_run(run_id)
        assert fresh_run["step_statuses"]["implement"]["status"] == "ready"


# ── Test: Full Chain E2E ──────────────────────────────────────────────────


class TestFullChainE2E:
    def test_three_step_chain_completes(self):
        """Full 3-step linear chain executes end-to-end."""
        from brynq.workflow import create_workflow, start_workflow, get_run
        from brynq.workflow_executor import WorkflowExecutor

        call_count = [0]

        def mock_fn(messages):
            call_count[0] += 1
            prompt = messages[-1]["content"]
            if "Analyze" in prompt:
                return "Analysis: issue in auth module"
            elif "analysis" in prompt.lower() or "implement" in prompt.lower():
                return "Implementation: fixed auth validation"
            elif "review" in prompt.lower() or "code" in prompt.lower():
                return "Review: LGTM, ship it"
            return f"Step {call_count[0]} output"

        wf = create_workflow("E2E", steps=_linear_steps())
        run = start_workflow(wf["workflow_id"])
        run_id = run["run_id"]

        executor = WorkflowExecutor(llm_fn=mock_fn)

        # Run enough ticks to complete all 3 steps
        for _ in range(5):
            executor.tick_once()

        # Verify the workflow completed
        final_run = get_run(run_id)
        assert final_run["status"] == "completed"
        assert final_run["steps_completed"] == 3

        # Verify all steps have results
        for step_id in ("analyze", "implement", "review"):
            assert final_run["step_statuses"][step_id]["status"] == "completed"
            assert final_run["step_statuses"][step_id]["result"] is not None

        # Verify LLM was called 3 times (once per step)
        assert call_count[0] == 3

    def test_parallel_then_merge(self):
        """Two parallel steps complete, then merge step executes."""
        from brynq.workflow import create_workflow, start_workflow, get_run
        from brynq.workflow_executor import WorkflowExecutor

        mock_fn, calls = _make_mock_llm(
            responses={
                "Research": "Research findings: topic is interesting",
                "Gather": "Data gathered: 100 records",
            },
            default="Synthesis complete",
        )

        wf = create_workflow("Parallel", steps=_parallel_steps())
        run = start_workflow(wf["workflow_id"])
        run_id = run["run_id"]

        executor = WorkflowExecutor(llm_fn=mock_fn)

        # Run enough ticks for all steps
        for _ in range(5):
            executor.tick_once()

        final_run = get_run(run_id)
        assert final_run["status"] == "completed"
        assert final_run["steps_completed"] == 3

        # The synthesize step should have received both upstream results
        synth_inputs = final_run["step_statuses"]["synthesize"].get("resolved_inputs", {})
        assert "research" in synth_inputs
        assert "data" in synth_inputs


# ── Test: Error Handling ──────────────────────────────────────────────────


class TestErrorHandling:
    def test_llm_error_fails_step_not_executor(self):
        """LLM error fails the step but the executor continues running."""
        from brynq.workflow import create_workflow, start_workflow, get_run
        from brynq.workflow_executor import WorkflowExecutor

        def failing_llm(messages):
            raise RuntimeError("LLM is down")

        wf = create_workflow("Fail", steps=_linear_steps())
        run = start_workflow(wf["workflow_id"])
        run_id = run["run_id"]

        executor = WorkflowExecutor(llm_fn=failing_llm)

        # Should not raise
        dispatched = executor.tick_once()

        # Step should be failed, workflow should be failed
        final_run = get_run(run_id)
        assert final_run["step_statuses"]["analyze"]["status"] == "failed"
        assert final_run["status"] == "failed"

    def test_executor_survives_corrupted_run_file(self, tmp_path, monkeypatch):
        """Executor gracefully skips unreadable run files."""
        from brynq.workflow_executor import WorkflowExecutor
        import brynq.workflow as wf

        # Write a corrupted run file
        runs_dir = wf.RUNS_DIR
        corrupted = runs_dir / "corrupted.json"
        corrupted.write_text("NOT VALID JSON{{{", encoding="utf-8")

        mock_fn, calls = _make_mock_llm()
        executor = WorkflowExecutor(llm_fn=mock_fn)

        # Should not raise
        dispatched = executor.tick_once()
        assert dispatched == 0
        assert len(calls) == 0

    def test_step_failure_does_not_crash_other_runs(self):
        """A failing step in one run does not affect other runs."""
        from brynq.workflow import create_workflow, start_workflow, get_run
        from brynq.workflow_executor import WorkflowExecutor

        call_count = [0]

        def selective_llm(messages):
            call_count[0] += 1
            prompt = messages[-1]["content"]
            if call_count[0] == 1:
                raise RuntimeError("First call fails")
            return "Success"

        # Create two identical workflows
        wf1 = create_workflow("WF1", steps=[
            {"id": "s1", "type": "task", "title": "Step 1",
             "config": {"prompt": "Do step 1"}},
        ])
        wf2 = create_workflow("WF2", steps=[
            {"id": "s1", "type": "task", "title": "Step 1",
             "config": {"prompt": "Do step 1"}},
        ])
        run1 = start_workflow(wf1["workflow_id"])
        run2 = start_workflow(wf2["workflow_id"])

        executor = WorkflowExecutor(llm_fn=selective_llm)
        executor.tick_once()  # First run fails, second succeeds

        r1 = get_run(run1["run_id"])
        r2 = get_run(run2["run_id"])

        # One should be failed, the other completed
        statuses = {r1["status"], r2["status"]}
        assert "failed" in statuses
        assert "completed" in statuses


# ── Test: Task-Mode Dispatch ──────────────────────────────────────────────


class TestTaskDispatch:
    def test_task_mode_creates_broker_task(self):
        """In task mode, the executor creates a broker task file."""
        from brynq.workflow import create_workflow, start_workflow
        from brynq.workflow_executor import WorkflowExecutor
        import brynq.workflow_executor as we

        mock_fn, _ = _make_mock_llm()  # Won't be called in task mode
        wf = create_workflow("TaskMode", steps=[
            {"id": "s1", "type": "task", "title": "Do something",
             "config": {"prompt": "Execute this."}},
        ])
        start_workflow(wf["workflow_id"])

        executor = WorkflowExecutor(dispatch_mode="task", llm_fn=mock_fn)
        dispatched = executor.tick_once()

        assert dispatched == 1

        # Check that a task file was created
        task_files = list(we.TASKS_DIR.glob("*.json"))
        assert len(task_files) == 1

        task = json.loads(task_files[0].read_text("utf-8"))
        assert task["status"] == "pending"
        assert task["workflow_run_id"] == wf["workflow_id"] or task.get("workflow_run_id")
        assert task["step_id"] == "s1"
        assert "[WF]" in task["title"]

    def test_task_mode_includes_resolved_prompt(self):
        """Task description contains the resolved prompt."""
        from brynq.workflow import create_workflow, start_workflow
        from brynq.workflow_executor import WorkflowExecutor
        import brynq.workflow_executor as we

        mock_fn, _ = _make_mock_llm()
        wf = create_workflow("Prompt", steps=[
            {"id": "s1", "type": "task", "title": "Step 1",
             "config": {"prompt": "Analyze the codebase for bugs."}},
        ])
        start_workflow(wf["workflow_id"])

        executor = WorkflowExecutor(dispatch_mode="task", llm_fn=mock_fn)
        executor.tick_once()

        task_files = list(we.TASKS_DIR.glob("*.json"))
        task = json.loads(task_files[0].read_text("utf-8"))
        assert "Analyze the codebase for bugs" in task["description"]


# ── Test: Prompt Resolution ───────────────────────────────────────────────


class TestPromptResolution:
    def test_fallback_to_title_when_no_prompt(self):
        """Steps without a prompt template use their title as the prompt."""
        from brynq.workflow import create_workflow, start_workflow, get_run
        from brynq.workflow_executor import WorkflowExecutor

        mock_fn, calls = _make_mock_llm(default="Done")
        wf = create_workflow("NoPrompt", steps=[
            {"id": "s1", "type": "task", "title": "Run the tests"},
        ])
        start_workflow(wf["workflow_id"])

        executor = WorkflowExecutor(llm_fn=mock_fn)
        executor.tick_once()

        assert len(calls) == 1
        assert "Run the tests" in calls[0]

    def test_variable_substitution_in_prompt(self):
        """Prompt templates get variables substituted from upstream results."""
        from brynq.workflow import create_workflow, start_workflow
        from brynq.workflow_executor import WorkflowExecutor

        mock_fn, calls = _make_mock_llm(
            responses={"First step": "result_from_step_a"},
            default="Final answer",
        )

        steps = [
            {"id": "a", "type": "task", "title": "First step",
             "config": {"prompt": "First step prompt"}},
            {"id": "b", "type": "task", "title": "Second step",
             "depends_on": ["a"],
             "inputs": {"prev_result": "a"},
             "config": {"prompt": "Use this: {prev_result}"}},
        ]
        wf = create_workflow("Vars", steps=steps)
        start_workflow(wf["workflow_id"])

        executor = WorkflowExecutor(llm_fn=mock_fn)
        executor.tick_once()  # Step a
        executor.tick_once()  # Step b

        assert len(calls) == 2
        assert "result_from_step_a" in calls[1]
