"""Tests for Runtime."""
import pytest

from logos_adk.runtime import Runtime, Supervisor, _default_runtime_holder
from logos_adk.scaling import (
    DistributedRuntime,
    LoadBalancer,
    PostgreSQLTaskQueue,
    SQLiteTaskQueue,
)
from logos_adk.types import Response, StreamDone, TextChunk, Usage


@pytest.fixture(autouse=True)
def _reset_default_runtime():
    """Reset default runtime between tests."""
    _default_runtime_holder.reset()
    yield
    _default_runtime_holder.reset()


def test_runtime_creation():
    rt = Runtime()
    assert rt is not None


def test_first_runtime_becomes_default():
    rt = Runtime()
    assert _default_runtime_holder.get() is rt


def test_second_runtime_is_isolated():
    rt1 = Runtime()
    rt2 = Runtime()
    assert _default_runtime_holder.get() is rt1
    assert rt2 is not rt1


def test_set_as_default():
    rt1 = Runtime()
    rt2 = Runtime()
    rt2.set_as_default()
    assert _default_runtime_holder.get() is rt2
    assert rt1 is not rt2


def test_get_or_create_default():
    rt = _default_runtime_holder.get_or_create()
    assert rt is not None
    assert _default_runtime_holder.get() is rt


def test_get_or_create_returns_existing():
    rt = Runtime()
    rt2 = _default_runtime_holder.get_or_create()
    assert rt2 is rt


def test_supervisor_defaults():
    supervisor = Supervisor()
    assert supervisor.on_failure == "restart"
    assert supervisor.max_retries == 3
    assert supervisor.on_max_retries == "escalate"


def test_supervisor_custom():
    supervisor = Supervisor(on_failure="skip", max_retries=5, on_max_retries="escalate")
    assert supervisor.on_failure == "skip"
    assert supervisor.max_retries == 5


def test_runtime_with_supervisor():
    supervisor = Supervisor(on_failure="restart", max_retries=2)
    runtime = Runtime(supervisor=supervisor)
    assert runtime.supervisor is supervisor


def test_runtime_spawn_registers_agent():
    from logos_adk.agent import Agent

    runtime = Runtime()
    agent = Agent("test", model="mock")
    runtime.spawn(agent)
    assert agent._runtime is runtime
    assert runtime.get_agent("test") is agent


def test_runtime_get_agent_missing_returns_none():
    runtime = Runtime()

    assert runtime.get_agent("missing") is None


def test_runtime_spawn_duplicate_name_raises():
    from logos_adk.agent import Agent

    runtime = Runtime()
    agent_one = Agent("bot", model="mock")
    agent_two = Agent("bot", model="mock")
    runtime.spawn(agent_one)

    with pytest.raises(ValueError, match="already registered"):
        runtime.spawn(agent_two)


def test_runtime_stop_agent():
    from logos_adk.agent import Agent

    runtime = Runtime()
    agent = Agent("test", model="mock")
    runtime.spawn(agent)
    runtime.stop(agent)
    assert runtime.get_agent("test") is None
    assert agent._runtime is None


@pytest.mark.asyncio
async def test_distributed_runtime_round_robins_replicas(tmp_path):
    from logos_adk.agent import Agent
    from logos_adk.providers.mock import MockProvider

    runtime = DistributedRuntime(
        load_balancer=LoadBalancer("round_robin"),
        task_queue=SQLiteTaskQueue(path=str(tmp_path / "tasks.db")),
    )
    runtime.register("bot", Agent("bot-a", model=MockProvider(responses=["a"])), replica_id="a")
    runtime.register("bot", Agent("bot-b", model=MockProvider(responses=["b"])), replica_id="b")

    first = await runtime.dispatch("bot", "hello")
    second = await runtime.dispatch("bot", "hello")

    assert first.content == "a"
    assert second.content == "b"


@pytest.mark.asyncio
async def test_distributed_runtime_sticky_sessions_pin_replica(tmp_path):
    from logos_adk.agent import Agent
    from logos_adk.providers.mock import MockProvider

    runtime = DistributedRuntime(
        load_balancer=LoadBalancer("round_robin"),
        task_queue=SQLiteTaskQueue(path=str(tmp_path / "tasks.db")),
        session_affinity=True,
    )
    runtime.register("bot", Agent("bot-a", model=MockProvider(responses=["a1", "a2"])), replica_id="a")
    runtime.register("bot", Agent("bot-b", model=MockProvider(responses=["b1", "b2"])), replica_id="b")

    first = await runtime.dispatch("bot", "hello", session_id="s1")
    second = await runtime.dispatch("bot", "hello-again", session_id="s1")
    other = await runtime.dispatch("bot", "hello", session_id="s2")

    assert [first.content, second.content] == ["a1", "a2"]
    assert other.content == "b1"
    assert runtime.stats()["session_affinity"]["enabled"] is True
    assert runtime.stats()["session_affinity"]["sessions"] == 2


@pytest.mark.asyncio
async def test_distributed_runtime_persists_queue_stats(tmp_path):
    from logos_adk.agent import Agent
    from logos_adk.providers.mock import MockProvider

    task_queue = SQLiteTaskQueue(path=str(tmp_path / "tasks.db"))
    runtime = DistributedRuntime(task_queue=task_queue)
    runtime.register("bot", Agent("bot", model=MockProvider(responses=["ok"])))

    response = await runtime.dispatch("bot", "hello")

    assert response.content == "ok"
    assert task_queue.stats()["completed"] == 1


@pytest.mark.asyncio
async def test_distributed_runtime_dispatch_unknown_agent_raises_lookup_error(tmp_path):
    runtime = DistributedRuntime(task_queue=SQLiteTaskQueue(path=str(tmp_path / "tasks.db")))

    with pytest.raises(LookupError, match="No replicas registered for agent 'missing'"):
        await runtime.dispatch("missing", "hello")


@pytest.mark.asyncio
async def test_distributed_runtime_stream_yields_chunks_and_clears_inflight(tmp_path):
    class StreamingReplica:
        def __init__(self, name: str):
            self.name = name
            self._runtime = None

        async def stream(self, prompt: str, *, session_id: str | None = None, **kwargs):
            yield TextChunk(text=f"{self.name}:{prompt}:{session_id}")
            yield StreamDone(
                response=Response(
                    content="done",
                    tool_calls=[],
                    usage=Usage(input_tokens=1, output_tokens=1),
                    raw={},
                )
            )

    runtime = DistributedRuntime(task_queue=SQLiteTaskQueue(path=str(tmp_path / "tasks.db")))
    runtime.register("bot", StreamingReplica("replica-a"), replica_id="a")

    chunks = [chunk async for chunk in runtime.stream("bot", "hello", session_id="s1")]

    assert [chunk.text for chunk in chunks if isinstance(chunk, TextChunk)] == [
        "replica-a:hello:s1"
    ]
    assert isinstance(chunks[-1], StreamDone)
    assert runtime.stats()["inflight"].get("a", 0) == 0


@pytest.mark.asyncio
async def test_distributed_runtime_stream_clears_inflight_after_generator_error(tmp_path):
    class BrokenStreamingReplica:
        def __init__(self):
            self.name = "broken"
            self._runtime = None

        async def stream(self, prompt: str, *, session_id: str | None = None, **kwargs):
            yield TextChunk(text="partial")
            raise RuntimeError("stream broke")

    runtime = DistributedRuntime(task_queue=SQLiteTaskQueue(path=str(tmp_path / "tasks.db")))
    runtime.register("bot", BrokenStreamingReplica(), replica_id="a")

    with pytest.raises(RuntimeError, match="stream broke"):
        async for _ in runtime.stream("bot", "hello", session_id="s1"):
            pass

    assert runtime.stats()["inflight"].get("a", 0) == 0


@pytest.mark.asyncio
async def test_distributed_runtime_stream_ui_uses_selected_replica_and_affinity(tmp_path):
    class UIStreamingReplica:
        def __init__(self, name: str):
            self.name = name
            self._runtime = None

        async def stream_ui_events(
            self,
            prompt: str,
            *,
            session_id: str | None = None,
            correlation_id: str | None = None,
            **kwargs,
        ):
            yield {
                "type": "graph_started",
                "event": "graph_started",
                "replica": self.name,
                "prompt": prompt,
                "session_id": session_id,
                "correlation_id": correlation_id,
            }
            yield {
                "type": "done",
                "event": "done",
                "replica": self.name,
                "status": "completed",
                "terminal": True,
            }

    runtime = DistributedRuntime(
        load_balancer=LoadBalancer("round_robin"),
        task_queue=SQLiteTaskQueue(path=str(tmp_path / "tasks.db")),
        session_affinity=True,
    )
    runtime.register("bot", UIStreamingReplica("replica-a"), replica_id="a")
    runtime.register("bot", UIStreamingReplica("replica-b"), replica_id="b")

    first = [
        event
        async for event in runtime.stream_ui(
            "bot",
            "hello",
            session_id="sticky-session",
            correlation_id="corr-1",
        )
    ]
    second = [
        event
        async for event in runtime.stream_ui(
            "bot",
            "again",
            session_id="sticky-session",
            correlation_id="corr-2",
        )
    ]

    assert first[0]["replica"] == second[0]["replica"]
    assert first[0]["prompt"] == "hello"
    assert first[0]["correlation_id"] == "corr-1"
    assert second[0]["correlation_id"] == "corr-2"
    assert runtime.stats()["session_affinity"]["sessions"] == 1


@pytest.mark.asyncio
async def test_distributed_runtime_stream_ui_missing_capability_raises_attribute_error(tmp_path):
    class PlainReplica:
        def __init__(self):
            self.name = "plain"
            self._runtime = None

        async def stream(self, prompt: str, *, session_id: str | None = None, **kwargs):
            yield TextChunk(text=prompt)

    runtime = DistributedRuntime(task_queue=SQLiteTaskQueue(path=str(tmp_path / "tasks.db")))
    runtime.register("bot", PlainReplica(), replica_id="a")

    with pytest.raises(
        AttributeError,
        match="does not support stream_ui_events",
    ):
        async for _ in runtime.stream_ui("bot", "hello", session_id="s1"):
            pass

    assert runtime.stats()["inflight"].get("a", 0) == 0


def test_distributed_runtime_defaults_to_sqlite_without_database_url(monkeypatch):
    monkeypatch.delenv("LOGOS_ADK_TASK_QUEUE_DSN", raising=False)
    monkeypatch.delenv("LOGOS_ADK_DATABASE_URL", raising=False)
    monkeypatch.delenv("DATABASE_URL", raising=False)

    runtime = DistributedRuntime()

    assert isinstance(runtime.task_queue, SQLiteTaskQueue)


def test_distributed_runtime_prefers_postgres_with_database_url(monkeypatch):
    monkeypatch.setenv("DATABASE_URL", "postgresql://user:pass@localhost/prod")

    runtime = DistributedRuntime()

    assert isinstance(runtime.task_queue, PostgreSQLTaskQueue)
    assert runtime.task_queue.dsn == "postgresql://user:pass@localhost/prod"
