"""Tests for the Gemma LLM API — the exact endpoint Hope's agent uses."""

import time

import pytest

from conftest import GEMMA_BASE_URL, GEMMA_MODEL

WEATHER_TOOL = {
    "type": "function",
    "function": {
        "name": "get_weather",
        "description": "Get the current weather and today's forecast for a location.",
        "parameters": {
            "type": "object",
            "properties": {"location": {"type": "string"}},
            "required": ["location"],
        },
    },
}


def _chat(llm_client, messages, *, tools=None, stream=False, max_tokens=200):
    payload = {
        "model": GEMMA_MODEL,
        "messages": messages,
        "max_tokens": max_tokens,
        "temperature": 0.3,
        "chat_template_kwargs": {"enable_thinking": False},
    }
    if tools:
        payload["tools"] = tools
    if stream:
        payload["stream"] = True
    return llm_client.post("/chat/completions", json=payload)


@pytest.mark.slow
def test_simple_completion(llm_client):
    resp = _chat(
        llm_client,
        [{"role": "user", "content": "What is 2+2? Answer with just the number."}],
        max_tokens=50,
    )
    assert resp.status_code == 200
    data = resp.json()
    content = (data["choices"][0]["message"].get("content") or "").strip()
    assert "4" in content


@pytest.mark.slow
def test_tool_call_weather(llm_client):
    resp = _chat(
        llm_client,
        [{"role": "user", "content": "What is the weather like in Raleigh right now?"}],
        tools=[WEATHER_TOOL],
        max_tokens=100,
    )
    assert resp.status_code == 200
    choice = resp.json()["choices"][0]
    assert choice["finish_reason"] == "tool_calls"
    tool_calls = choice["message"]["tool_calls"]
    assert any(tc["function"]["name"] == "get_weather" for tc in tool_calls)


@pytest.mark.slow
def test_thinking_disabled(llm_client):
    """A normally-slow prompt must answer fast when enable_thinking=false."""
    messages = [
        {"role": "user", "content": "Count from 1 to 20, then tell me the sum."},
    ]
    start = time.monotonic()
    resp = _chat(llm_client, messages, max_tokens=300)
    elapsed = time.monotonic() - start
    assert resp.status_code == 200
    content = (resp.json()["choices"][0]["message"].get("content") or "").strip()
    assert len(content) > 10
    # With thinking enabled this prompt takes well over 3s on the xNAS.
    assert elapsed < 3, f"response took {elapsed:.1f}s — thinking may be enabled"


@pytest.mark.slow
def test_streaming(llm_client):
    resp = _chat(
        llm_client,
        [{"role": "user", "content": "Say the words: the quick brown fox jumps over the lazy dog."}],
        stream=True,
        max_tokens=100,
    )
    assert resp.status_code == 200

    chunks = []
    for line in resp.iter_lines():
        if not line or not line.startswith("data: "):
            continue
        payload = line[6:].strip()
        if payload == "[DONE]":
            break
        chunk = __import__("json").loads(payload)
        delta = (chunk.get("choices") or [{}])[0].get("delta", {})
        if delta.get("content"):
            chunks.append(delta["content"])

    assert len(chunks) > 1, f"expected multiple streamed chunks, got {len(chunks)}"
    assembled = "".join(chunks)
    assert "quick brown fox" in assembled.lower()