"""Tests for GemmaLLM context compaction (agent.GemmaLLM._compact_context).

Compaction runs a real side LLM call to the xNAS, so these are marked slow.
The agent module is imported inside the container via docker exec.
"""

import pytest

from conftest import GEMMA_BASE_URL, GEMMA_MODEL, docker_exec_python

# {num_messages} and {do_chat} are filled per test; all other braces are literal.
TEMPLATE = '''
import logging, sys
from livekit.agents import llm as lk_llm
sys.path.insert(0, "/opt/voice-agent")
from agent import GemmaLLM

logging.basicConfig(level=logging.INFO, stream=sys.stdout)

llm = GemmaLLM(
    model="{model}",
    base_url="{base_url}",
    api_key="not-needed",
    max_completion_tokens=100,
    extra_body={{"chat_template_kwargs": {{"enable_thinking": False}}}},
)

items = [lk_llm.ChatMessage(role="system", content=["You are Hope, a voice assistant."])]
for i in range({num_messages}):
    items.append(lk_llm.ChatMessage(role="user", content=["user message number " + str(i)]))
    items.append(lk_llm.ChatMessage(role="assistant", content=["agent reply number " + str(i)]))

ctx = lk_llm.ChatContext(items)
print("BUILT ctx with " + str(len(ctx.items)) + " items")

import asyncio

async def main():
    try:
        compacted = llm._compact_context(ctx, None)
    except Exception as e:
        print("COMPACT_ERROR " + repr(e))
        return
    print("AFTER items=" + str(len(compacted.items)))
    for it in compacted.items:
        role = getattr(it, "role", None)
        content = GemmaLLM._msg_content_text(it)
        print("ITEM " + str(role) + ": " + content[:120])

    {do_chat}

asyncio.run(main())
'''

CHAT_SNIPPET = '    stream = llm.chat(chat_ctx=compacted)\n    collected = await stream.collect()\n    print("CHAT_OK text_len=" + str(len(collected.text)))'


def _run(num_messages: int, *, do_chat: bool = False) -> str:
    code = TEMPLATE.format(
        model=GEMMA_MODEL,
        base_url=GEMMA_BASE_URL,
        num_messages=num_messages,
        do_chat=CHAT_SNIPPET if do_chat else "pass",
    )
    proc = docker_exec_python(code, timeout=120, check=False)
    return proc.stdout + (("\n" + proc.stderr) if proc.stderr else "")


@pytest.mark.slow
def test_compaction_triggers_on_size():
    """A context with > 24 items must compact without crashing and chat() must still work."""
    # 13 pairs = 26 non-system + 1 system = 27 items (> threshold of 24)
    out = _run(13, do_chat=True)
    assert "BUILT ctx with 27 items" in out
    assert "Context compacted:" in out
    assert "CHAT_OK text_len=" in out


@pytest.mark.slow
def test_compaction_preserves_system_prompt():
    out = _run(28)
    lines = [l for l in out.splitlines() if l.startswith("ITEM system:")]
    assert lines, f"no system message survived compaction:\n{out}"
    assert "You are Hope" in lines[0]


def test_short_context_no_compaction():
    """Context with non_sys <= KEEP_RECENT must NOT be compacted."""
    # system + 4 user/assistant pairs = 9 items, non_sys=8 <= KEEP_RECENT(10)
    out = _run(4)
    assert "BUILT ctx with 9 items" in out
    # _compact_context should return the context unchanged (early return)
    assert "AFTER items=9" in out, f"expected no compaction:\n{out}"