"""
tests/test_locomo_isolation.py — Tests for LoCoMo benchmark isolation (#1177).

Verifies that:
  1. EvalHarness uses the provided owner_id in remember() and retrieve() calls
  2. run_eval generates a locomo-eval-<ts> owner_id when --isolate (default)
  3. Cleanup (_cleanup_eval_owner) hard-deletes all claims + entities for the owner
  4. --no-isolate passes no owner_id to the harness (uses plugin default)
"""
from __future__ import annotations

import asyncio
import time
from dataclasses import dataclass, field
from typing import Any, List, Optional
from unittest.mock import AsyncMock, MagicMock, patch
import pytest

from cortex.eval.harness import EvalHarness
from cortex.eval.run_eval import _LOCOMO_EVAL_PREFIX, _cleanup_eval_owner


# ---------------------------------------------------------------------------
# Minimal mock plugin
# ---------------------------------------------------------------------------

@dataclass
class _MockRememberResult:
    session_id: str = "test"
    claims_stored: int = 1
    claims_extracted: int = 1
    ok: bool = True
    errors: List[str] = field(default_factory=list)


@dataclass
class _MockRetrievalResult:
    formatted_context: str = "Jordan is a software engineer"
    memories: List[Any] = field(default_factory=list)


class MockCortexPlugin:
    """Minimal mock of CortexPlugin that records calls."""

    def __init__(self):
        self.remember_calls: List[dict] = []
        self.retrieve_calls: List[dict] = []
        self.llm = None
        self.config = MagicMock()
        self.config.owner_id = "production"

    async def remember(self, conversation, session_id=None, owner_id=None, **kwargs):
        self.remember_calls.append({"session_id": session_id, "owner_id": owner_id})
        return _MockRememberResult(session_id=session_id or "")

    async def retrieve(self, query, owner_id=None, **kwargs):
        self.retrieve_calls.append({"query": query, "owner_id": owner_id})
        return _MockRetrievalResult()


# ---------------------------------------------------------------------------
# Minimal dataset fixture
# ---------------------------------------------------------------------------

def _minimal_dataset():
    """Return a tiny LoCoMoDataset with one session and one QA pair."""
    from cortex.eval.locomo import LoCoMoSession, LoCoMoDataset, LoCoMoQAPair

    session = LoCoMoSession(
        session_id="test-session-001",
        conversation=[
            {"role": "user", "content": "My name is Jordan."},
            {"role": "assistant", "content": "Nice to meet you, Jordan."},
        ],
        qa_pairs=[
            LoCoMoQAPair(
                question_id="q1",
                question="What is Jordan's name?",
                answer="Jordan",
            )
        ],
    )
    dataset = LoCoMoDataset(name="test", sessions=[session])
    return dataset


# ---------------------------------------------------------------------------
# Test 1: owner_id flows into remember() and retrieve()
# ---------------------------------------------------------------------------

@pytest.mark.asyncio
async def test_harness_passes_owner_id_to_remember_and_retrieve():
    """EvalHarness must propagate owner_id to both remember() and retrieve()."""
    plugin = MockCortexPlugin()
    test_owner = "locomo-eval-9999999"

    harness = EvalHarness(plugin=plugin, owner_id=test_owner)
    dataset = _minimal_dataset()
    await harness.run_locomo_dataset(dataset)

    # remember() must have been called with our owner_id
    assert len(plugin.remember_calls) > 0, "remember() was never called"
    for call in plugin.remember_calls:
        assert call["owner_id"] == test_owner, (
            f"remember() called with owner_id={call['owner_id']!r}, expected {test_owner!r}"
        )

    # retrieve() must have been called with our owner_id
    assert len(plugin.retrieve_calls) > 0, "retrieve() was never called"
    for call in plugin.retrieve_calls:
        assert call["owner_id"] == test_owner, (
            f"retrieve() called with owner_id={call['owner_id']!r}, expected {test_owner!r}"
        )


# ---------------------------------------------------------------------------
# Test 2: no owner_id → None is passed (plugin uses its own default)
# ---------------------------------------------------------------------------

@pytest.mark.asyncio
async def test_harness_no_owner_id_passes_none():
    """When owner_id is not set, None is passed so plugin uses its own default."""
    plugin = MockCortexPlugin()

    harness = EvalHarness(plugin=plugin)  # no owner_id
    dataset = _minimal_dataset()
    await harness.run_locomo_dataset(dataset)

    for call in plugin.remember_calls:
        assert call["owner_id"] is None, (
            f"expected owner_id=None when not set, got {call['owner_id']!r}"
        )
    for call in plugin.retrieve_calls:
        assert call["owner_id"] is None, (
            f"expected owner_id=None when not set, got {call['owner_id']!r}"
        )


# ---------------------------------------------------------------------------
# Test 3: default owner_id prefix is locomo-eval-<timestamp>
# ---------------------------------------------------------------------------

def test_locomo_eval_prefix():
    """The auto-generated owner_id must start with 'locomo-eval-'."""
    ts = int(time.time())
    generated = f"{_LOCOMO_EVAL_PREFIX}-{ts}"
    assert generated.startswith("locomo-eval-"), f"prefix wrong: {generated}"
    assert str(ts) in generated


# ---------------------------------------------------------------------------
# Test 4: _cleanup_eval_owner deletes claims and entities
# ---------------------------------------------------------------------------

@pytest.mark.asyncio
async def test_cleanup_deletes_claims_and_entities():
    """_cleanup_eval_owner must hard-delete all claims + entities for the owner."""
    # Build a mock storage that returns fake claim IDs and entities
    mock_storage = MagicMock()

    # Mock _get_conn as an async context manager that returns a connection
    mock_conn = MagicMock()

    # execute() returns an async context manager with rows, also awaitable
    class _AsyncCMRows:
        def __init__(self, rows):
            self._rows = rows
        async def __aenter__(self):
            return self
        async def __aexit__(self, *_):
            pass
        def __await__(self):
            return self.__aenter__().__await__()
        async def fetchall(self):
            return self._rows

    class _AsyncCMNone:
        async def __aenter__(self):
            return self
        async def __aexit__(self, *_):
            pass
        def __await__(self):
            return self.__aenter__().__await__()

    # conn.execute returns a context manager whose fetchall returns claim/entity rows
    claim_rows = [{"id": "claim-aaa"}, {"id": "claim-bbb"}]
    entity_rows = [{"id": "entity-xxx"}, {"id": "entity-yyy"}]

    _execute_calls = []

    def _execute_side_effect(sql, params=()):
        _execute_calls.append((sql, params))
        if "SELECT id FROM semantic_claim" in sql:
            return _AsyncCMRows(claim_rows)
        if "SELECT id FROM entity" in sql:
            return _AsyncCMRows(entity_rows)
        return _AsyncCMNone()

    mock_conn.execute = MagicMock(side_effect=_execute_side_effect)
    mock_conn.commit = AsyncMock()

    mock_storage._get_conn = AsyncMock(return_value=mock_conn)
    mock_storage._write_lock = asyncio.Lock()
    mock_storage._vec_loaded = False
    mock_storage.list_entities = AsyncMock(return_value=[])

    mock_plugin = MagicMock()
    mock_plugin.storage = mock_storage

    test_owner = "locomo-eval-cleanup-test"
    await _cleanup_eval_owner(mock_plugin, test_owner, verbose=False)

    # _cleanup_eval_owner uses bulk SQL DELETEs, not delete_claim.
    # Verify that DELETE statements were issued for claims and entities.
    delete_sqls = [sql for sql, _ in _execute_calls if "DELETE" in sql.upper()]
    assert len(delete_sqls) >= 2, (
        f"expected at least 2 DELETE SQL calls (claims + entities), "
        f"got {len(delete_sqls)}: {delete_sqls}"
    )


# ---------------------------------------------------------------------------
# Test 5: _cleanup_eval_owner is safe when storage is None
# ---------------------------------------------------------------------------

@pytest.mark.asyncio
async def test_cleanup_safe_when_no_storage():
    """_cleanup_eval_owner must not raise when plugin.storage is None."""
    mock_plugin = MagicMock()
    mock_plugin.storage = None

    # Should complete without raising
    await _cleanup_eval_owner(mock_plugin, "locomo-eval-nostorage", verbose=False)


# ---------------------------------------------------------------------------
# Test 6: --isolate flag triggers auto-generated owner_id (arg parsing)
# ---------------------------------------------------------------------------

def test_argparse_isolate_default_is_true():
    """--isolate must be the default (True) so isolation is on without any flag."""
    from cortex.eval.run_eval import _build_parser

    parser = _build_parser()

    # No isolation flags provided → should default to isolate=True
    args = parser.parse_args(["--benchmark", "locomo", "--sample"])
    assert args.isolate is True, f"expected isolate=True by default, got {args.isolate}"


def test_argparse_no_isolate_sets_false():
    """--no-isolate must set isolate=False."""
    from cortex.eval.run_eval import _build_parser

    parser = _build_parser()
    args = parser.parse_args(["--benchmark", "locomo", "--sample", "--no-isolate"])
    assert args.isolate is False, f"expected isolate=False, got {args.isolate}"


def test_argparse_isolate_explicit_sets_true():
    """--isolate explicitly still sets isolate=True."""
    from cortex.eval.run_eval import _build_parser

    parser = _build_parser()
    args = parser.parse_args(["--benchmark", "locomo", "--sample", "--isolate"])
    assert args.isolate is True, f"expected isolate=True, got {args.isolate}"
