# type: ignore
"""Single-node integration tests.

Run with:
    uv run pytest tests/test_1node.py -v
"""

from __future__ import annotations

import time

import pytest
from exo_tools.harness import is_model_downloaded, place_instance

from .framework import DEFAULT_MODEL, InstanceSpec


@pytest.mark.cluster(count=1)
@pytest.mark.instance(DEFAULT_MODEL)
def test_place_instance_and_chat(session):
    resp = session.chat("Say hello in one sentence.")
    assert len(resp) > 0


@pytest.mark.cluster(count=1)
@pytest.mark.instance(DEFAULT_MODEL)
def test_chat_multiple_turns(session):
    first_reply = session.chat("What is 2 + 2?")
    assert len(first_reply) > 0

    second_reply = session.multi_turn(
        [
            {"role": "user", "content": "What is 2 + 2?"},
            {"role": "assistant", "content": first_reply},
            {"role": "user", "content": "Now multiply that by 3."},
        ]
    )
    assert len(second_reply) > 0


@pytest.mark.cluster(count=1)
@pytest.mark.instance(DEFAULT_MODEL)
def test_delete_instance(session):
    from exo_tools.harness import wait_for_instance_gone

    session.client.request_json("DELETE", f"/instance/{session.instance_id}")
    wait_for_instance_gone(session.client, session.instance_id, timeout=30.0)
    assert len(session.instances) == 0, (
        f"Expected no instances, found {len(session.instances)}"
    )


@pytest.mark.cluster(count=1)
def test_download_from_scratch(session):
    """Ensure the model is not on the cluster, then place an instance to
    trigger a fresh download and verify inference.
    """
    node_id = next(iter(session.state.get("nodeIdentities", {})))

    # Delete any existing download — the API call is idempotent
    session.client.request_json("DELETE", f"/download/{node_id}/{DEFAULT_MODEL}")

    # Poll until the model is gone (it may already be gone)
    deadline = time.time() + 60.0
    while time.time() < deadline:
        if not is_model_downloaded(session.client, DEFAULT_MODEL):
            break
        time.sleep(2.0)
    else:
        raise AssertionError(f"Expected {DEFAULT_MODEL} to be deleted from cluster")

    place_instance(session.client, DEFAULT_MODEL, timeout=900.0)
    session.instance_spec = InstanceSpec(model_id=DEFAULT_MODEL)
    resp = session.chat("Say hello in one sentence.")
    assert len(resp) > 0