"""
Tests for session key continuity (#1179).

Validates that session_key provides stable conversation identity
across session_id changes (e.g., gateway restarts).
"""

from __future__ import annotations

import pytest
import pytest_asyncio

from unittest.mock import AsyncMock

from cortex.hippo.config import HippoConfig
from cortex.hippo.engine import HippoEngine
from cortex.hippo.models import ConversationRecord


@pytest.fixture
def hippo_config() -> HippoConfig:
    return HippoConfig(
        enabled=True,
        db_path=":memory:",
        fresh_tail_count=4,
        max_tokens=10_000,
        context_threshold=0.75,
        context_window_tokens=100,
        leaf_trigger_margin=100,
    )


@pytest_asyncio.fixture
async def engine(hippo_config: HippoConfig) -> HippoEngine:
    eng = HippoEngine(config=hippo_config)
    await eng.initialize(
        db_path=":memory:",
        owner_id="test-owner",
        conversation_id="test-session",
    )
    yield eng
    await eng.shutdown()


# ── Session Key Tests ───────────────────────────────────────────────────


@pytest.mark.asyncio
async def test_session_key_reuses_conversation(engine: HippoEngine):
    """Same session_key with different session_id → same conversation."""
    conv1 = await engine.conversation_store.get_or_create_conversation(
        "owner-a", "session-1", session_key="stable-key-1"
    )
    conv2 = await engine.conversation_store.get_or_create_conversation(
        "owner-a", "session-2", session_key="stable-key-1"
    )
    assert conv1.conversation_id == conv2.conversation_id


@pytest.mark.asyncio
async def test_session_key_creates_new_for_new_key(engine: HippoEngine):
    """Different session_key with different session_id → different conversation."""
    conv1 = await engine.conversation_store.get_or_create_conversation(
        "owner-a", "session-1", session_key="key-alpha"
    )
    conv2 = await engine.conversation_store.get_or_create_conversation(
        "owner-a", "session-2", session_key="key-beta"
    )
    assert conv1.conversation_id != conv2.conversation_id


@pytest.mark.asyncio
async def test_no_session_key_falls_back_to_session_id(engine: HippoEngine):
    """Without session_key, lookup uses session_id (backward compat)."""
    conv1 = await engine.conversation_store.get_or_create_conversation(
        "owner-b", "session-x"
    )
    conv2 = await engine.conversation_store.get_or_create_conversation(
        "owner-b", "session-x"
    )
    assert conv1.conversation_id == conv2.conversation_id


@pytest.mark.asyncio
async def test_session_key_updates_session_id(engine: HippoEngine):
    """When session_key matches, session_id is updated to current value."""
    conv1 = await engine.conversation_store.get_or_create_conversation(
        "owner-c", "old-session", session_key="persistent-key"
    )
    assert conv1.session_id == "old-session"

    conv2 = await engine.conversation_store.get_or_create_conversation(
        "owner-c", "new-session", session_key="persistent-key"
    )
    assert conv2.conversation_id == conv1.conversation_id
    assert conv2.session_id == "new-session"


@pytest.mark.asyncio
async def test_session_key_column_exists(engine: HippoEngine):
    """Migration: session_key column exists in hippo_conversations."""
    async with engine._db.conn.execute(
        "PRAGMA table_info(hippo_conversations)"
    ) as cur:
        columns = [row[1] async for row in cur]
    assert "session_key" in columns


@pytest.mark.asyncio
async def test_engine_threads_session_key():
    """Engine.initialize() threads session_key to conversation store."""
    eng = HippoEngine(HippoConfig(db_path=":memory:"))
    conv = await eng.initialize(
        db_path=":memory:",
        owner_id="test-owner",
        conversation_id="sess-1",
        session_key="eng-key-1",
    )
    assert conv.session_key == "eng-key-1"

    # Re-initialize with same key, different session_id
    eng2 = HippoEngine(HippoConfig(db_path=":memory:"))
    # Need same DB — use engine's db directly
    # Instead: create a second engine on the same DB path
    # For in-memory, we test through conversation_store
    conv2 = await eng.conversation_store.get_or_create_conversation(
        "test-owner", "sess-2", session_key="eng-key-1"
    )
    assert conv2.conversation_id == conv.conversation_id
    assert conv2.session_id == "sess-2"
    await eng.shutdown()


@pytest.mark.asyncio
async def test_router_accepts_session_key():
    """HTTP endpoint models accept session_key field."""
    from cortex.hippo.router import (
        BootstrapRequest,
        AfterTurnRequest,
        IngestRequest,
    )

    # BootstrapRequest
    req = BootstrapRequest(session_id="s1", session_key="k1")
    assert req.session_key == "k1"

    # AfterTurnRequest
    req2 = AfterTurnRequest(
        session_id="s1",
        messages=[],
        session_key="k2",
    )
    assert req2.session_key == "k2"

    # IngestRequest
    req3 = IngestRequest(
        conversation_id=1,
        role="user",
        content="hi",
        session_key="k3",
    )
    assert req3.session_key == "k3"

    # Without session_key (backward compat)
    req4 = BootstrapRequest(session_id="s2")
    assert req4.session_key is None


@pytest.mark.asyncio
async def test_after_turn_with_session_key(engine: HippoEngine):
    """session_key flows through after_turn without error."""
    conv = engine._conversation
    # after_turn should accept session_key without raising
    await engine.after_turn(
        conversation_id=conv.conversation_id,
        messages=[
            {"role": "user", "content": "hello"},
            {"role": "assistant", "content": "hi there"},
        ],
        pre_prompt_count=0,
        session_key="after-turn-key",
    )
    # Verify messages were ingested
    msgs = await engine.conversation_store.get_messages(
        conv.conversation_id, "test-owner"
    )
    assert len(msgs) >= 2


@pytest.mark.asyncio
async def test_session_key_stored_on_create(engine: HippoEngine):
    """New conversation stores session_key."""
    conv = await engine.conversation_store.get_or_create_conversation(
        "owner-d", "sess-new", session_key="brand-new-key"
    )
    assert conv.session_key == "brand-new-key"

    # Verify via direct DB query
    async with engine._db.conn.execute(
        "SELECT session_key FROM hippo_conversations WHERE conversation_id = ?",
        (conv.conversation_id,),
    ) as cur:
        row = await cur.fetchone()
    assert row[0] == "brand-new-key"


@pytest.mark.asyncio
async def test_session_key_none_backward_compat(engine: HippoEngine):
    """Conversations without session_key still work."""
    conv = await engine.conversation_store.get_or_create_conversation(
        "owner-e", "sess-old"
    )
    assert conv.session_key is None

    # Can still be found by session_id
    conv2 = await engine.conversation_store.get_or_create_conversation(
        "owner-e", "sess-old"
    )
    assert conv2.conversation_id == conv.conversation_id
