"""Phase 8: Auto Memory Extraction tests.

Tests the extraction pipeline, ZTM integration,
PersistentHandler implementation, and MCP tool definitions.
"""
import json
import uuid

import pytest

from conftest import SmallMockAIClient
from trustmem.extraction.pipeline import (
    ExtractionPipeline,
    _char_ngrams,
    _minhash_jaccard,
    _minhash_signature,
    _parse_json_safe,
)
from trustmem.extraction.prompts import (
    BATCH_CONFLICT_PROMPT,
    CONFLICT_SIMILARITY_THRESHOLD,
    MINHASH_NOOP_THRESHOLD,
    MINHASH_NUM_PERM,
    UNIFIED_EXTRACTION_PROMPT,
    format_conversation,
)
from trustmem.handler.persistent import PersistentHandler
from trustmem.ztm.constants import DEFAULT_CONFIDENCE


# ============================================================
# Mock AI Clients
# ============================================================


class ExtractionMockAIClient(SmallMockAIClient):
    """AI client that returns structured JSON for extraction tests."""

    def llm_judge(self, prompt: str) -> str:
        # Stage 1: Extraction
        if "extract" in prompt.lower() or "atomic facts" in prompt.lower():
            return json.dumps({
                "facts": [
                    {"content": "user plans Tokyo trip next month", "type": "episodic", "confidence": "high"},
                    {"content": "travel budget is 150000 yen", "type": "semantic", "confidence": "high"},
                    {"content": "prefers hotels near train stations", "type": "semantic", "confidence": "high"},
                ]
            })
        # Stage 3: Conflict resolution
        if "compare" in prompt.lower() or "action" in prompt.lower():
            return json.dumps({
                "results": [
                    {"index": 0, "action": "ADD", "reason": "new information"}
                ]
            })
        return "no_contradiction"


class EmptyExtractionMockAIClient(SmallMockAIClient):
    """AI client that returns empty facts (no memorable content)."""

    def llm_judge(self, prompt: str) -> str:
        if "extract" in prompt.lower() or "atomic facts" in prompt.lower():
            return '{"facts": []}'
        return "no_contradiction"


# ============================================================
# Fixtures
# ============================================================


@pytest.fixture
def persistent_handler():
    ai = ExtractionMockAIClient()
    h = PersistentHandler(database_url=":memory:", ai_client=ai)
    return h


@pytest.fixture
def sample_messages():
    return [
        {"role": "user", "content": "I'm planning a trip to Tokyo next month. Budget is 150,000 yen and I want a hotel near a train station."},
        {"role": "assistant", "content": "I'd be happy to help you find a hotel in Tokyo! Let me search for options near major train stations within your budget."},
        {"role": "user", "content": "Also, I have a cat allergy, so please make sure the hotel is pet-free."},
    ]


@pytest.fixture
def short_messages():
    """Messages too short for extraction."""
    return [
        {"role": "user", "content": "Hi"},
        {"role": "assistant", "content": "Hello!"},
    ]


# ============================================================
# Test: Extraction Prompts
# ============================================================


class TestExtractionPrompts:
    """Tests for prompts.py."""

    def test_format_conversation(self, sample_messages):
        result = format_conversation(sample_messages)
        assert "user:" in result
        assert "assistant:" in result
        assert "Tokyo" in result
        assert "cat allergy" in result

    def test_format_conversation_empty(self):
        result = format_conversation([])
        assert result == ""

    def test_unified_prompt_has_placeholder(self):
        assert "{conversation}" in UNIFIED_EXTRACTION_PROMPT

    def test_batch_conflict_prompt_has_placeholder(self):
        assert "{pairs_json}" in BATCH_CONFLICT_PROMPT

    def test_conflict_similarity_threshold(self):
        assert CONFLICT_SIMILARITY_THRESHOLD == 0.60


# ============================================================
# Test: JSON Parsing
# ============================================================


class TestParseJsonSafe:
    """Tests for _parse_json_safe utility."""

    def test_valid_json(self):
        result = _parse_json_safe('{"facts": ["a", "b"]}')
        assert result == {"facts": ["a", "b"]}

    def test_json_with_markdown_fences(self):
        text = '```json\n{"facts": ["a"]}\n```'
        result = _parse_json_safe(text)
        assert result == {"facts": ["a"]}

    def test_json_with_plain_fences(self):
        text = '```\n{"facts": []}\n```'
        result = _parse_json_safe(text)
        assert result == {"facts": []}

    def test_json_with_surrounding_text(self):
        text = 'Here is the result: {"facts": ["test"]} hope that helps'
        result = _parse_json_safe(text)
        assert result == {"facts": ["test"]}

    def test_invalid_json(self):
        result = _parse_json_safe("this is not json at all")
        assert result is None

    def test_empty_string(self):
        result = _parse_json_safe("")
        assert result is None

    def test_none_input(self):
        result = _parse_json_safe("")
        assert result is None

    def test_whitespace_only(self):
        result = _parse_json_safe("   \n  ")
        assert result is None


# ============================================================
# Test: Extraction Pipeline
# ============================================================


class TestExtractionPipeline:
    """Tests for ExtractionPipeline."""

    def test_extract_with_mock_client(self, sample_messages):
        ai = ExtractionMockAIClient()
        pipeline = ExtractionPipeline(ai_client=ai)
        result = pipeline.run(messages=sample_messages, user_id="test-agent")

        assert "facts" in result
        assert "llm_calls" in result
        assert len(result["facts"]) == 3
        assert result["llm_calls"] >= 1

    def test_extract_fact_structure(self, sample_messages):
        ai = ExtractionMockAIClient()
        pipeline = ExtractionPipeline(ai_client=ai)
        result = pipeline.run(messages=sample_messages, user_id="test-agent")

        for fact in result["facts"]:
            assert "content" in fact
            assert "confidence" in fact
            assert "action" in fact
            assert fact["confidence"] in ("high", "medium", "low")
            assert fact["action"] in ("ADD", "NOOP", "CONTRADICT")

    def test_empty_extraction(self):
        ai = EmptyExtractionMockAIClient()
        pipeline = ExtractionPipeline(ai_client=ai)
        result = pipeline.run(
            messages=[{"role": "user", "content": "Hello"}],
            user_id="test-agent",
        )
        assert result["facts"] == []

    def test_no_vectorstore_all_add(self, sample_messages):
        """Without vectorstore, all facts default to ADD."""
        ai = ExtractionMockAIClient()
        pipeline = ExtractionPipeline(ai_client=ai, vectorstore=None)
        result = pipeline.run(messages=sample_messages, user_id="test-agent")

        for fact in result["facts"]:
            assert fact["action"] == "ADD"
            assert fact["conflict_with"] is None


# ============================================================
# Test: ZTM Integration
# ============================================================


class TestZTMIntegration:
    """Tests for ZTM constants and pipeline integration."""

    def test_default_confidence(self):
        assert DEFAULT_CONFIDENCE == "medium"

    def test_default_confidence_is_valid_level(self):
        from trustmem.ztm.constants import CONFIDENCE_LEVELS
        assert DEFAULT_CONFIDENCE in CONFIDENCE_LEVELS


# ============================================================
# Test: MCP Tool Definition
# ============================================================


# ============================================================
# Test: MinHash Functions (Phase A)
# ============================================================


class TestMinHashFunctions:
    """Tests for MinHash utility functions."""

    def test_char_ngrams_basic(self):
        result = _char_ngrams("hello")
        assert result == {"hel", "ell", "llo"}

    def test_char_ngrams_short(self):
        """Text shorter than n returns the whole text as single shingle."""
        result = _char_ngrams("hi")
        assert result == {"hi"}

    def test_char_ngrams_empty(self):
        result = _char_ngrams("")
        assert result == set()

    def test_char_ngrams_case_insensitive(self):
        result_lower = _char_ngrams("Hello")
        result_upper = _char_ngrams("HELLO")
        assert result_lower == result_upper

    def test_char_ngrams_japanese(self):
        result = _char_ngrams("猫アレルギー")
        assert "猫アレ" in result
        assert "アレル" in result
        assert "レルギ" in result
        assert "ルギー" in result

    def test_minhash_identical(self):
        """Identical text → Jaccard == 1.0."""
        text = "猫アレルギーがある"
        shingles = _char_ngrams(text)
        sig = _minhash_signature(shingles)
        assert _minhash_jaccard(sig, sig) == 1.0

    def test_minhash_different(self):
        """Completely different text → low Jaccard."""
        sig_a = _minhash_signature(_char_ngrams("猫アレルギーがある"))
        sig_b = _minhash_signature(_char_ngrams("Snowflakeのコスト最適化"))
        jaccard = _minhash_jaccard(sig_a, sig_b)
        assert jaccard < 0.3

    def test_minhash_similar(self):
        """Slightly different text → non-zero Jaccard, but not 1.0."""
        sig_a = _minhash_signature(_char_ngrams("東京のホテルを予約したい"))
        sig_b = _minhash_signature(_char_ngrams("東京のホテルを検索したい"))
        jaccard = _minhash_jaccard(sig_a, sig_b)
        # Share most 3-grams except around the differing characters
        assert 0.2 < jaccard < 1.0

    def test_minhash_empty_shingles(self):
        sig = _minhash_signature(set())
        assert len(sig) == MINHASH_NUM_PERM
        assert all(v == 0 for v in sig)

    def test_minhash_signature_length(self):
        sig = _minhash_signature(_char_ngrams("test"))
        assert len(sig) == MINHASH_NUM_PERM

    def test_minhash_constants(self):
        assert MINHASH_NOOP_THRESHOLD == 0.98
        assert MINHASH_NUM_PERM == 128


# ============================================================
# Test: MinHash Pre-filter in Pipeline (Phase A)
# ============================================================


class ConflictTrackingMockAIClient(SmallMockAIClient):
    """AI client that tracks conflict resolution calls."""

    def __init__(self, dim: int = 8):
        super().__init__(dim)
        self.conflict_calls = 0

    def llm_judge(self, prompt: str) -> str:
        # Stage 1: Extraction — return a single fact
        if "extract" in prompt.lower() or "atomic facts" in prompt.lower():
            return json.dumps({
                "facts": [
                    {"content": "猫アレルギーがある", "type": "semantic"},
                ]
            })
        # Stage 3: Conflict resolution
        if "compare" in prompt.lower() or "action" in prompt.lower():
            self.conflict_calls += 1
            return json.dumps({
                "results": [
                    {"index": 0, "action": "CONTRADICT", "reason": "test contradict"}
                ]
            })
        return "no_contradiction"


class TestMinHashPreFilter:
    """Tests for MinHash pre-filtering in the pipeline."""

    def test_exact_duplicate_skips_llm(self):
        """Identical existing content → MinHash auto-NOOP, no LLM call for Stage 3."""
        ai = ConflictTrackingMockAIClient()
        from trustmem.storage.sqlite_store import SQLiteStore
        from trustmem.vectorstore.sqlite import SQLiteVectorStore
        vs = SQLiteVectorStore(SQLiteStore(":memory:"))

        # Pre-populate vectorstore with an existing memory that has content in payload
        existing_content = "猫アレルギーがある"
        vec = ai.embed(existing_content)
        vs.upsert(
            id="existing-mem-1",
            vector=vec,
            payload={
                "user_id": "test-agent",
                "content": existing_content,
                "memory_type": "semantic",
                "memory_id": "existing-mem-1",
                "created_at": "2024-01-01T00:00:00",
            },
        )

        pipeline = ExtractionPipeline(ai_client=ai, vectorstore=vs)
        result = pipeline.run(
            messages=[{"role": "user", "content": "猫アレルギーがある"}],
            user_id="test-agent",
        )

        # Should have 1 fact with NOOP action (MinHash auto-detected)
        assert len(result["facts"]) == 1
        fact = result["facts"][0]
        assert fact["action"] == "NOOP"
        assert "MinHash" in fact["reason"]

        # LLM conflict resolution should NOT have been called
        assert ai.conflict_calls == 0

    def test_different_facts_use_llm(self):
        """Different content → MinHash passes through → LLM called."""
        ai = ConflictTrackingMockAIClient()
        from trustmem.storage.sqlite_store import SQLiteStore
        from trustmem.vectorstore.sqlite import SQLiteVectorStore
        vs = SQLiteVectorStore(SQLiteStore(":memory:"))

        # Pre-populate with a similar but different memory
        existing_content = "犬アレルギーがある"
        vec = ai.embed(existing_content)
        vs.upsert(
            id="existing-mem-2",
            vector=vec,
            payload={
                "user_id": "test-agent",
                "content": existing_content,
                "memory_type": "semantic",
                "memory_id": "existing-mem-2",
                "created_at": "2024-01-01T00:00:00",
            },
        )

        pipeline = ExtractionPipeline(ai_client=ai, vectorstore=vs)
        result = pipeline.run(
            messages=[{"role": "user", "content": "猫アレルギーがある"}],
            user_id="test-agent",
        )

        # If vectorstore returns a match above 0.80 threshold,
        # and MinHash says it's different, LLM should be called
        if result["facts"][0]["action"] != "ADD":
            # LLM was called for conflict resolution
            assert ai.conflict_calls >= 1

    def test_empty_existing_content_fallback(self):
        """Empty existing_content (legacy data) → skip MinHash → LLM fallback."""
        ai = ConflictTrackingMockAIClient()
        from trustmem.storage.sqlite_store import SQLiteStore
        from trustmem.vectorstore.sqlite import SQLiteVectorStore
        vs = SQLiteVectorStore(SQLiteStore(":memory:"))

        # Pre-populate with legacy payload (no content field)
        vec = ai.embed("猫アレルギーがある")
        vs.upsert(
            id="legacy-mem-1",
            vector=vec,
            payload={
                "user_id": "test-agent",
                "memory_type": "semantic",
                # No "content" key — legacy format
            },
        )

        pipeline = ExtractionPipeline(ai_client=ai, vectorstore=vs)
        result = pipeline.run(
            messages=[{"role": "user", "content": "猫アレルギーがある"}],
            user_id="test-agent",
        )

        # With empty content, MinHash can't auto-NOOP → falls through to LLM
        # The LLM should have been called if similarity threshold was met
        assert len(result["facts"]) == 1


# ============================================================
# Test: VectorStore Payload Content (Phase A)
# ============================================================


class TestPayloadContent:
    """Tests for VectorStore payload containing content field."""

    def test_vectorstore_payload_has_content(self, persistent_handler):
        """After add_memory, VectorStore payload should contain content."""
        persistent_handler.add_memory(
            content="テスト用メモリコンテンツ",
            user_id="test-agent",
        )

        # Check the vectorstore has the content in payload
        vs = persistent_handler._vs
        assert vs.count() >= 1

        # Search for the memory
        vec = persistent_handler._ai.embed("テスト用メモリコンテンツ")
        results = vs.search(query_vector=vec, limit=1)
        assert len(results) >= 1
        payload = results[0].get("payload", {})
        assert payload.get("content") == "テスト用メモリコンテンツ"
        assert payload.get("memory_id") is not None
        assert payload.get("created_at") is not None

    def test_few_shot_examples_in_prompt(self):
        """BATCH_CONFLICT_PROMPT should contain few-shot examples."""
        assert "NOOP" in BATCH_CONFLICT_PROMPT
        assert "ADD" in BATCH_CONFLICT_PROMPT
        assert "CONTRADICT" in BATCH_CONFLICT_PROMPT
        assert "Example 1" in BATCH_CONFLICT_PROMPT
        assert "Example 2" in BATCH_CONFLICT_PROMPT
        assert "Example 3" in BATCH_CONFLICT_PROMPT
        assert "Example 4" in BATCH_CONFLICT_PROMPT


# ============================================================
# Test: CONTRADICT Action in Pipeline (Stage 3)
# ============================================================


class ContradictMockAIClient(SmallMockAIClient):
    """AI client that returns CONTRADICT for conflict resolution."""

    def llm_judge(self, prompt: str) -> str:
        if "extract" in prompt.lower() or "atomic facts" in prompt.lower():
            return json.dumps({
                "facts": [
                    {"content": "駅近ホテルが嫌い", "type": "semantic", "confidence": "high"},
                ]
            })
        if "compare" in prompt.lower() or "action" in prompt.lower():
            return json.dumps({
                "results": [
                    {"index": 0, "action": "CONTRADICT", "reason": "opposite preference"}
                ]
            })
        return "no_contradiction"


class TestContradictAction:
    """Tests for CONTRADICT action in the extraction pipeline."""

    def test_contradict_action_passes_through(self):
        """Stage 3 CONTRADICT should be preserved in pipeline result."""
        ai = ContradictMockAIClient()
        from trustmem.storage.sqlite_store import SQLiteStore
        from trustmem.vectorstore.sqlite import SQLiteVectorStore
        vs = SQLiteVectorStore(SQLiteStore(":memory:"))

        # Pre-populate vectorstore with similar but contradictory memory
        existing_content = "駅近ホテルが好き"
        vec = ai.embed(existing_content)
        vs.upsert(
            id="existing-mem-contradict",
            vector=vec,
            payload={
                "user_id": "test-agent",
                "content": existing_content,
                "memory_type": "semantic",
                "memory_id": "existing-mem-contradict",
                "created_at": "2024-01-01T00:00:00",
            },
        )

        pipeline = ExtractionPipeline(ai_client=ai, vectorstore=vs)
        result = pipeline.run(
            messages=[{"role": "user", "content": "駅近ホテルが嫌い"}],
            user_id="test-agent",
        )

        assert len(result["facts"]) == 1
        fact = result["facts"][0]
        # If vectorstore returned a match above threshold, CONTRADICT should be the action
        if fact.get("conflict_with"):
            assert fact["action"] == "CONTRADICT"
            assert fact["conflict_with"] == "existing-mem-contradict"

    def test_contradict_valid_action_in_pipeline(self):
        """CONTRADICT should be accepted as a valid action (not downgraded to NOOP)."""
        from trustmem.extraction.pipeline import ExtractionPipeline
        ai = ContradictMockAIClient()
        pipeline = ExtractionPipeline(ai_client=ai, vectorstore=None)

        # Test _resolve_conflicts directly
        pairs = [{
            "index": 0,
            "new_fact": "駅近ホテルが嫌い",
            "existing_content": "駅近ホテルが好き",
            "existing_id": "test-id",
            "existing_created_at": "2024-01-01",
            "similarity": 0.85,
        }]
        results, calls = pipeline._resolve_conflicts(pairs)
        assert len(results) == 1
        assert results[0]["action"] == "CONTRADICT"
