"""
Tests for the DAG executor.

Uses a mock LLM that returns deterministic answers. Covers:
- linear pipeline
- parallel pipeline
- mixed DAG
- retry with success
- retries exhausted → FAILED
- dependency skip
- resume from a checkpoint
- rate-limiter delay
"""

import asyncio
import sys
import time
from pathlib import Path

# path to the source file (tests run from tests/)
sys.path.insert(0, str(Path(__file__).parent.parent))

from src.pipeline.dag_executor import (
    DAGExecutor,
    RateLimitException,
)
from src.pipeline.models import (
    HarvestContext,
    SubTask,
    TaskPlan,
    TaskState,
)


# ─── Mock-LLM ──────────────────────────────────────────────────────

class MockLLM:
    """Deterministic mock LLM for tests.

    Behaviour configurable per task_id:
    - "success" (default): returns dummy output immediately
    - "fail_then_succeed": fails on the first attempt, succeeds on the second
    - "always_fail": always raises
    - "rate_limit_then_succeed": once RateLimit, then success
    - "delay:N": wait N seconds, then success
    - "json": returns JSON output
    """

    def __init__(self, behavior: dict[str, str] = None):
        self.behavior = behavior or {}
        self.call_log: list[tuple[str, int]] = []  # (task_id, attempt)
        self.call_counts: dict[str, int] = {}

    async def call(self, prompt: str, meta: dict) -> tuple[str, dict]:
        """Mock implementation of the LLM interface."""
        task_id = meta.get("task_id", "unknown")
        self.call_counts[task_id] = self.call_counts.get(task_id, 0) + 1
        attempt = self.call_counts[task_id]
        self.call_log.append((task_id, attempt))

        behavior = self.behavior.get(task_id, "success")

        if behavior == "success":
            return f"Output for {task_id}", {"tokens_used": 100}

        if behavior == "always_fail":
            raise RuntimeError(f"mock error for {task_id}")

        if behavior == "fail_then_succeed":
            if attempt == 1:
                raise RuntimeError(f"let the first attempt fail ({task_id})")
            return f"Output for {task_id} nach Retry", {"tokens_used": 100}

        if behavior == "rate_limit_then_succeed":
            if attempt == 1:
                raise RateLimitException(f"429 for {task_id}")
            return f"Output for {task_id} nach Rate-Limit", {"tokens_used": 100}

        if behavior.startswith("delay:"):
            seconds = float(behavior.split(":")[1])
            await asyncio.sleep(seconds)
            return f"Output for {task_id} nach {seconds}s", {"tokens_used": 100}

        if behavior == "json":
            return '```json\n{"result": "ok", "value": 42}\n```', {"tokens_used": 100}

        return f"Default output for {task_id}", {"tokens_used": 100}


# ─── Test helpers ───────────────────────────────────────────────────

def make_context() -> HarvestContext:
    """Create an empty HarvestContext."""
    return HarvestContext(query="test")


def make_task(task_id: str, deps: list[str] = None, max_retries: int = 0) -> SubTask:
    """Helper for creating a SubTask."""
    return SubTask(
        id=task_id,
        phase="test",
        description=f"Test {task_id}",
        prompt_template="dummy",
        prompt_params={"task_id": task_id},
        depends_on=deps or [],
        max_retries=max_retries,
    )


PROMPTS = {"dummy": "Task: {task_id}"}


# ─── Tests ─────────────────────────────────────────────────────────

async def test_linear_pipeline():
    """3 tasks in sequence: A → B → C."""
    print("Test: linear pipeline ...", end=" ")
    plan = TaskPlan(use_case="test", tasks=[
        make_task("A"),
        make_task("B", deps=["A"]),
        make_task("C", deps=["B"]),
    ])
    mock = MockLLM()
    executor = DAGExecutor(llm_call=mock.call, max_parallel=5, rate_limit=100)
    ctx = make_context()
    results = await executor.execute(plan, PROMPTS, ctx)

    assert len(results) == 3
    assert all(r.state == TaskState.DONE for r in results.values())
    assert ctx.completed_layers == [0, 1, 2]
    # linear order: A before B before C
    a_time = results["A"].finished_at
    b_time = results["B"].started_at
    assert a_time <= b_time, f"A finished at {a_time}, B started at {b_time}"
    print("✓")


async def test_parallel_layer():
    """5 tasks in parallel in one layer."""
    print("Test: parallel layer ...", end=" ")
    plan = TaskPlan(use_case="test", tasks=[
        make_task(f"T{i}") for i in range(5)
    ])
    mock = MockLLM({f"T{i}": "delay:0.1" for i in range(5)})
    executor = DAGExecutor(llm_call=mock.call, max_parallel=5, rate_limit=100)
    ctx = make_context()

    start = time.monotonic()
    await executor.execute(plan, PROMPTS, ctx)
    elapsed = time.monotonic() - start

    # 5 parallel tasks of 0.1 s each should run in ~0.1-0.3 s, not 0.5 s
    assert elapsed < 0.4, f"parallelism does not work: {elapsed:.2f}s"
    assert all(r.state == TaskState.DONE for r in ctx.task_results.values())
    print(f"✓ ({elapsed:.2f}s for 5 tasks of 0.1s)")


async def test_max_parallel_limit():
    """With max_parallel=2, 4 tasks of 0.1 s each should take ~0.2 s, not 0.1 s."""
    print("Test: max_parallel-Limit ...", end=" ")
    plan = TaskPlan(use_case="test", tasks=[
        make_task(f"T{i}") for i in range(4)
    ])
    mock = MockLLM({f"T{i}": "delay:0.1" for i in range(4)})
    executor = DAGExecutor(llm_call=mock.call, max_parallel=2, rate_limit=100)
    ctx = make_context()

    start = time.monotonic()
    await executor.execute(plan, PROMPTS, ctx)
    elapsed = time.monotonic() - start

    # 4 tasks with at most 2 in parallel of 0.1 s each = ~0.2 s
    assert 0.15 < elapsed < 0.35, f"max_parallel=2 does not hold: {elapsed:.2f}s"
    print(f"✓ ({elapsed:.2f}s)")


async def test_mixed_dag():
    """A → (B, C) → D — a DAG with branching."""
    print("Test: mixed DAG (A → B,C → D) ...", end=" ")
    plan = TaskPlan(use_case="test", tasks=[
        make_task("A"),
        make_task("B", deps=["A"]),
        make_task("C", deps=["A"]),
        make_task("D", deps=["B", "C"]),
    ])
    mock = MockLLM()
    executor = DAGExecutor(llm_call=mock.call, max_parallel=5, rate_limit=100)
    ctx = make_context()
    await executor.execute(plan, PROMPTS, ctx)

    assert all(r.state == TaskState.DONE for r in ctx.task_results.values())
    assert ctx.completed_layers == [0, 1, 2]  # 3 layers
    # B and C after A
    assert ctx.task_results["B"].started_at >= ctx.task_results["A"].finished_at
    # D after B and C
    assert ctx.task_results["D"].started_at >= ctx.task_results["B"].finished_at
    assert ctx.task_results["D"].started_at >= ctx.task_results["C"].finished_at
    print("✓")


async def test_retry_success():
    """The task fails on the first attempt, succeeds on the second."""
    print("Test: retry with success ...", end=" ")
    plan = TaskPlan(use_case="test", tasks=[
        make_task("R", max_retries=2),
    ])
    mock = MockLLM({"R": "fail_then_succeed"})
    executor = DAGExecutor(llm_call=mock.call, max_parallel=5, rate_limit=100)
    ctx = make_context()
    await executor.execute(plan, PROMPTS, ctx)

    result = ctx.task_results["R"]
    assert result.state == TaskState.DONE, f"State: {result.state}"
    assert result.retries_used == 1, f"Retries: {result.retries_used}"
    assert mock.call_counts["R"] == 2
    print("✓")


async def test_retry_exhausted():
    """The task always fails, the retries are used up."""
    print("Test: retries exhausted → FAILED ...", end=" ")
    plan = TaskPlan(use_case="test", tasks=[
        make_task("F", max_retries=2),
    ])
    mock = MockLLM({"F": "always_fail"})
    executor = DAGExecutor(llm_call=mock.call, max_parallel=5, rate_limit=100)
    ctx = make_context()
    await executor.execute(plan, PROMPTS, ctx)

    result = ctx.task_results["F"]
    assert result.state == TaskState.FAILED
    assert "mock error" in result.error
    assert mock.call_counts["F"] == 3  # 1 + 2 retries
    print("✓")


async def test_dependency_skip():
    """The task is SKIPPED if a dependency FAILED."""
    print("Test: dependency skip ...", end=" ")
    plan = TaskPlan(use_case="test", tasks=[
        make_task("A"),
        make_task("B", deps=["A"]),  # depends on A
        make_task("C", deps=["B"]),  # depends on B
    ])
    mock = MockLLM({"A": "always_fail"})
    executor = DAGExecutor(llm_call=mock.call, max_parallel=5, rate_limit=100)
    ctx = make_context()
    await executor.execute(plan, PROMPTS, ctx)

    assert ctx.task_results["A"].state == TaskState.FAILED
    assert ctx.task_results["B"].state == TaskState.SKIPPED
    assert ctx.task_results["C"].state == TaskState.SKIPPED
    # B and C must not have been called
    assert "B" not in mock.call_counts
    assert "C" not in mock.call_counts
    print("✓")


async def test_rate_limit_recovery():
    """RateLimitException → back-off → retry → success."""
    print("Test: Rate-Limit Recovery ...", end=" ")
    plan = TaskPlan(use_case="test", tasks=[
        make_task("RL", max_retries=2),
    ])
    mock = MockLLM({"RL": "rate_limit_then_succeed"})
    executor = DAGExecutor(
        llm_call=mock.call,
        max_parallel=5,
        rate_limit=100,
    )
    # shorten the back-off for the test
    executor.rate_limiter.backoff_initial = 0.05
    ctx = make_context()
    await executor.execute(plan, PROMPTS, ctx)

    result = ctx.task_results["RL"]
    assert result.state == TaskState.DONE
    assert mock.call_counts["RL"] == 2
    print("✓")


async def test_resume_from_layer():
    """Resume: the first layer is skipped, earlier results stay."""
    print("Test: resume from a checkpoint ...", end=" ")
    plan = TaskPlan(use_case="test", tasks=[
        make_task("A"),
        make_task("B", deps=["A"]),
    ])
    # first run normally
    mock1 = MockLLM()
    executor1 = DAGExecutor(llm_call=mock1.call, max_parallel=5, rate_limit=100)
    ctx = make_context()
    await executor1.execute(plan, PROMPTS, ctx)
    assert mock1.call_counts == {"A": 1, "B": 1}
    assert ctx.completed_layers == [0, 1]

    # Now resume from layer 1 (i.e. only B anew, keep A)
    # a new executor with a different mock, so that we can measure who is called
    mock2 = MockLLM()
    executor2 = DAGExecutor(llm_call=mock2.call, max_parallel=5, rate_limit=100)
    await executor2.execute(plan, PROMPTS, ctx, resume_from_layer=1)

    # A must not have been called again
    assert "A" not in mock2.call_counts, f"A was called unexpectedly: {mock2.call_counts}"
    # B is run again
    assert mock2.call_counts == {"B": 1}, f"Mock2 calls: {mock2.call_counts}"
    print("✓")


async def test_progress_callback():
    """Progress events are sent correctly."""
    print("Test: Progress-Callback ...", end=" ")
    events = []

    async def cb(event: str, data: dict):
        events.append((event, data))

    plan = TaskPlan(use_case="test", tasks=[
        make_task("A"),
        make_task("B", deps=["A"]),
    ])
    mock = MockLLM()
    executor = DAGExecutor(
        llm_call=mock.call, max_parallel=5, rate_limit=100,
        progress_callback=cb,
    )
    ctx = make_context()
    await executor.execute(plan, PROMPTS, ctx)

    event_types = [e[0] for e in events]
    assert "layer_start" in event_types
    assert "task_start" in event_types
    assert "task_done" in event_types
    assert "layer_done" in event_types
    print(f"✓ ({len(events)} Events)")


async def test_persist_callback():
    """The persist callback is called after every layer."""
    print("Test: Persist-Callback ...", end=" ")
    persist_calls = []

    async def persist(ctx: HarvestContext, layer_idx: int):
        persist_calls.append((layer_idx, len(ctx.task_results)))

    plan = TaskPlan(use_case="test", tasks=[
        make_task("A"),
        make_task("B", deps=["A"]),
        make_task("C", deps=["B"]),
    ])
    mock = MockLLM()
    executor = DAGExecutor(
        llm_call=mock.call, max_parallel=5, rate_limit=100,
        persist_callback=persist,
    )
    ctx = make_context()
    await executor.execute(plan, PROMPTS, ctx)

    assert len(persist_calls) == 3  # 3 layers
    assert [c[0] for c in persist_calls] == [0, 1, 2]
    print("✓")


async def test_json_parsing():
    """JSON output is parsed automatically."""
    print("Test: JSON-Parsing ...", end=" ")
    plan = TaskPlan(use_case="test", tasks=[make_task("J")])
    mock = MockLLM({"J": "json"})
    executor = DAGExecutor(llm_call=mock.call, max_parallel=5, rate_limit=100)
    ctx = make_context()
    await executor.execute(plan, PROMPTS, ctx)

    result = ctx.task_results["J"]
    assert result.state == TaskState.DONE
    assert result.parsed_output == {"result": "ok", "value": 42}
    print("✓")


async def test_invalid_max_parallel():
    """The constructor rejects an invalid max_parallel."""
    print("Test: Invalid max_parallel ...", end=" ")
    mock = MockLLM()
    try:
        DAGExecutor(llm_call=mock.call, max_parallel=0)
        assert False, "should raise an exception"
    except ValueError:
        pass
    try:
        DAGExecutor(llm_call=mock.call, max_parallel=20)
        assert False, "should raise an exception"
    except ValueError:
        pass
    print("✓")


async def test_stop_check_between_layers():
    """A stop check between layers aborts the pipeline cleanly."""
    print("Test: stop between layers ...", end=" ")
    plan = TaskPlan(use_case="test", tasks=[
        make_task("A"),
        make_task("B", deps=["A"]),
        make_task("C", deps=["B"]),
    ])

    # stop flag that becomes active after the first layer
    stop_flag = {"stop": False}

    async def stop_after_first(event: str, data: dict):
        if event == "layer_done" and data.get("layer") == 1:
            stop_flag["stop"] = True

    mock = MockLLM()
    executor = DAGExecutor(
        llm_call=mock.call,
        max_parallel=5,
        rate_limit=100,
        progress_callback=stop_after_first,
        stop_check=lambda: stop_flag["stop"],
    )
    ctx = make_context()
    await executor.execute(plan, PROMPTS, ctx)

    # A must have run
    assert ctx.task_results["A"].state == TaskState.DONE
    # B and C must be SKIPPED
    assert ctx.task_results["B"].state == TaskState.SKIPPED, \
        f"B state: {ctx.task_results['B'].state}"
    assert ctx.task_results["C"].state == TaskState.SKIPPED
    # B must not have been called
    assert "B" not in mock.call_counts
    assert "C" not in mock.call_counts
    print("✓")


async def test_stop_before_first_layer():
    """Stop before the first layer → all tasks SKIPPED."""
    print("Test: stop before the first layer ...", end=" ")
    plan = TaskPlan(use_case="test", tasks=[
        make_task("A"),
        make_task("B", deps=["A"]),
    ])
    mock = MockLLM()
    executor = DAGExecutor(
        llm_call=mock.call,
        max_parallel=5,
        rate_limit=100,
        stop_check=lambda: True,  # abort immediately
    )
    ctx = make_context()
    await executor.execute(plan, PROMPTS, ctx)

    # all tasks must be SKIPPED
    assert ctx.task_results["A"].state == TaskState.SKIPPED
    assert ctx.task_results["B"].state == TaskState.SKIPPED
    # no LLM calls
    assert not mock.call_counts
    print("✓")


# ─── Test-Runner ───────────────────────────────────────────────────

async def main():
    print("=" * 60)
    print("DAG-Executor Tests")
    print("=" * 60)
    tests = [
        test_linear_pipeline,
        test_parallel_layer,
        test_max_parallel_limit,
        test_mixed_dag,
        test_retry_success,
        test_retry_exhausted,
        test_dependency_skip,
        test_rate_limit_recovery,
        test_resume_from_layer,
        test_progress_callback,
        test_persist_callback,
        test_json_parsing,
        test_invalid_max_parallel,
        test_stop_check_between_layers,
        test_stop_before_first_layer,
    ]
    failed = 0
    for test in tests:
        try:
            await test()
        except AssertionError as e:
            print(f"❌ {test.__name__}: {e}")
            failed += 1
        except Exception as e:
            print(f"💥 {test.__name__}: {type(e).__name__}: {e}")
            import traceback; traceback.print_exc()
            failed += 1

    print("=" * 60)
    if failed:
        print(f"❌ {failed}/{len(tests)} tests failed")
        sys.exit(1)
    else:
        print(f"✅ All {len(tests)} tests passed")


if __name__ == "__main__":
    asyncio.run(main())