"""调度服务单元测试。"""
import pytest

from app.constants import TaskStatus
from app.models.schemas import AgentInfo, TaskInfo
from app.repository.agent_repo import AgentRepo
from app.repository.task_repo import TaskRepo
from app.services.scheduler import Scheduler


@pytest.mark.asyncio
async def test_dispatch_binds_min_load_agent(redis):
    task_repo = TaskRepo(redis)
    agent_repo = AgentRepo(redis)

    # 两个 compile Agent，b2 负载更低
    await agent_repo.upsert(AgentInfo(agent_id="b1", endpoint="e1", agent_tags=["compile"], max_concurrent=2, current_load=1))
    await agent_repo.upsert(AgentInfo(agent_id="b2", endpoint="e2", agent_tags=["compile"], max_concurrent=2, current_load=0))

    await task_repo.create(TaskInfo(request_id="t1", task_type="compile", task_tags=["compile"]))

    sched = Scheduler(redis, task_repo, agent_repo)
    ok = await sched.dispatch("t1")
    assert ok is True

    task = await task_repo.get("t1")
    assert task.agent_id == "b2"
    assert task.status == TaskStatus.RUNNING
    assert "t1" in await task_repo.running_tasks()


@pytest.mark.asyncio
async def test_dispatch_no_agent_stays_pending(redis):
    task_repo = TaskRepo(redis)
    agent_repo = AgentRepo(redis)
    await task_repo.create(TaskInfo(request_id="t1", task_type="compile", task_tags=["compile"]))

    sched = Scheduler(redis, task_repo, agent_repo)
    ok = await sched.dispatch("t1")
    assert ok is False
    task = await task_repo.get("t1")
    assert task.status == TaskStatus.PENDING


@pytest.mark.asyncio
async def test_dispatch_skips_full_agent(redis):
    task_repo = TaskRepo(redis)
    agent_repo = AgentRepo(redis)
    # 唯一 agent 已满载
    await agent_repo.upsert(AgentInfo(agent_id="b1", endpoint="e1", agent_tags=["compile"], max_concurrent=1, current_load=1))
    await task_repo.create(TaskInfo(request_id="t1", task_type="compile", task_tags=["compile"]))

    sched = Scheduler(redis, task_repo, agent_repo)
    assert await sched.dispatch("t1") is False


@pytest.mark.asyncio
async def test_dispatch_pending(redis):
    task_repo = TaskRepo(redis)
    agent_repo = AgentRepo(redis)
    await agent_repo.upsert(AgentInfo(agent_id="b1", endpoint="e1", agent_tags=["compile"], max_concurrent=2))
    await task_repo.create(TaskInfo(request_id="t1", task_type="compile", task_tags=["compile"]))
    await task_repo.create(TaskInfo(request_id="t2", task_type="compile", task_tags=["compile"]))

    sched = Scheduler(redis, task_repo, agent_repo)
    n = await sched.dispatch_pending(10)
    assert n == 2