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

from app.constants import AgentStatus, 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


class _StubRelay:
    """模拟 relay：不发起真实网络请求，直接返回推送成功。"""

    def __init__(self, ok: bool = True):
        self.ok = ok

    async def dispatch_command(self, agent_id: str, request_id: str, payload: dict) -> bool:
        return self.ok


def _sched(redis, task_repo, agent_repo):
    return Scheduler(redis, task_repo, agent_repo, relay=_StubRelay())


@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 = _sched(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_pick_candidate_lowest_priority(redis):
    agent_repo = AgentRepo(redis)

    # 三个 compile Agent：b1 负载最低但优先级高(最不被优先)，b2 负载高但优先级低(最优先)
    await agent_repo.upsert(AgentInfo(agent_id="b1", endpoint="e1", agent_tags=["compile"], max_concurrent=2, current_load=0, priority=5))
    await agent_repo.upsert(AgentInfo(agent_id="b2", endpoint="e2", agent_tags=["compile"], max_concurrent=2, current_load=1, priority=1))
    await agent_repo.upsert(AgentInfo(agent_id="b3", endpoint="e3", agent_tags=["compile"], max_concurrent=2, current_load=0, priority=3))

    sched = Scheduler(redis)
    picked = await sched.pick_candidate(["compile"])
    assert picked is not None
    assert picked.agent_id == "b2"  # 优先级最小优先分配


@pytest.mark.asyncio
async def test_pick_candidate_same_priority_min_load(redis):
    agent_repo = AgentRepo(redis)

    # 同优先级(默认3)时，负载更低者优先
    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))

    sched = Scheduler(redis)
    picked = await sched.pick_candidate(["compile"])
    assert picked is not None
    assert picked.agent_id == "b2"


@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 = _sched(redis, task_repo, agent_repo)
    n = await sched.dispatch_pending(10)
    assert n == 2


@pytest.mark.asyncio
async def test_dispatch_excludes_offline_unavailable_stopping(redis):
    """调度排除 offline（心跳失联）、unavailable（手动不可用）与 stopping 的 agent。"""
    task_repo = TaskRepo(redis)
    agent_repo = AgentRepo(redis)
    # b1 手动不可用，b2 正在停止，b3 离线，b4 就绪
    await agent_repo.upsert(AgentInfo(agent_id="b1", endpoint="e1", agent_tags=["compile"], max_concurrent=2))
    await agent_repo.upsert(AgentInfo(agent_id="b2", endpoint="e2", agent_tags=["compile"], max_concurrent=2))
    await agent_repo.upsert(AgentInfo(agent_id="b3", endpoint="e3", agent_tags=["compile"], max_concurrent=2))
    await agent_repo.upsert(AgentInfo(agent_id="b4", endpoint="e4", agent_tags=["compile"], max_concurrent=2))
    await agent_repo.set_unavailable("b1")
    await agent_repo.update_status("b2", AgentStatus.STOPPING)
    await agent_repo.update_status("b3", AgentStatus.OFFLINE)

    await task_repo.create(TaskInfo(request_id="t1", task_type="compile", task_tags=["compile"]))
    sched = _sched(redis, task_repo, agent_repo)
    ok = await sched.dispatch("t1")
    assert ok is True
    task = await task_repo.get("t1")
    assert task.agent_id == "b4"


@pytest.mark.asyncio
async def test_dispatch_processing_agent_with_capacity(redis):
    """processing 的 agent 若还有并发容量，仍可接收新任务。"""
    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 agent_repo.update_status("b1", AgentStatus.PROCESSING)
    await task_repo.create(TaskInfo(request_id="t1", task_type="compile", task_tags=["compile"]))
    sched = _sched(redis, task_repo, agent_repo)
    assert await sched.dispatch("t1") is True
    assert (await task_repo.get("t1")).agent_id == "b1"


@pytest.mark.asyncio
async def test_dispatch_marks_agent_processing(redis):
    """调度下发成功后，agent 状态流转为 processing。"""
    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"]))
    sched = _sched(redis, task_repo, agent_repo)
    assert await sched.dispatch("t1") is True
    assert (await agent_repo.get("b1")).status == AgentStatus.PROCESSING


@pytest.mark.asyncio
async def test_dispatch_push_fail_reverts_agent_to_ready(redis):
    """推送失败时任务回退 pending，agent 释放负载并回到 ready。"""
    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"]))
    sched = Scheduler(redis, task_repo, agent_repo, relay=_StubRelay(ok=False))
    assert await sched.dispatch("t1") is False
    task = await task_repo.get("t1")
    assert task.status == TaskStatus.PENDING
    assert not task.agent_id
    a = await agent_repo.get("b1")
    assert a.status == AgentStatus.READY
    assert a.current_load == 0