from fastapi.testclient import TestClient

from ai_orchestrator.main import app


def test_health_endpoint_returns_ok() -> None:
    client = TestClient(app)

    response = client.get("/health")

    assert response.status_code == 200
    assert response.json() == {"status": "ok"}


def test_post_tasks_creates_planned_task() -> None:
    client = TestClient(app)

    response = client.post(
        "/tasks",
        json={
            "project_id": "default",
            "goal": "Prepare a plan",
            "inputs": {"source": "test"},
            "execution_mode": "agent_graph",
        },
    )

    payload = response.json()

    assert response.status_code == 200
    assert payload["status"] == "running"
    assert payload["progress"]["total"] == 2


def test_cancel_endpoint_marks_task_cancelled() -> None:
    client = TestClient(app)

    created = client.post(
        "/tasks",
        json={
            "project_id": "default",
            "goal": "Cancel me",
            "inputs": {},
            "execution_mode": "agent_graph",
        },
    ).json()
    response = client.post(f"/tasks/{created['task_id']}/cancel", json={"reason": "stop"})

    assert response.status_code == 200
    assert response.json()["status"] == "cancelled"


def test_worker_registration_is_exposed_via_list_endpoint() -> None:
    client = TestClient(app)

    registration = client.post(
        "/workers/register",
        json={
            "worker_id": "worker_home_pc",
            "name": "Home PC",
            "capabilities": ["file.read", "command.run"],
            "version": "0.1.0",
            "machine": "DESKTOP-1",
            "os": "windows",
        },
    )
    listing = client.get("/workers")

    assert registration.status_code == 200
    assert listing.status_code == 200
    assert any(worker["worker_id"] == "worker_home_pc" for worker in listing.json()["workers"])


def test_worker_command_dispatch_and_poll_cycle() -> None:
    client = TestClient(app)

    registration = client.post(
        "/workers/register",
        json={
            "worker_id": "worker_home_pc",
            "name": "Home PC",
            "capabilities": ["file.read", "command.run"],
            "version": "0.1.0",
            "machine": "DESKTOP-1",
            "os": "windows",
        },
    ).json()
    dispatch = client.post(
        f"/workers/{registration['session_id']}/commands",
        json={
            "task_id": "task_1",
            "node_id": "node_1",
            "command_name": "file.read",
            "args": {"path": "D:/test.txt"},
            "timeout_ms": 30000,
            "policy_context": {"approved": True},
        },
    )
    polled = client.get(f"/workers/{registration['session_id']}/commands")

    assert dispatch.status_code == 200
    assert dispatch.json()["status"] == "queued"
    assert len(polled.json()["commands"]) == 1


def test_worker_websocket_heartbeat_and_poll() -> None:
    client = TestClient(app)
    registration = client.post(
        "/workers/register",
        json={
            "worker_id": "worker_socket_pc",
            "name": "Socket PC",
            "capabilities": ["file.read"],
            "version": "0.1.0",
            "machine": "DESKTOP-2",
            "os": "windows",
        },
    ).json()
    client.post(
        f"/workers/{registration['session_id']}/commands",
        json={
            "task_id": "task_2",
            "node_id": "node_2",
            "command_name": "file.read",
            "args": {"path": "D:/socket.txt"},
            "timeout_ms": 30000,
            "policy_context": {"approved": True},
        },
    )

    with client.websocket_connect(f"/workers/ws/{registration['session_id']}") as websocket:
        websocket.send_json({"event_type": "heartbeat"})
        heartbeat_ack = websocket.receive_json()
        websocket.send_json({"event_type": "poll"})
        commands = websocket.receive_json()

    assert heartbeat_ack["type"] == "ack"
    assert heartbeat_ack["status"] in {"online", "busy"}
    assert commands["type"] == "commands"
    assert len(commands["commands"]) == 1