"""Agent 池存储层单元测试。"""
import time

import pytest

from app.constants import AgentStatus
from app.models.schemas import AgentInfo


@pytest.mark.asyncio
async def test_upsert_and_get(agent_repo):
    a = AgentInfo(agent_id="a1", endpoint="agent-1:5000", agent_tags=["compile", "test"], max_concurrent=2)
    created = await agent_repo.upsert(a)
    assert created is True

    got = await agent_repo.get("a1")
    assert got is not None
    assert got.agent_tags == ["compile", "test"]
    assert got.status == AgentStatus.READY
    assert got.priority == 3  # 默认优先级


@pytest.mark.asyncio
async def test_priority_roundtrip_and_update(agent_repo):
    a = AgentInfo(agent_id="a1", endpoint="e1", agent_tags=["compile"], priority=1)
    await agent_repo.upsert(a)
    assert (await agent_repo.get("a1")).priority == 1

    assert await agent_repo.update_priority("a1", 5) is True
    assert (await agent_repo.get("a1")).priority == 5

    assert await agent_repo.update_priority("missing", 2) is False


@pytest.mark.asyncio
async def test_beat_and_stale(agent_repo):
    a = AgentInfo(agent_id="a1", endpoint="agent-1:5000", agent_tags=["compile"])
    await agent_repo.upsert(a)
    assert await agent_repo.beat("a1", 1) is True

    # 心跳为最新，不应 stale
    stale = await agent_repo.stale_agents(1)
    assert "a1" not in stale


@pytest.mark.asyncio
async def test_by_tags(agent_repo):
    a1 = AgentInfo(agent_id="a1", endpoint="e1", agent_tags=["compile"], max_concurrent=2)
    a2 = AgentInfo(agent_id="a2", endpoint="e2", agent_tags=["compile", "test"], max_concurrent=2)
    a3 = AgentInfo(agent_id="a3", endpoint="e3", agent_tags=["test"], max_concurrent=2)
    await agent_repo.upsert(a1)
    await agent_repo.upsert(a2)
    await agent_repo.upsert(a3)

    ids = {a.agent_id for a in await agent_repo.by_tags(["compile"])}
    assert ids == {"a1", "a2"}

    ids = {a.agent_id for a in await agent_repo.by_tags(["compile", "test"])}
    assert ids == {"a2"}


@pytest.mark.asyncio
async def test_by_tags_excludes_offline_unavailable_stopping(agent_repo):
    """调度候选排除 offline/unavailable/stopping，允许 ready/processing（processing 仍有并发容量）。"""
    a1 = AgentInfo(agent_id="a1", endpoint="e1", agent_tags=["compile"], max_concurrent=2)
    a2 = AgentInfo(agent_id="a2", endpoint="e2", agent_tags=["compile"], max_concurrent=2)
    a3 = AgentInfo(agent_id="a3", endpoint="e3", agent_tags=["compile"], max_concurrent=2)
    a4 = AgentInfo(agent_id="a4", endpoint="e4", agent_tags=["compile"], max_concurrent=2)
    a5 = AgentInfo(agent_id="a5", endpoint="e5", agent_tags=["compile"], max_concurrent=2)
    await agent_repo.upsert(a1)
    await agent_repo.upsert(a2)
    await agent_repo.upsert(a3)
    await agent_repo.upsert(a4)
    await agent_repo.upsert(a5)
    await agent_repo.update_status("a2", AgentStatus.OFFLINE)      # 离线（心跳失联）不可调度
    await agent_repo.set_unavailable("a3")                          # 手动不可用不可调度
    await agent_repo.update_status("a4", AgentStatus.STOPPING)     # 停止中不可调度
    await agent_repo.update_status("a5", AgentStatus.PROCESSING)   # 处理中但还有容量，可调度

    ids = {a.agent_id for a in await agent_repo.by_tags(["compile"])}
    assert ids == {"a1", "a5"}


@pytest.mark.asyncio
async def test_unavailable_heartbeat_not_revive(agent_repo):
    """手动置不可用（unavailable）后，即使发心跳也不复活。"""
    a1 = AgentInfo(agent_id="a1", endpoint="e1", agent_tags=["compile"], max_concurrent=2)
    await agent_repo.upsert(a1)
    await agent_repo.set_unavailable("a1")
    assert (await agent_repo.get("a1")).status == AgentStatus.UNAVAILABLE
    # 心跳不应复活
    await agent_repo.beat("a1", 0)
    assert (await agent_repo.get("a1")).status == AgentStatus.UNAVAILABLE
    # 恢复可用
    await agent_repo.set_available("a1")
    assert (await agent_repo.get("a1")).status == AgentStatus.READY


@pytest.mark.asyncio
async def test_unavailable_distinct_from_offline(agent_repo):
    """离线（心跳失联）与不可用（手动禁用）是两个独立状态。"""
    a1 = AgentInfo(agent_id="a1", endpoint="e1", agent_tags=["compile"], max_concurrent=2)
    a2 = AgentInfo(agent_id="a2", endpoint="e2", agent_tags=["compile"], max_concurrent=2)
    await agent_repo.upsert(a1)
    await agent_repo.upsert(a2)
    # a1 心跳失联→离线；a2 手动禁用→不可用
    await agent_repo.update_status("a1", AgentStatus.OFFLINE)
    await agent_repo.set_unavailable("a2")
    assert (await agent_repo.get("a1")).status == AgentStatus.OFFLINE
    assert (await agent_repo.get("a2")).status == AgentStatus.UNAVAILABLE
    # 发心跳：离线的 a1 恢复 ready，不可用的 a2 保持 unavailable
    await agent_repo.beat("a1", 0)
    await agent_repo.beat("a2", 0)
    assert (await agent_repo.get("a1")).status == AgentStatus.READY
    assert (await agent_repo.get("a2")).status == AgentStatus.UNAVAILABLE


@pytest.mark.asyncio
async def test_mark_idle_keeps_unavailable(agent_repo):
    """任务结束后 mark_idle：unavailable 的 agent 保持不可用，其余置回 ready。"""
    a1 = AgentInfo(agent_id="a1", endpoint="e1", agent_tags=["compile"], max_concurrent=2)
    a2 = AgentInfo(agent_id="a2", endpoint="e2", agent_tags=["compile"], max_concurrent=2)
    await agent_repo.upsert(a1)
    await agent_repo.upsert(a2)
    await agent_repo.set_unavailable("a1")
    await agent_repo.update_status("a2", AgentStatus.PROCESSING)
    await agent_repo.mark_idle("a1")
    await agent_repo.mark_idle("a2")
    assert (await agent_repo.get("a1")).status == AgentStatus.UNAVAILABLE
    assert (await agent_repo.get("a2")).status == AgentStatus.READY


@pytest.mark.asyncio
async def test_heartbeat_recovers_stale_offline(agent_repo):
    """超时离线（非手动）的 agent 在发心跳后恢复为 ready。"""
    a1 = AgentInfo(agent_id="a1", endpoint="e1", agent_tags=["compile"], max_concurrent=2)
    await agent_repo.upsert(a1)
    await agent_repo.update_status("a1", AgentStatus.OFFLINE)
    assert (await agent_repo.get("a1")).status == AgentStatus.OFFLINE
    await agent_repo.beat("a1", 0)
    assert (await agent_repo.get("a1")).status == AgentStatus.READY


@pytest.mark.asyncio
async def test_heartbeat_keeps_busy_states(agent_repo):
    """processing/stopping 状态不因心跳改变。"""
    for st in (AgentStatus.PROCESSING, AgentStatus.STOPPING):
        a = AgentInfo(agent_id=f"a-{st.value}", endpoint="e1", agent_tags=["compile"], max_concurrent=2)
        await agent_repo.upsert(a)
        await agent_repo.update_status(a.agent_id, st)
        await agent_repo.beat(a.agent_id, 1)
        assert (await agent_repo.get(a.agent_id)).status == st


@pytest.mark.asyncio
async def test_remove(agent_repo):
    a1 = AgentInfo(agent_id="a1", endpoint="e1", agent_tags=["compile"])
    await agent_repo.upsert(a1)
    await agent_repo.remove("a1")
    assert await agent_repo.get("a1") is None
    assert await agent_repo.by_tags(["compile"]) == []