"""
Unit tests for OrchestratorActor (orchestrator/runner.py).

Tests cover:
- GoalState: all_collected, pending_count properties
- OrchestratorActor lifecycle: decompose → dispatch → collect → synthesize
- Timeout and partial collection behaviour
- Checkpoint triggering
- Error handling (parse errors, decomposition failure)
- Concurrent goal processing (max_concurrent_goals)

All tests use InMemoryBus -- no NATS or external infrastructure required.
"""

from __future__ import annotations

import asyncio
import json
import os
import tempfile
import time
from typing import Any

import pytest
import yaml

from loom.bus.memory import InMemoryBus
from loom.core.messages import (
    OrchestratorGoal,
    TaskMessage,
    TaskResult,
    TaskStatus,
)
from loom.orchestrator.runner import GoalState, OrchestratorActor
from loom.orchestrator.store import InMemoryCheckpointStore
from loom.worker.backends import LLMBackend

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


class MockOrchestratorBackend(LLMBackend):
    """Returns configurable responses depending on system prompt content."""

    def __init__(self, decompose_response: str, synthesis_response: str = "{}"):
        self._decompose = decompose_response
        self._synthesis = synthesis_response

    async def complete(self, system_prompt, user_message, max_tokens, temperature, **kw):
        # Route to decomposition or synthesis based on system prompt
        if "task decomposition" in system_prompt.lower():
            content = self._decompose
        else:
            content = self._synthesis
        return {
            "content": content,
            "model": "mock-orch",
            "prompt_tokens": 100,
            "completion_tokens": 50,
        }


def _write_config(
    available_workers: list[dict] | None = None,
    timeout_seconds: int = 5,
    max_concurrent_goals: int | None = None,
) -> str:
    """Write a minimal orchestrator config to a temp file."""
    config = {
        "name": "test-orchestrator",
        "timeout_seconds": timeout_seconds,
        "max_concurrent_tasks": 5,
        "available_workers": available_workers
        or [
            {
                "name": "summarizer",
                "description": "Summarizes text",
                "input_schema": {"type": "object", "required": ["text"]},
                "default_model_tier": "local",
            },
        ],
    }
    if max_concurrent_goals is not None:
        config["max_concurrent_goals"] = max_concurrent_goals
    fd, path = tempfile.mkstemp(suffix=".yaml")
    with os.fdopen(fd, "w") as f:
        yaml.dump(config, f)
    return path


def _make_goal_data(instruction: str = "Summarize this document") -> dict[str, Any]:
    goal = OrchestratorGoal(instruction=instruction, context={"text": "Hello world"})
    return goal.model_dump(mode="json")


def _make_result_data(
    task_id: str,
    worker_type: str = "summarizer",
    status: TaskStatus = TaskStatus.COMPLETED,
    output: dict | None = None,
) -> dict[str, Any]:
    result = TaskResult(
        task_id=task_id,
        worker_type=worker_type,
        status=status,
        output=output or {"summary": "Test summary"},
        model_used="mock",
        processing_time_ms=100,
        token_usage={"prompt_tokens": 50, "completion_tokens": 30},
    )
    return result.model_dump(mode="json")


# ---------------------------------------------------------------------------
# GoalState tests
# ---------------------------------------------------------------------------


class TestGoalState:
    def test_all_collected_false_when_empty(self):
        goal = OrchestratorGoal(instruction="test")
        state = GoalState(goal=goal)
        assert state.all_collected is False

    def test_all_collected_false_when_partial(self):
        goal = OrchestratorGoal(instruction="test")
        state = GoalState(goal=goal)
        task = TaskMessage(worker_type="summarizer", payload={})
        state.dispatched_tasks[task.task_id] = task
        assert state.all_collected is False

    def test_all_collected_true_when_complete(self):
        goal = OrchestratorGoal(instruction="test")
        state = GoalState(goal=goal)

        task = TaskMessage(worker_type="summarizer", payload={})
        state.dispatched_tasks[task.task_id] = task

        result = TaskResult(
            task_id=task.task_id,
            worker_type="summarizer",
            status=TaskStatus.COMPLETED,
            output={"data": "test"},
        )
        state.collected_results[task.task_id] = result

        assert state.all_collected is True

    def test_pending_count(self):
        goal = OrchestratorGoal(instruction="test")
        state = GoalState(goal=goal)

        for i in range(3):
            task = TaskMessage(worker_type="summarizer", payload={})
            state.dispatched_tasks[task.task_id] = task

        assert state.pending_count == 3

        # Collect one result
        first_id = next(iter(state.dispatched_tasks.keys()))
        state.collected_results[first_id] = TaskResult(
            task_id=first_id,
            worker_type="summarizer",
            status=TaskStatus.COMPLETED,
            output={},
        )
        assert state.pending_count == 2

    def test_start_time_is_set(self):
        goal = OrchestratorGoal(instruction="test")
        before = time.monotonic()
        state = GoalState(goal=goal)
        after = time.monotonic()
        assert before <= state.start_time <= after


# ---------------------------------------------------------------------------
# OrchestratorActor handle_message tests
# ---------------------------------------------------------------------------


class TestHandleMessage:
    @pytest.mark.asyncio
    async def test_invalid_goal_data_does_not_crash(self):
        """Malformed goal data is handled gracefully."""
        config_path = _write_config()
        try:
            backend = MockOrchestratorBackend("[]")
            bus = InMemoryBus()
            actor = OrchestratorActor(
                actor_id="test-orch",
                config_path=config_path,
                backend=backend,
                nats_url="nats://localhost:4222",
            )
            actor._bus = bus
            await bus.connect()

            # Pass garbage -- should not raise
            await actor.handle_message({"invalid": True})
        finally:
            os.unlink(config_path)

    @pytest.mark.asyncio
    async def test_empty_decomposition_publishes_failure(self):
        """When decomposition produces no subtasks, a FAILED result is published."""
        config_path = _write_config()
        try:
            backend = MockOrchestratorBackend("[]")  # Empty plan
            bus = InMemoryBus()
            actor = OrchestratorActor(
                actor_id="test-orch",
                config_path=config_path,
                backend=backend,
            )
            actor._bus = bus
            await bus.connect()

            goal_data = _make_goal_data()
            goal_id = goal_data["goal_id"]
            sub = await bus.subscribe(f"loom.results.{goal_id}")

            await actor.handle_message(goal_data)

            result = await sub.__anext__()
            assert result["status"] == TaskStatus.FAILED.value
            assert "no subtasks" in result["error"].lower()
        finally:
            os.unlink(config_path)

    @pytest.mark.asyncio
    async def test_goal_state_cleaned_up_after_completion(self):
        """Goal state is removed from _active_goals after processing."""
        config_path = _write_config(timeout_seconds=1)
        try:
            # Return a valid subtask plan
            plan = json.dumps(
                [
                    {
                        "worker_type": "summarizer",
                        "payload": {"text": "test"},
                    }
                ]
            )
            backend = MockOrchestratorBackend(plan)
            bus = InMemoryBus()
            actor = OrchestratorActor(
                actor_id="test-orch",
                config_path=config_path,
                backend=backend,
            )
            actor._bus = bus
            await bus.connect()

            goal_data = _make_goal_data()
            # Will timeout on collection since no worker responds, but state
            # should be cleaned up regardless
            await actor.handle_message(goal_data)

            assert len(actor._active_goals) == 0
        finally:
            os.unlink(config_path)


# ---------------------------------------------------------------------------
# _record_in_history tests
# ---------------------------------------------------------------------------


class TestRecordInHistory:
    @pytest.mark.asyncio
    async def test_history_accumulates(self):
        config_path = _write_config()
        try:
            backend = MockOrchestratorBackend("[]")
            actor = OrchestratorActor(
                actor_id="test-orch",
                config_path=config_path,
                backend=backend,
            )

            goal = OrchestratorGoal(instruction="Test goal")
            goal_state = GoalState(goal=goal)
            results = [
                TaskResult(
                    task_id="t1",
                    worker_type="summarizer",
                    status=TaskStatus.COMPLETED,
                    output={"summary": "done"},
                    processing_time_ms=100,
                ),
            ]
            synthesis = {"confidence": "high"}

            await actor._record_in_history(goal_state, results, synthesis)
            assert len(goal_state.conversation_history) == 1

            entry = goal_state.conversation_history[0]
            assert entry["goal_id"] == goal.goal_id
            assert entry["subtask_count"] == 1
            assert entry["synthesis_confidence"] == "high"
        finally:
            os.unlink(config_path)

    @pytest.mark.asyncio
    async def test_history_records_failures(self):
        config_path = _write_config()
        try:
            backend = MockOrchestratorBackend("[]")
            actor = OrchestratorActor(
                actor_id="test-orch",
                config_path=config_path,
                backend=backend,
            )

            goal = OrchestratorGoal(instruction="Test")
            goal_state = GoalState(goal=goal)
            results = [
                TaskResult(
                    task_id="t1",
                    worker_type="summarizer",
                    status=TaskStatus.FAILED,
                    error="timeout",
                    processing_time_ms=0,
                ),
            ]
            await actor._record_in_history(goal_state, results, {})

            entry = goal_state.conversation_history[0]
            assert entry["results"][0]["error"] == "timeout"
        finally:
            os.unlink(config_path)


# ---------------------------------------------------------------------------
# Concurrent goal processing tests
# ---------------------------------------------------------------------------


class TestConcurrentGoals:
    @pytest.mark.asyncio
    async def test_default_max_concurrent_goals_is_one(self):
        """Without config, max_concurrent_goals defaults to 1."""
        config_path = _write_config()
        try:
            backend = MockOrchestratorBackend("[]")
            actor = OrchestratorActor(
                actor_id="test-orch",
                config_path=config_path,
                backend=backend,
            )
            assert actor.max_concurrent == 1
        finally:
            os.unlink(config_path)

    @pytest.mark.asyncio
    async def test_max_concurrent_goals_from_config(self):
        """Config value is passed through to BaseActor.max_concurrent."""
        config_path = _write_config(max_concurrent_goals=4)
        try:
            backend = MockOrchestratorBackend("[]")
            actor = OrchestratorActor(
                actor_id="test-orch",
                config_path=config_path,
                backend=backend,
            )
            assert actor.max_concurrent == 4
        finally:
            os.unlink(config_path)

    @pytest.mark.asyncio
    async def test_concurrent_goals_have_isolated_history(self):
        """Each GoalState maintains its own conversation history."""
        config_path = _write_config(max_concurrent_goals=4)
        try:
            backend = MockOrchestratorBackend("[]")
            actor = OrchestratorActor(
                actor_id="test-orch",
                config_path=config_path,
                backend=backend,
            )

            goal_states = []

            async def record_one(i: int):
                goal = OrchestratorGoal(instruction=f"Goal {i}")
                gs = GoalState(goal=goal)
                goal_states.append(gs)
                results = [
                    TaskResult(
                        task_id=f"t{i}",
                        worker_type="summarizer",
                        status=TaskStatus.COMPLETED,
                        output={"n": i},
                        processing_time_ms=10,
                    ),
                ]
                await actor._record_in_history(gs, results, {"confidence": "high"})

            # Fire 20 concurrent writes — each to its own GoalState
            await asyncio.gather(*(record_one(i) for i in range(20)))

            assert len(goal_states) == 20
            # Each GoalState should have exactly 1 entry (no cross-contamination)
            for gs in goal_states:
                assert len(gs.conversation_history) == 1
        finally:
            os.unlink(config_path)

    @pytest.mark.asyncio
    async def test_bus_injection_via_constructor(self):
        """The bus= keyword argument is forwarded to BaseActor."""
        config_path = _write_config()
        try:
            backend = MockOrchestratorBackend("[]")
            bus = InMemoryBus()
            actor = OrchestratorActor(
                actor_id="test-orch",
                config_path=config_path,
                backend=backend,
                bus=bus,
            )
            assert actor._bus is bus
        finally:
            os.unlink(config_path)


# ---------------------------------------------------------------------------
# Per-goal state isolation tests
# ---------------------------------------------------------------------------


class TestGoalIsolation:
    def test_goalstate_conversation_history_defaults_empty(self):
        """New GoalState has empty conversation_history."""
        goal = OrchestratorGoal(instruction="test")
        state = GoalState(goal=goal)
        assert state.conversation_history == []
        assert state.checkpoint_counter == 0

    def test_goalstate_history_not_shared(self):
        """Two GoalState instances do not share the same history list."""
        goal_a = OrchestratorGoal(instruction="A")
        goal_b = OrchestratorGoal(instruction="B")
        state_a = GoalState(goal=goal_a)
        state_b = GoalState(goal=goal_b)

        state_a.conversation_history.append({"goal_id": "a"})
        assert len(state_b.conversation_history) == 0

    def test_checkpoint_counter_per_goal(self):
        """Checkpoint counters are independent across GoalState instances."""
        goal_a = OrchestratorGoal(instruction="A")
        goal_b = OrchestratorGoal(instruction="B")
        state_a = GoalState(goal=goal_a)
        state_b = GoalState(goal=goal_b)

        state_a.checkpoint_counter += 1
        state_a.checkpoint_counter += 1
        state_b.checkpoint_counter += 1

        assert state_a.checkpoint_counter == 2
        assert state_b.checkpoint_counter == 1

    @pytest.mark.asyncio
    async def test_record_in_history_writes_to_goal_state(self):
        """_record_in_history writes to the GoalState, not the actor."""
        config_path = _write_config()
        try:
            backend = MockOrchestratorBackend("[]")
            actor = OrchestratorActor(
                actor_id="test-orch",
                config_path=config_path,
                backend=backend,
            )

            goal = OrchestratorGoal(instruction="Test")
            gs = GoalState(goal=goal)
            results = [
                TaskResult(
                    task_id="t1",
                    worker_type="summarizer",
                    status=TaskStatus.COMPLETED,
                    output={"data": "x"},
                    processing_time_ms=10,
                ),
            ]
            await actor._record_in_history(gs, results, {"confidence": "high"})

            assert len(gs.conversation_history) == 1
            assert gs.conversation_history[0]["goal_id"] == goal.goal_id
            # Actor should have no shared history attribute
            assert not hasattr(actor, "_conversation_history")
        finally:
            os.unlink(config_path)


# ---------------------------------------------------------------------------
# Full lifecycle tests
# ---------------------------------------------------------------------------


class TestFullLifecycle:
    @pytest.mark.asyncio
    async def test_full_goal_lifecycle_with_simulated_worker(self):
        """Full decompose → dispatch → collect → synthesize → publish cycle."""
        # Backend that returns a valid single-task plan
        plan = json.dumps(
            [
                {
                    "worker_type": "summarizer",
                    "payload": {"text": "test content"},
                }
            ]
        )
        synthesis = json.dumps(
            {
                "summary": "synthesized result",
                "confidence": "high",
            }
        )
        backend = MockOrchestratorBackend(plan, synthesis)

        config_path = _write_config(timeout_seconds=5)
        try:
            bus = InMemoryBus()
            await bus.connect()

            actor = OrchestratorActor(
                actor_id="test-full",
                config_path=config_path,
                backend=backend,
                bus=bus,
            )

            goal_data = _make_goal_data("Summarize the document")
            goal_id = goal_data["goal_id"]

            # Subscribe for the final result
            result_sub = await bus.subscribe(f"loom.results.{goal_id}")

            # Pre-subscribe the worker BEFORE handle_message starts.
            # This mirrors real deployments where workers are already running.
            worker_sub = await bus.subscribe("loom.tasks.incoming")
            ready = asyncio.Event()

            async def worker_simulator():
                ready.set()
                async for data in worker_sub:
                    task = TaskMessage(**data)
                    # Small delay to let orchestrator set up result subscription
                    await asyncio.sleep(0.05)
                    result = TaskResult(
                        task_id=task.task_id,
                        parent_task_id=task.parent_task_id,
                        worker_type=task.worker_type,
                        status=TaskStatus.COMPLETED,
                        output={"summary": "Worker produced this"},
                        model_used="mock",
                        processing_time_ms=50,
                        token_usage={"prompt_tokens": 10, "completion_tokens": 5},
                    )
                    await bus.publish(
                        f"loom.results.{task.parent_task_id}",
                        result.model_dump(mode="json"),
                    )
                    await worker_sub.unsubscribe()
                    break

            worker_task = asyncio.create_task(worker_simulator())
            await ready.wait()
            await actor.handle_message(goal_data)
            await worker_task

            # Verify final result was published. The result subject receives
            # both worker intermediate results AND the final orchestrator result.
            # The final result has task_id == goal_id.
            final = None
            for _ in range(5):
                msg = await asyncio.wait_for(result_sub.__anext__(), timeout=2.0)
                if msg["task_id"] == goal_id:
                    final = msg
                    break

            assert final is not None, "Final orchestrator result not found"
            assert final["status"] == TaskStatus.COMPLETED.value
            assert final["output"] is not None
        finally:
            os.unlink(config_path)

    @pytest.mark.asyncio
    async def test_subtask_limit_truncates_plan(self):
        """When decomposition returns more subtasks than max_concurrent_tasks, truncate."""
        # Return 10 subtasks
        plan = json.dumps(
            [{"worker_type": "summarizer", "payload": {"text": f"chunk {i}"}} for i in range(10)]
        )
        backend = MockOrchestratorBackend(plan)

        config_path = _write_config(timeout_seconds=1)
        try:
            bus = InMemoryBus()
            await bus.connect()
            actor = OrchestratorActor(
                actor_id="test-limit",
                config_path=config_path,
                backend=backend,
                bus=bus,
            )
            # max_concurrent_tasks defaults to 5 from _write_config

            # Subscribe to tasks to count how many were dispatched
            task_sub = await bus.subscribe("loom.tasks.incoming")

            goal_data = _make_goal_data("Process many chunks")
            await actor.handle_message(goal_data)

            # Collect dispatched tasks (should be capped at 5)
            dispatched = []
            for _ in range(5):
                try:
                    msg = await asyncio.wait_for(task_sub.__anext__(), timeout=0.5)
                    dispatched.append(msg)
                except (TimeoutError, StopAsyncIteration):
                    break
            assert len(dispatched) == 5
        finally:
            os.unlink(config_path)

    @pytest.mark.asyncio
    async def test_decomposition_error_publishes_failure(self):
        """When decomposition fails, a FAILED result is published."""
        # Backend that raises RuntimeError — the decomposer catches this
        # and re-raises as ValueError/RuntimeError which the orchestrator
        # catches and publishes as FAILED.

        class FailingBackend(LLMBackend):
            async def complete(self, system_prompt, user_message, max_tokens, temperature, **kw):
                raise RuntimeError("LLM unavailable")

        config_path = _write_config(timeout_seconds=1)
        try:
            bus = InMemoryBus()
            await bus.connect()
            actor = OrchestratorActor(
                actor_id="test-decomp-fail",
                config_path=config_path,
                backend=FailingBackend(),
                bus=bus,
            )

            goal_data = _make_goal_data("This will fail")
            goal_id = goal_data["goal_id"]
            result_sub = await bus.subscribe(f"loom.results.{goal_id}")

            await actor.handle_message(goal_data)

            result = await asyncio.wait_for(result_sub.__anext__(), timeout=2.0)
            assert result["status"] == TaskStatus.FAILED.value
            # Either an error message about decomposition/orchestrator failure,
            # or "no subtasks" if decomposer caught the error and returned empty
            assert result["error"] is not None
        finally:
            os.unlink(config_path)

    @pytest.mark.asyncio
    async def test_collection_timeout_returns_partial_results(self):
        """When timeout expires before all results arrive, partial results are synthesized."""
        # Plan with 3 subtasks, only 1 will respond
        plan = json.dumps(
            [{"worker_type": "summarizer", "payload": {"text": f"chunk {i}"}} for i in range(3)]
        )
        synthesis = json.dumps({"partial": True, "confidence": "low"})
        backend = MockOrchestratorBackend(plan, synthesis)

        config_path = _write_config(timeout_seconds=1)
        try:
            bus = InMemoryBus()
            await bus.connect()
            actor = OrchestratorActor(
                actor_id="test-timeout",
                config_path=config_path,
                backend=backend,
                bus=bus,
            )

            goal_data = _make_goal_data("Partial timeout test")
            goal_id = goal_data["goal_id"]
            result_sub = await bus.subscribe(f"loom.results.{goal_id}")

            # Pre-subscribe the worker before handle_message
            worker_sub = await bus.subscribe("loom.tasks.incoming")

            # Worker only responds to first task
            async def partial_worker():
                data = await worker_sub.__anext__()
                task = TaskMessage(**data)
                # Small delay to let orchestrator set up result subscription
                await asyncio.sleep(0.05)
                result = TaskResult(
                    task_id=task.task_id,
                    parent_task_id=task.parent_task_id,
                    worker_type=task.worker_type,
                    status=TaskStatus.COMPLETED,
                    output={"summary": "partial"},
                    processing_time_ms=10,
                )
                await bus.publish(
                    f"loom.results.{task.parent_task_id}",
                    result.model_dump(mode="json"),
                )
                # Don't respond to remaining tasks — let timeout fire
                await worker_sub.unsubscribe()

            worker_task = asyncio.create_task(partial_worker())
            await actor.handle_message(goal_data)
            await worker_task

            # Should still get a final result (synthesized from partial)
            final = await asyncio.wait_for(result_sub.__anext__(), timeout=3.0)
            assert final["status"] == TaskStatus.COMPLETED.value
        finally:
            os.unlink(config_path)


# ---------------------------------------------------------------------------
# Checkpoint tests
# ---------------------------------------------------------------------------


class TestCheckpointing:
    @pytest.mark.asyncio
    async def test_checkpoint_triggered_when_threshold_exceeded(self):
        """Checkpoint is created when conversation history exceeds token threshold."""
        config_path = _write_config()
        try:
            backend = MockOrchestratorBackend("[]")
            store = InMemoryCheckpointStore()
            actor = OrchestratorActor(
                actor_id="test-ckpt",
                config_path=config_path,
                backend=backend,
                checkpoint_store=store,
            )

            # Configure a very low threshold so checkpoint triggers
            actor._checkpoint_manager.token_threshold = 10

            goal = OrchestratorGoal(instruction="Test checkpoint")
            goal_state = GoalState(goal=goal)

            # Add enough history to trigger
            for i in range(5):
                results = [
                    TaskResult(
                        task_id=f"t{i}",
                        worker_type="summarizer",
                        status=TaskStatus.COMPLETED,
                        output={"data": "x" * 100},
                        processing_time_ms=10,
                    ),
                ]
                await actor._record_in_history(goal_state, results, {"confidence": "high"})

            import structlog

            log = structlog.get_logger().bind(goal_id=goal.goal_id)
            await actor._maybe_checkpoint(goal_state, log)

            # Checkpoint should have been created
            assert goal_state.checkpoint_counter == 1
            # History should be trimmed to recent window
            assert (
                len(goal_state.conversation_history) <= actor._checkpoint_manager.recent_window_size
            )
        finally:
            os.unlink(config_path)

    @pytest.mark.asyncio
    async def test_checkpoint_failure_is_non_fatal(self):
        """If checkpoint store raises, orchestrator continues without crashing."""

        class FailingStore(InMemoryCheckpointStore):
            async def save(self, checkpoint):
                raise RuntimeError("Store unavailable")

        config_path = _write_config()
        try:
            backend = MockOrchestratorBackend("[]")
            actor = OrchestratorActor(
                actor_id="test-ckpt-fail",
                config_path=config_path,
                backend=backend,
                checkpoint_store=FailingStore(),
            )
            actor._checkpoint_manager.token_threshold = 10

            goal = OrchestratorGoal(instruction="Test")
            goal_state = GoalState(goal=goal)

            for i in range(5):
                results = [
                    TaskResult(
                        task_id=f"t{i}",
                        worker_type="summarizer",
                        status=TaskStatus.COMPLETED,
                        output={"data": "x" * 100},
                        processing_time_ms=10,
                    ),
                ]
                await actor._record_in_history(goal_state, results, {})

            import structlog

            log = structlog.get_logger().bind(goal_id=goal.goal_id)

            # Should not raise
            await actor._maybe_checkpoint(goal_state, log)
        finally:
            os.unlink(config_path)

    @pytest.mark.asyncio
    async def test_no_checkpoint_when_no_store(self):
        """_maybe_checkpoint is a no-op when no checkpoint store is configured."""
        config_path = _write_config()
        try:
            backend = MockOrchestratorBackend("[]")
            actor = OrchestratorActor(
                actor_id="test-no-ckpt",
                config_path=config_path,
                backend=backend,
                checkpoint_store=None,
            )

            goal = OrchestratorGoal(instruction="Test")
            goal_state = GoalState(goal=goal)

            import structlog

            log = structlog.get_logger().bind(goal_id=goal.goal_id)
            await actor._maybe_checkpoint(goal_state, log)
            assert goal_state.checkpoint_counter == 0
        finally:
            os.unlink(config_path)


# ---------------------------------------------------------------------------
# on_reload tests (lines 278-281)
# ---------------------------------------------------------------------------


class TestOnReload:
    @pytest.mark.asyncio
    async def test_on_reload_updates_timeout_and_concurrency(self, tmp_path):
        """Lines 278-281: on_reload re-reads config from disk."""
        import yaml

        config_path = _write_config(timeout_seconds=10)
        try:
            backend = MockOrchestratorBackend("[]")
            actor = OrchestratorActor(
                actor_id="test-reload",
                config_path=config_path,
                backend=backend,
            )

            assert actor._task_timeout == 10.0

            # Write an updated config with a different timeout.
            updated = {
                "name": "test-orchestrator",
                "timeout_seconds": 42,
                "max_concurrent_tasks": 3,
                "available_workers": [],
            }
            with open(config_path, "w") as f:
                yaml.dump(updated, f)

            await actor.on_reload()

            assert actor._task_timeout == 42.0
            assert actor._max_concurrent_tasks == 3
        finally:
            os.unlink(config_path)

    @pytest.mark.asyncio
    async def test_on_reload_uses_defaults_when_keys_absent(self, tmp_path):
        """on_reload falls back to defaults (300s, 5 tasks) when keys are missing."""
        import yaml

        config_path = _write_config(timeout_seconds=5)
        try:
            backend = MockOrchestratorBackend("[]")
            actor = OrchestratorActor(
                actor_id="test-reload-defaults",
                config_path=config_path,
                backend=backend,
            )

            # Write a config without timeout_seconds or max_concurrent_tasks.
            minimal = {"name": "test-orchestrator", "available_workers": []}
            with open(config_path, "w") as f:
                yaml.dump(minimal, f)

            await actor.on_reload()

            assert actor._task_timeout == 300.0
            assert actor._max_concurrent_tasks == 5
        finally:
            os.unlink(config_path)


# ---------------------------------------------------------------------------
# handle_message exception handler (lines 404-407)
# ---------------------------------------------------------------------------


class TestHandleMessageExceptionPath:
    @pytest.mark.asyncio
    async def test_unexpected_exception_during_dispatch_publishes_failure(self):
        """Lines 404-407: unexpected exception during _dispatch_subtasks publishes FAILED."""
        config_path = _write_config(timeout_seconds=5)
        try:
            plan = json.dumps([{"worker_type": "summarizer", "payload": {"text": "test"}}])
            backend = MockOrchestratorBackend(plan)
            bus = InMemoryBus()
            await bus.connect()
            actor = OrchestratorActor(
                actor_id="test-exc",
                config_path=config_path,
                backend=backend,
                bus=bus,
            )

            goal_data = _make_goal_data("Exception test")
            goal_id = goal_data["goal_id"]
            result_sub = await bus.subscribe(f"loom.results.{goal_id}")

            # Patch _dispatch_subtasks to raise an unexpected exception.
            async def _boom(*args, **kwargs):
                raise RuntimeError("unexpected crash")

            actor._dispatch_subtasks = _boom

            await actor.handle_message(goal_data)

            result = await asyncio.wait_for(result_sub.__anext__(), timeout=2.0)
            assert result["status"] == TaskStatus.FAILED.value
            assert "Orchestrator error" in result["error"]
            assert "unexpected crash" in result["error"]
            # Goal state must be cleaned up.
            assert goal_id not in actor._active_goals
        finally:
            os.unlink(config_path)


# ---------------------------------------------------------------------------
# _collect_results early_exit path (lines 567-568)
# ---------------------------------------------------------------------------


class TestCollectResultsEarlyExit:
    @pytest.mark.asyncio
    async def test_early_exit_branch_logged_when_callback_signals_stop(self):
        """Lines 567-568: early_exited branch triggers log message."""
        config_path = _write_config(timeout_seconds=5)
        try:
            bus = InMemoryBus()
            await bus.connect()
            backend = MockOrchestratorBackend("[]")
            actor = OrchestratorActor(
                actor_id="test-early-exit",
                config_path=config_path,
                backend=backend,
                bus=bus,
            )

            goal = OrchestratorGoal(instruction="Early exit test")
            task = TaskMessage(worker_type="summarizer", payload={"text": "hi"})
            goal_state = GoalState(goal=goal)
            goal_state.dispatched_tasks[task.task_id] = task

            # Callback that signals early exit after first result.
            async def stop_after_first(result, collected, expected):
                return True  # signal stop

            import structlog

            log = structlog.get_logger().bind(goal_id=goal.goal_id)

            # Publish a result so the stream can collect it and trigger early exit.
            async def publisher():
                await asyncio.sleep(0.05)
                result = TaskResult(
                    task_id=task.task_id,
                    parent_task_id=goal.goal_id,
                    worker_type="summarizer",
                    status=TaskStatus.COMPLETED,
                    output={"summary": "done"},
                )
                await bus.publish(
                    f"loom.results.{goal.goal_id}",
                    result.model_dump(mode="json"),
                )

            pub_task = asyncio.create_task(publisher())

            results = await actor._collect_results(goal_state, log, on_result=stop_after_first)

            await pub_task

            # We should have gotten exactly 1 result (the one we published).
            assert len(results) == 1
            assert results[0].task_id == task.task_id
        finally:
            os.unlink(config_path)
