"""
Tests for workflow context passing between steps.

Covers: step inputs, result routing, multi-level chaining,
parallel fan-in, variable substitution, original_input, and
missing input handling.
"""

import json
import pytest


@pytest.fixture(autouse=True)
def _isolate(tmp_path, monkeypatch):
    workflows = tmp_path / "workflows"
    defs = workflows / "definitions"
    runs = workflows / "runs"
    for d in (workflows, defs, runs):
        d.mkdir(parents=True, exist_ok=True)

    import brynq.server as srv
    monkeypatch.setattr(srv, "BROKER_DIR", tmp_path)

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


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


def _two_step_with_input():
    """A produces a result, B consumes it."""
    return [
        {"id": "analyze", "type": "task", "title": "Analyze data"},
        {
            "id": "review",
            "type": "task",
            "title": "Review analysis",
            "depends_on": ["analyze"],
            "inputs": {"analysis": "analyze"},
            "config": {"prompt": "Review this analysis:\n{analysis}"},
        },
    ]


def _three_step_chain():
    """A -> B -> C linear chain with context passing at each hop."""
    return [
        {
            "id": "extract",
            "type": "task",
            "title": "Extract data",
            "config": {"prompt": "Extract data from {original_input}"},
        },
        {
            "id": "transform",
            "type": "task",
            "title": "Transform data",
            "depends_on": ["extract"],
            "inputs": {"raw_data": "extract"},
            "config": {"prompt": "Transform this raw data:\n{raw_data}"},
        },
        {
            "id": "summarize",
            "type": "task",
            "title": "Summarize results",
            "depends_on": ["transform"],
            "inputs": {"transformed": "transform"},
            "config": {"prompt": "Summarize:\n{transformed}"},
        },
    ]


def _parallel_join():
    """Two parallel steps feed into a single join step."""
    return [
        {"id": "research_a", "type": "task", "title": "Research topic A"},
        {"id": "research_b", "type": "task", "title": "Research topic B"},
        {
            "id": "combine",
            "type": "task",
            "title": "Combine research",
            "depends_on": ["research_a", "research_b"],
            "inputs": {"topic_a": "research_a", "topic_b": "research_b"},
            "config": {
                "prompt": "Combine findings:\nA: {topic_a}\nB: {topic_b}"
            },
        },
    ]


# ── Test: Step Inputs Definition ──────────────────────────────────────────


class TestStepInputsDefinition:
    def test_inputs_preserved_in_workflow(self):
        from brynq.workflow import create_workflow
        wf = create_workflow("Context Test", steps=_two_step_with_input())
        assert "error" not in wf
        review_step = next(s for s in wf["steps"] if s["id"] == "review")
        assert review_step["inputs"] == {"analysis": "analyze"}

    def test_empty_inputs_default(self):
        from brynq.workflow import create_workflow
        steps = [{"id": "solo", "type": "task", "title": "Solo step"}]
        wf = create_workflow("No Inputs", steps=steps)
        assert wf["steps"][0]["inputs"] == {}

    def test_invalid_input_source_rejected(self):
        from brynq.workflow import create_workflow
        steps = [
            {"id": "a", "type": "task"},
            {
                "id": "b",
                "type": "task",
                "depends_on": ["a"],
                "inputs": {"data": "nonexistent_step"},
            },
        ]
        result = create_workflow("Bad Ref", steps=steps)
        assert "error" in result
        assert "nonexistent_step" in result["error"]


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


class TestResultRouting:
    def test_result_routed_to_downstream_step(self):
        from brynq.workflow import create_workflow, start_workflow, complete_step, get_run
        wf = create_workflow("Route", steps=_two_step_with_input())
        run = start_workflow(wf["workflow_id"])
        rid = run["run_id"]

        complete_step(rid, "analyze", "The data shows a 15% increase")
        updated = get_run(rid)

        resolved = updated["step_statuses"]["review"]["resolved_inputs"]
        assert resolved["analysis"] == "The data shows a 15% increase"

    def test_result_not_routed_to_uninvolved_step(self):
        from brynq.workflow import create_workflow, start_workflow, complete_step, get_run
        steps = [
            {"id": "a", "type": "task"},
            {"id": "b", "type": "task"},
            {
                "id": "c",
                "type": "task",
                "depends_on": ["a"],
                "inputs": {"from_a": "a"},
            },
        ]
        wf = create_workflow("Isolated", steps=steps)
        run = start_workflow(wf["workflow_id"])
        rid = run["run_id"]

        complete_step(rid, "a", "result from A")
        updated = get_run(rid)

        # b should NOT have any resolved inputs from a
        assert updated["step_statuses"]["b"]["resolved_inputs"] == {}
        # c should have the input
        assert updated["step_statuses"]["c"]["resolved_inputs"]["from_a"] == "result from A"

    def test_resolved_inputs_initialized_empty(self):
        from brynq.workflow import create_workflow, start_workflow, get_run
        wf = create_workflow("Init", steps=_two_step_with_input())
        run = start_workflow(wf["workflow_id"])
        updated = get_run(run["run_id"])

        for sid, status in updated["step_statuses"].items():
            assert "resolved_inputs" in status
            assert isinstance(status["resolved_inputs"], dict)


# ── Test: Multi-Level Chaining (A -> B -> C) ─────────────────────────────


class TestMultiLevelChaining:
    def test_three_step_chain_context_flows(self):
        from brynq.workflow import create_workflow, start_workflow, complete_step, get_run
        wf = create_workflow("Chain", steps=_three_step_chain())
        run = start_workflow(wf["workflow_id"])
        rid = run["run_id"]

        # Step 1: extract
        complete_step(rid, "extract", "raw CSV data with 100 rows")
        state = get_run(rid)
        assert state["step_statuses"]["transform"]["resolved_inputs"]["raw_data"] == "raw CSV data with 100 rows"

        # Step 2: transform
        complete_step(rid, "transform", "cleaned DataFrame, 95 valid rows")
        state = get_run(rid)
        assert state["step_statuses"]["summarize"]["resolved_inputs"]["transformed"] == "cleaned DataFrame, 95 valid rows"

        # Step 3: summarize
        result = complete_step(rid, "summarize", "Summary: 95% data quality")
        assert result["workflow_status"] == "completed"

    def test_intermediate_results_preserved(self):
        from brynq.workflow import create_workflow, start_workflow, complete_step, get_run
        wf = create_workflow("Preserve", steps=_three_step_chain())
        run = start_workflow(wf["workflow_id"])
        rid = run["run_id"]

        complete_step(rid, "extract", "step 1 output")
        complete_step(rid, "transform", "step 2 output")
        complete_step(rid, "summarize", "step 3 output")

        final = get_run(rid)
        # Each step's own result is stored
        assert final["step_statuses"]["extract"]["result"] == "step 1 output"
        assert final["step_statuses"]["transform"]["result"] == "step 2 output"
        assert final["step_statuses"]["summarize"]["result"] == "step 3 output"


# ── Test: Parallel Steps Feeding Into Join ────────────────────────────────


class TestParallelJoin:
    def test_parallel_results_merged_into_join(self):
        from brynq.workflow import create_workflow, start_workflow, complete_step, get_run
        wf = create_workflow("Parallel Join", steps=_parallel_join())
        run = start_workflow(wf["workflow_id"])
        rid = run["run_id"]

        complete_step(rid, "research_a", "Topic A findings: significant correlation")
        complete_step(rid, "research_b", "Topic B findings: no effect observed")

        state = get_run(rid)
        resolved = state["step_statuses"]["combine"]["resolved_inputs"]
        assert resolved["topic_a"] == "Topic A findings: significant correlation"
        assert resolved["topic_b"] == "Topic B findings: no effect observed"

    def test_partial_parallel_results(self):
        from brynq.workflow import create_workflow, start_workflow, complete_step, get_run
        wf = create_workflow("Partial", steps=_parallel_join())
        run = start_workflow(wf["workflow_id"])
        rid = run["run_id"]

        # Only complete one of the parallel steps
        complete_step(rid, "research_a", "A done")
        state = get_run(rid)

        resolved = state["step_statuses"]["combine"]["resolved_inputs"]
        assert resolved["topic_a"] == "A done"
        # topic_b not yet resolved
        assert "topic_b" not in resolved

    def test_join_step_ready_only_when_all_deps_complete(self):
        from brynq.workflow import create_workflow, start_workflow, complete_step, get_run
        wf = create_workflow("Ready Check", steps=_parallel_join())
        run = start_workflow(wf["workflow_id"])
        rid = run["run_id"]

        result1 = complete_step(rid, "research_a", "A result")
        assert "combine" not in result1["newly_ready"]

        result2 = complete_step(rid, "research_b", "B result")
        assert "combine" in result2["newly_ready"]


# ── Test: Variable Substitution ───────────────────────────────────────────


class TestVariableSubstitution:
    def test_basic_substitution(self):
        from brynq.workflow import (
            create_workflow, start_workflow, complete_step, resolve_step_prompt,
        )
        wf = create_workflow("Substitute", steps=_two_step_with_input())
        run = start_workflow(wf["workflow_id"])
        rid = run["run_id"]

        complete_step(rid, "analyze", "Revenue grew 20% YoY")
        result = resolve_step_prompt(rid, "review")
        assert "error" not in result
        assert result["prompt"] == "Review this analysis:\nRevenue grew 20% YoY"

    def test_original_input_substitution(self):
        from brynq.workflow import (
            create_workflow, start_workflow, resolve_step_prompt,
        )
        wf = create_workflow("Original", steps=_three_step_chain())
        run = start_workflow(
            wf["workflow_id"],
            inputs={"original_input": "quarterly_report.csv"},
        )
        rid = run["run_id"]

        result = resolve_step_prompt(rid, "extract")
        assert "error" not in result
        assert result["prompt"] == "Extract data from quarterly_report.csv"

    def test_multiple_variables_in_one_prompt(self):
        from brynq.workflow import (
            create_workflow, start_workflow, complete_step, resolve_step_prompt,
        )
        wf = create_workflow("Multi Var", steps=_parallel_join())
        run = start_workflow(wf["workflow_id"])
        rid = run["run_id"]

        complete_step(rid, "research_a", "Finding A")
        complete_step(rid, "research_b", "Finding B")

        result = resolve_step_prompt(rid, "combine")
        assert "error" not in result
        assert "Finding A" in result["prompt"]
        assert "Finding B" in result["prompt"]
        assert result["prompt"] == "Combine findings:\nA: Finding A\nB: Finding B"

    def test_unresolved_placeholder_left_as_is(self):
        from brynq.workflow import (
            create_workflow, start_workflow, resolve_step_prompt,
        )
        wf = create_workflow("Unresolved", steps=_two_step_with_input())
        run = start_workflow(wf["workflow_id"])
        rid = run["run_id"]

        # Don't complete "analyze" — the {analysis} placeholder stays
        result = resolve_step_prompt(rid, "review")
        assert "error" not in result
        assert "{analysis}" in result["prompt"]

    def test_no_prompt_template_returns_error(self):
        from brynq.workflow import (
            create_workflow, start_workflow, resolve_step_prompt,
        )
        steps = [
            {"id": "no_prompt", "type": "task", "title": "No prompt step"},
        ]
        wf = create_workflow("No Prompt", steps=steps)
        run = start_workflow(wf["workflow_id"])
        rid = run["run_id"]

        result = resolve_step_prompt(rid, "no_prompt")
        assert "error" in result
        assert "no prompt template" in result["error"]

    def test_resolve_nonexistent_run(self):
        from brynq.workflow import resolve_step_prompt
        result = resolve_step_prompt("fake_run", "fake_step")
        assert "error" in result

    def test_resolve_nonexistent_step(self):
        from brynq.workflow import (
            create_workflow, start_workflow, resolve_step_prompt,
        )
        steps = [{"id": "real", "type": "task"}]
        wf = create_workflow("Real", steps=steps)
        run = start_workflow(wf["workflow_id"])
        result = resolve_step_prompt(run["run_id"], "ghost_step")
        assert "error" in result


# ── Test: Original Input ──────────────────────────────────────────────────


class TestOriginalInput:
    def test_original_input_stored_on_run(self):
        from brynq.workflow import create_workflow, start_workflow, get_run
        steps = [{"id": "s1", "type": "task"}]
        wf = create_workflow("OI", steps=steps)
        run = start_workflow(
            wf["workflow_id"],
            inputs={"original_input": "Analyze Q4 earnings"},
        )
        loaded = get_run(run["run_id"])
        assert loaded["original_input"] == "Analyze Q4 earnings"

    def test_original_input_defaults_empty(self):
        from brynq.workflow import create_workflow, start_workflow, get_run
        steps = [{"id": "s1", "type": "task"}]
        wf = create_workflow("OI Default", steps=steps)
        run = start_workflow(wf["workflow_id"])
        loaded = get_run(run["run_id"])
        assert loaded["original_input"] == ""

    def test_original_input_available_in_all_steps(self):
        from brynq.workflow import (
            create_workflow, start_workflow, complete_step, resolve_step_prompt,
        )
        steps = [
            {
                "id": "first",
                "type": "task",
                "config": {"prompt": "Process: {original_input}"},
            },
            {
                "id": "second",
                "type": "task",
                "depends_on": ["first"],
                "config": {"prompt": "Continue with: {original_input}"},
            },
        ]
        wf = create_workflow("OI All Steps", steps=steps)
        run = start_workflow(
            wf["workflow_id"],
            inputs={"original_input": "user request"},
        )
        rid = run["run_id"]

        r1 = resolve_step_prompt(rid, "first")
        assert r1["prompt"] == "Process: user request"

        complete_step(rid, "first", "done")
        r2 = resolve_step_prompt(rid, "second")
        assert r2["prompt"] == "Continue with: user request"


# ── Test: Missing Input Handling ──────────────────────────────────────────


class TestMissingInputHandling:
    def test_step_with_inputs_before_source_completes(self):
        """Resolved inputs should be empty until the source step completes."""
        from brynq.workflow import create_workflow, start_workflow, get_run
        wf = create_workflow("Before", steps=_two_step_with_input())
        run = start_workflow(wf["workflow_id"])
        state = get_run(run["run_id"])
        assert state["step_statuses"]["review"]["resolved_inputs"] == {}

    def test_prompt_with_mix_of_resolved_and_unresolved(self):
        """When some inputs are resolved and others are not, only resolved
        ones get substituted."""
        from brynq.workflow import (
            create_workflow, start_workflow, complete_step, resolve_step_prompt,
        )
        wf = create_workflow("Mix", steps=_parallel_join())
        run = start_workflow(wf["workflow_id"])
        rid = run["run_id"]

        # Only complete research_a
        complete_step(rid, "research_a", "A result")
        result = resolve_step_prompt(rid, "combine")
        assert "A result" in result["prompt"]
        assert "{topic_b}" in result["prompt"]


# ── Test: Backward Compatibility ──────────────────────────────────────────


class TestBackwardCompatibility:
    def test_workflow_without_inputs_still_works(self):
        """Steps without inputs field should function identically to before."""
        from brynq.workflow import create_workflow, start_workflow, complete_step, get_run
        steps = [
            {"id": "plan", "type": "task", "title": "Plan"},
            {"id": "build", "type": "task", "title": "Build", "depends_on": ["plan"]},
            {"id": "test", "type": "task", "title": "Test", "depends_on": ["build"]},
        ]
        wf = create_workflow("Legacy", steps=steps)
        run = start_workflow(wf["workflow_id"])
        rid = run["run_id"]

        result1 = complete_step(rid, "plan", "Plan done")
        assert "build" in result1["newly_ready"]

        result2 = complete_step(rid, "build", "Build done")
        assert "test" in result2["newly_ready"]

        result3 = complete_step(rid, "test", "Test passed")
        assert result3["workflow_status"] == "completed"

    def test_existing_step_statuses_fields_preserved(self):
        """The new resolved_inputs field coexists with existing fields."""
        from brynq.workflow import create_workflow, start_workflow, complete_step, get_run
        steps = [{"id": "s1", "type": "task", "title": "Step 1"}]
        wf = create_workflow("Fields", steps=steps)
        run = start_workflow(wf["workflow_id"])
        rid = run["run_id"]

        complete_step(rid, "s1", "my result")
        state = get_run(rid)
        s1 = state["step_statuses"]["s1"]
        assert s1["status"] == "completed"
        assert s1["result"] == "my result"
        assert s1["completed_at"] is not None
        assert s1["error"] is None
        assert "resolved_inputs" in s1
