"""Phase 8.5: Zero Trust Knowledge Graph tests.

Tests entity helpers, extraction pipeline entity support,
ZTM integration, PersistentHandler KG methods,
SQLiteStore entity CRUD, and MCP tool definitions.
"""
import json
import uuid

import pytest

from conftest import SmallMockAIClient
from trustmem.extraction.entities import (
    MENTION_CONFIDENCE_PROMOTE_THRESHOLD,
    VALID_ENTITY_TYPES,
    build_entity_edge,
    build_entity_node,
    compute_edge_id,
    compute_entity_id,
    determine_entity_confidence,
    normalize_predicate,
)
from trustmem.extraction.pipeline import ExtractionPipeline, _parse_json_safe
from trustmem.extraction.prompts import UNIFIED_EXTRACTION_WITH_ENTITIES_PROMPT
from trustmem.handler.persistent import PersistentHandler
from trustmem.constants import VALID_ENTITY_TYPES as MODEL_ENTITY_TYPES
from trustmem.constants import VALID_EDGE_DIRECTIONS
from trustmem.storage.sqlite_store import SQLiteStore
from trustmem.ztm.constants import DEFAULT_CONFIDENCE


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


class KGExtractionMockAIClient(SmallMockAIClient):
    """AI client that returns entities and relations from extraction."""

    def llm_judge(self, prompt: str) -> str:
        if "extract" in prompt.lower() or "atomic facts" in prompt.lower():
            return json.dumps({
                "facts": [
                    {"content": "Alice works at Acme Corp", "type": "semantic", "confidence": "high"},
                    {"content": "Acme Corp is in Tokyo", "type": "semantic", "confidence": "high"},
                ],
                "entities": [
                    {"name": "Alice", "type": "person"},
                    {"name": "Acme Corp", "type": "organization"},
                    {"name": "Tokyo", "type": "location"},
                ],
                "relations": [
                    {"subject": "Alice", "predicate": "works at", "object": "Acme Corp"},
                    {"subject": "Acme Corp", "predicate": "located in", "object": "Tokyo"},
                ],
            })
        if "compare" in prompt.lower() or "action" in prompt.lower():
            return json.dumps({
                "results": [
                    {"index": 0, "action": "ADD", "reason": "new info"}
                ]
            })
        return "no_contradiction"


# ============================================================
# TestEntityHelpers
# ============================================================


class TestEntityHelpers:
    """Unit tests for extraction/entities.py helper functions."""

    def test_compute_entity_id_deterministic(self):
        id1 = compute_entity_id("Alice", "person", "agent1")
        id2 = compute_entity_id("Alice", "person", "agent1")
        assert id1 == id2
        assert id1.startswith("ent_")
        assert len(id1) == 44  # "ent_" + 40 hex chars

    def test_compute_entity_id_case_insensitive(self):
        id1 = compute_entity_id("Alice", "person", "agent1")
        id2 = compute_entity_id("alice", "person", "agent1")
        assert id1 == id2

    def test_compute_entity_id_different_types(self):
        id1 = compute_entity_id("Alice", "person", "agent1")
        id2 = compute_entity_id("Alice", "organization", "agent1")
        assert id1 != id2

    def test_compute_edge_id_deterministic(self):
        id1 = compute_edge_id("ent_a", "works_at", "ent_b")
        id2 = compute_edge_id("ent_a", "works_at", "ent_b")
        assert id1 == id2
        assert id1.startswith("edge_")

    def test_compute_edge_id_different_direction(self):
        id1 = compute_edge_id("ent_a", "works_at", "ent_b")
        id2 = compute_edge_id("ent_b", "works_at", "ent_a")
        assert id1 != id2

    def test_normalize_predicate_basic(self):
        assert normalize_predicate("Works At") == "works_at"
        assert normalize_predicate("located-in") == "located_in"
        assert normalize_predicate("  born  in  ") == "born_in"

    def test_normalize_predicate_preserves_japanese(self):
        assert normalize_predicate("勤務先") == "勤務先"
        assert normalize_predicate("出身 地") == "出身_地"

    def test_normalize_predicate_consecutive_separators(self):
        assert normalize_predicate("a---b") == "a_b"
        assert normalize_predicate("a   b") == "a_b"

    def test_valid_entity_types(self):
        expected = {"person", "organization", "location", "product", "concept", "event", "other"}
        assert VALID_ENTITY_TYPES == expected

    def test_determine_entity_confidence(self):
        assert determine_entity_confidence(1) == "medium"
        assert determine_entity_confidence(2) == "medium"
        assert determine_entity_confidence(3) == "high"
        assert determine_entity_confidence(10) == "high"

    def test_mention_confidence_promote_threshold(self):
        assert MENTION_CONFIDENCE_PROMOTE_THRESHOLD == 3

    def test_build_entity_node(self):
        node = build_entity_node("Alice", "person", "agent1")
        assert node["entity_id"].startswith("ent_")
        assert node["name"] == "Alice"
        assert node["entity_type"] == "person"
        assert node["user_id"] == "agent1"
        assert node["confidence"] == DEFAULT_CONFIDENCE
        assert node["mention_count"] == 1
        assert node["is_active"] is True

    def test_build_entity_node_confidence_always_default(self):
        """Confidence is always DEFAULT_CONFIDENCE for new entities."""
        node = build_entity_node("Alice", "person", "agent1")
        assert node["confidence"] == DEFAULT_CONFIDENCE

    def test_build_entity_edge(self):
        edge = build_entity_edge(
            "ent_a", "Alice", "works_at", "ent_b", "Acme", "agent1"
        )
        assert edge["edge_id"].startswith("edge_")
        assert edge["subject_id"] == "ent_a"
        assert edge["object_id"] == "ent_b"
        assert edge["predicate"] == "works_at"
        assert edge["mention_count"] == 1

    def test_build_entity_edge_normalizes_predicate(self):
        edge = build_entity_edge(
            "ent_a", "Alice", "Works At", "ent_b", "Acme", "agent1"
        )
        assert edge["predicate"] == "works_at"


# ============================================================
# TestExtractionPipelineEntities
# ============================================================


class TestExtractionPipelineEntities:
    """Test extraction pipeline entity support."""

    def test_pipeline_returns_entities_and_relations(self):
        ai = KGExtractionMockAIClient(dim=8)
        pipeline = ExtractionPipeline(ai_client=ai)
        result = pipeline.run(
            messages=[
                {"role": "user", "content": "Alice works at Acme Corp in Tokyo"},
                {"role": "assistant", "content": "Got it, noted."},
            ],
            user_id="test",
            extract_entities=True,
        )
        assert "entities" in result
        assert "relations" in result
        assert len(result["entities"]) >= 1
        assert len(result["relations"]) >= 1

    def test_pipeline_entities_have_name_and_type(self):
        ai = KGExtractionMockAIClient(dim=8)
        pipeline = ExtractionPipeline(ai_client=ai)
        result = pipeline.run(
            messages=[{"role": "user", "content": "Test content"}],
            user_id="test",
            extract_entities=True,
        )
        for ent in result.get("entities", []):
            assert "name" in ent
            assert "type" in ent

    def test_pipeline_relations_have_triple(self):
        ai = KGExtractionMockAIClient(dim=8)
        pipeline = ExtractionPipeline(ai_client=ai)
        result = pipeline.run(
            messages=[{"role": "user", "content": "Test content"}],
            user_id="test",
            extract_entities=True,
        )
        for rel in result.get("relations", []):
            assert "subject" in rel
            assert "predicate" in rel
            assert "object" in rel

    def test_pipeline_without_entities(self):
        ai = SmallMockAIClient(dim=8)
        pipeline = ExtractionPipeline(ai_client=ai)
        result = pipeline.run(
            messages=[{"role": "user", "content": "Hello"}],
            user_id="test",
            extract_entities=False,
        )
        assert result["entities"] == []
        assert result["relations"] == []

    def test_unified_extraction_with_entities_prompt_exists(self):
        assert "{conversation}" in UNIFIED_EXTRACTION_WITH_ENTITIES_PROMPT
        assert "entities" in UNIFIED_EXTRACTION_WITH_ENTITIES_PROMPT.lower()
        assert "relations" in UNIFIED_EXTRACTION_WITH_ENTITIES_PROMPT.lower()


# ============================================================
# TestEntityZTM
# ============================================================


class TestEntityZTM:
    """ZTM integration tests for knowledge graph."""

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

    def test_initial_confidence_medium(self):
        confidence = determine_entity_confidence(1)
        assert confidence == "medium"

    def test_model_entity_types_match(self):
        assert MODEL_ENTITY_TYPES == VALID_ENTITY_TYPES

    def test_valid_edge_directions(self):
        assert VALID_EDGE_DIRECTIONS == {"outgoing", "incoming", "both"}


# ============================================================
# TestSQLiteStoreEntities
# ============================================================


class TestSQLiteStoreEntities:
    """SQLiteStore CRUD tests for entity_nodes and entity_edges."""

    @pytest.fixture
    def store(self):
        s = SQLiteStore(":memory:")
        yield s
        s.close()

    def _ensure_nodes(self, store, *node_specs):
        """Create entity nodes for FK constraints. node_specs: (id, name, user_id)"""
        for eid, name, uid in node_specs:
            store.put_entity(build_entity_node(name, "person", uid))
            # Override entity_id if it doesn't match the deterministic one
            existing = store.get_entity_by_name(name.lower(), uid)
            if existing and existing["entity_id"] != eid:
                # Put with explicit ID
                node = build_entity_node(name, "person", uid)
                node["entity_id"] = eid
                store.put_entity(node)

    def _ensure_edge_nodes(self, store, subject_id, subject_name, object_id, object_name, user_id="agent1"):
        """Ensure both subject and object entity nodes exist."""
        for eid, name in [(subject_id, subject_name), (object_id, object_name)]:
            node = build_entity_node(name, "person", user_id)
            node["entity_id"] = eid
            store.put_entity(node)

    def test_put_get_entity(self, store):
        ent = build_entity_node("Alice", "person", "agent1")
        store.put_entity(ent)
        fetched = store.get_entity(ent["entity_id"])
        assert fetched is not None
        assert fetched["name"] == "Alice"
        assert fetched["entity_type"] == "person"
        assert fetched["is_active"] is True

    def test_get_entity_not_found(self, store):
        assert store.get_entity("nonexistent") is None

    def test_get_entity_by_name(self, store):
        ent = build_entity_node("Alice", "person", "agent1")
        store.put_entity(ent)
        fetched = store.get_entity_by_name("alice", "agent1")
        assert fetched is not None
        assert fetched["entity_id"] == ent["entity_id"]

    def test_get_entity_by_name_case_insensitive(self, store):
        ent = build_entity_node("Alice", "person", "agent1")
        store.put_entity(ent)
        fetched = store.get_entity_by_name("ALICE", "agent1")
        assert fetched is not None

    def test_delete_entity(self, store):
        ent = build_entity_node("Alice", "person", "agent1")
        store.put_entity(ent)
        result = store.delete_entity(ent["entity_id"])
        assert result is True
        fetched = store.get_entity(ent["entity_id"])
        assert fetched["is_active"] is False

    def test_put_get_entity_edge(self, store):
        self._ensure_edge_nodes(store, "ent_a", "Alice", "ent_b", "Acme")
        edge = build_entity_edge("ent_a", "Alice", "works_at", "ent_b", "Acme", "agent1")
        store.put_entity_edge(edge)
        fetched = store.get_entity_edge(edge["edge_id"])
        assert fetched is not None
        assert fetched["predicate"] == "works_at"
        assert fetched["subject_name"] == "Alice"
        assert fetched["object_name"] == "Acme"

    def test_delete_entity_edge(self, store):
        self._ensure_edge_nodes(store, "ent_a", "Alice", "ent_b", "Acme")
        edge = build_entity_edge("ent_a", "Alice", "works_at", "ent_b", "Acme", "agent1")
        store.put_entity_edge(edge)
        result = store.delete_entity_edge(edge["edge_id"])
        assert result is True
        fetched = store.get_entity_edge(edge["edge_id"])
        assert fetched["is_active"] is False

    def test_search_entity_edges_both(self, store):
        self._ensure_edge_nodes(store, "ent_a", "Alice", "ent_b", "Acme")
        self._ensure_edge_nodes(store, "ent_c", "Bob", "ent_a", "Alice")
        edge1 = build_entity_edge("ent_a", "Alice", "works_at", "ent_b", "Acme", "agent1")
        edge2 = build_entity_edge("ent_c", "Bob", "knows", "ent_a", "Alice", "agent1")
        store.put_entity_edge(edge1)
        store.put_entity_edge(edge2)
        results = store.search_entity_edges("agent1", "Alice", direction="both")
        assert len(results) == 2

    def test_search_entity_edges_outgoing(self, store):
        self._ensure_edge_nodes(store, "ent_a", "Alice", "ent_b", "Acme")
        self._ensure_edge_nodes(store, "ent_c", "Bob", "ent_a", "Alice")
        edge1 = build_entity_edge("ent_a", "Alice", "works_at", "ent_b", "Acme", "agent1")
        edge2 = build_entity_edge("ent_c", "Bob", "knows", "ent_a", "Alice", "agent1")
        store.put_entity_edge(edge1)
        store.put_entity_edge(edge2)
        results = store.search_entity_edges("agent1", "Alice", direction="outgoing")
        assert len(results) == 1
        assert results[0]["predicate"] == "works_at"

    def test_search_entity_edges_with_predicate_filter(self, store):
        self._ensure_edge_nodes(store, "ent_a", "Alice", "ent_b", "Acme")
        self._ensure_edge_nodes(store, "ent_a", "Alice", "ent_c", "Tokyo")
        edge1 = build_entity_edge("ent_a", "Alice", "works_at", "ent_b", "Acme", "agent1")
        edge2 = build_entity_edge("ent_a", "Alice", "lives_in", "ent_c", "Tokyo", "agent1")
        store.put_entity_edge(edge1)
        store.put_entity_edge(edge2)
        results = store.search_entity_edges("agent1", "Alice", predicate="works_at")
        assert len(results) == 1
        assert results[0]["predicate"] == "works_at"

    def test_get_entity_stats(self, store):
        ent1 = build_entity_node("Alice", "person", "agent1")
        ent2 = build_entity_node("Acme", "organization", "agent1")
        edge = build_entity_edge(ent1["entity_id"], "Alice", "works_at", ent2["entity_id"], "Acme", "agent1")
        store.put_entity(ent1)
        store.put_entity(ent2)
        store.put_entity_edge(edge)
        stats = store.get_entity_stats("agent1")
        assert stats["total_entities"] == 2
        assert stats["total_edges"] == 1
        assert "confidence_distribution" in stats
        assert "entity_type_distribution" in stats

    def test_clear_all_includes_entities(self, store):
        ent = build_entity_node("Alice", "person", "agent1")
        store.put_entity(ent)
        self._ensure_edge_nodes(store, "ent_a", "Alice", "ent_b", "Acme")
        edge = build_entity_edge("ent_a", "Alice", "works_at", "ent_b", "Acme", "agent1")
        store.put_entity_edge(edge)
        store.clear_all()
        assert store.get_entity(ent["entity_id"]) is None
        assert store.get_entity_edge(edge["edge_id"]) is None


# ============================================================
# TestEntityCRUDPersistent
# ============================================================


class TestEntityCRUDPersistent:
    """PersistentHandler KG method tests."""

    @pytest.fixture
    def handler(self):
        ai = SmallMockAIClient(dim=8)
        h = PersistentHandler(database_url=":memory:", ai_client=ai)
        return h

    def test_add_entity_new(self, handler):
        result = handler.add_entity("Alice", "person", user_id="agent1")
        assert result["entity_id"].startswith("ent_")
        assert result["is_new"] is True

    def test_add_entity_merge(self, handler):
        r1 = handler.add_entity("Alice", "person", user_id="agent1")
        r2 = handler.add_entity("Alice", "person", user_id="agent1")
        assert r2["is_new"] is False
        assert r2["mention_count"] == 2

    def test_add_entity_edge(self, handler):
        result = handler.add_entity_edge("Alice", "person", "works_at", "Acme", "organization", user_id="agent1")
        assert result["edge_id"].startswith("edge_")
        assert result["is_new"] is True

    def test_search_entity_relations(self, handler):
        handler.add_entity_edge("Alice", "person", "works_at", "Acme", "organization", user_id="agent1")
        results = handler.search_entity_relations("Alice", user_id="agent1")
        assert len(results) >= 1

    def test_get_entity_graph(self, handler):
        handler.add_entity_edge("Alice", "person", "works_at", "Acme", "organization", user_id="agent1")
        graph = handler.get_entity_graph("Alice", hops=1, user_id="agent1")
        assert graph["center"] == "Alice"
        assert graph["node_count"] >= 1

    def test_delete_entity(self, handler):
        result = handler.add_entity("Alice", "person", user_id="agent1")
        assert handler.delete_entity(result["entity_id"]) is True

    def test_delete_entity_edge(self, handler):
        result = handler.add_entity_edge("Alice", "person", "works_at", "Acme", "organization", user_id="agent1")
        assert handler.delete_entity_edge(result["edge_id"]) is True

    def test_get_entity_stats(self, handler):
        handler.add_entity("Alice", "person", user_id="agent1")
        handler.add_entity_edge("Alice", "person", "works_at", "Acme", "organization", user_id="agent1")
        stats = handler.get_entity_stats(user_id="agent1")
        assert stats["total_entities"] >= 2
        assert stats["total_edges"] >= 1

    def test_reset_clears_entities(self, handler):
        handler.add_entity("Alice", "person", user_id="agent1")
        handler.reset_stores()
        stats = handler.get_entity_stats(user_id="agent1")
        assert stats["total_entities"] == 0


# ============================================================
# TestIngestConversationWithEntities
# ============================================================


class TestRememberWithEntities:
    """Test remember() entity/relation extraction."""

    @pytest.fixture
    def handler(self):
        ai = KGExtractionMockAIClient(dim=8)
        h = PersistentHandler(database_url=":memory:", ai_client=ai)
        return h

    def test_remember_creates_entities(self, handler):
        result = handler.remember(
            user_id="agent1",
            messages=[
                {"role": "user", "content": "Alice works at Acme Corp in Tokyo"},
                {"role": "assistant", "content": "Got it."},
            ],
        )
        assert len(result.get("entities", [])) >= 1
        assert result["stats"]["entities_added"] >= 1

    def test_remember_creates_relations(self, handler):
        result = handler.remember(
            user_id="agent1",
            messages=[
                {"role": "user", "content": "Alice works at Acme Corp in Tokyo"},
                {"role": "assistant", "content": "Got it."},
            ],
        )
        assert len(result.get("relations", [])) >= 1
        assert result["stats"]["relations_added"] >= 1

    def test_remember_dry_run_no_entities(self, handler):
        result = handler.remember(
            user_id="agent1",
            messages=[
                {"role": "user", "content": "Alice works at Acme Corp"},
                {"role": "assistant", "content": "Got it."},
            ],
            dry_run=True,
        )
        assert len(result.get("entities", [])) == 0
        assert len(result.get("relations", [])) == 0


# ============================================================
# TestMCPToolsKnowledgeGraph
# ============================================================


