"""Unit tests for PgSessionStore (sequence_number + archive)."""

import pytest

from navi.core.pg_session_store import PgSessionStore
from navi.core.session import Session
from navi.llm.base import Message
from tests.conftest_factory import FakeConnection, FakePool


@pytest.mark.asyncio
async def test_save_assigns_sequence_numbers_to_new_messages():
    conn = FakeConnection()
    conn.enqueue("OK")  # UPDATE sessions
    conn.enqueue(None)  # executemany UPDATE (empty)
    conn.enqueue("OK")  # INSERT executemany
    conn.enqueue("OK")  # UPDATE next_sequence

    pool = FakePool(conn)
    store = PgSessionStore(pool)
    store._initialized = True

    session = Session(profile_id="test")
    session.messages = [
        Message(role="user", content="hi"),
        Message(role="assistant", content="hello"),
    ]
    for m in session.messages:
        assert m.sequence_number == -1

    await store.save(session)

    assert session.messages[0].sequence_number == 0
    assert session.messages[1].sequence_number == 1
    assert session.db_next_sequence == 2


@pytest.mark.asyncio
async def test_save_updates_existing_rows_by_sequence_number():
    conn = FakeConnection()
    conn.enqueue("OK")  # UPDATE sessions
    conn.enqueue(None)  # executemany UPDATE
    conn.enqueue("OK")  # INSERT (empty)
    conn.enqueue("OK")  # UPDATE next_sequence

    pool = FakePool(conn)
    store = PgSessionStore(pool)
    store._initialized = True

    session = Session(profile_id="test")
    m1 = Message(role="user", content="hi")
    m1.sequence_number = 0
    m2 = Message(role="assistant", content="hello")
    m2.sequence_number = 1
    session.messages = [m1, m2]
    session.db_next_sequence = 2

    await store.save(session)

    executemany_calls = [c for c in conn.calls if c[0] == "executemany"]
    assert len(executemany_calls) == 1
    assert "UPDATE session_messages" in executemany_calls[0][1]


@pytest.mark.asyncio
async def test_archive_old_messages_sends_correct_sql():
    conn = FakeConnection()
    conn.enqueue("INSERT 0 3")  # copy to archive
    conn.enqueue("OK")  # DELETE
    conn.enqueue("OK")  # UPDATE archive_threshold

    pool = FakePool(conn)
    store = PgSessionStore(pool)
    store._initialized = True

    n = await store.archive_old_messages("sid-1", keep_seq_threshold=100)
    assert n == 3

    execute_calls = [c for c in conn.calls if c[0] == "execute"]
    assert any("DELETE FROM session_messages" in c[1] for c in execute_calls)
    assert any("UPDATE sessions SET archive_threshold" in c[1] for c in execute_calls)
