"""
tests/test_integration.py — Layer 4: Full end-to-end integration tests.

Uses REAL SQLite storage (:memory:), REAL config (defaults.toml).
Mocks ONLY external LLM API calls.
"""
from __future__ import annotations

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

# LLM provider types for realistic mocks
from cortex.llm.provider import (
    CompletionResponse,
    EmbedResponse,
    ExtractionResponse,
)

# ---------------------------------------------------------------------------
# Mock LLM Provider
# ---------------------------------------------------------------------------

class MockLLMProvider:
    """
    Mock LLM provider that returns realistic responses for integration tests.
    
    - complete(): returns CompletionResponse with JSON for extraction prompts,
                  plain text otherwise
    - embed(): returns EmbedResponse with deterministic fake embeddings
    - extract(): returns ExtractionResponse with realistic data
    """

    DEFAULT_COMPLETION_MODEL = "mock-model"
    DEFAULT_EMBED_MODEL = "mock-embed"
    _usage_log: list = []

    def __init__(self):
        self._usage_log = []
        self._complete_calls = 0
        self._embed_calls = 0

    async def complete(
        self,
        system: str = "",
        user: str = "",
        model: str | None = None,
        response_format: str = "text",
        max_tokens: int | None = None,
        temperature: float = 0.0,
        **kwargs,
    ) -> CompletionResponse:
        self._complete_calls += 1
        
        # If asking for JSON (extraction), return realistic claim extraction
        if response_format == "json" or "extract" in system.lower():
            text = json.dumps({
                "claims": [
                    {
                        "action": "ADD",
                        "subject": "Andrew",
                        "predicate": "lives_in",
                        "object_value": "Bangkok",
                        "confidence": 0.95,
                        "temporal": "current",
                        "entities_mentioned": [
                            {"name": "Andrew", "entity_type": "person", "aliases": []},
                            {"name": "Bangkok", "entity_type": "location", "aliases": []}
                        ],
                        "sensitivity": "normal"
                    },
                    {
                        "action": "ADD",
                        "subject": "Andrew",
                        "predicate": "name",
                        "object_value": "Andrew",
                        "confidence": 0.99,
                        "temporal": "permanent",
                        "entities_mentioned": [
                            {"name": "Andrew", "entity_type": "person", "aliases": []}
                        ],
                        "sensitivity": "normal"
                    }
                ]
            })
        elif "relationship" in system.lower() or "classify" in system.lower():
            # Reconciliation LLM calls
            text = json.dumps({"relationship": "NEW", "confidence": 0.9})
        else:
            text = "This is a mock LLM response."
        
        return CompletionResponse(
            text=text,
            model=model or "mock-model",
            tokens_in=100,
            tokens_out=50,
            cost_estimate=0.001,
            latency_ms=10.0,
        )

    async def extract(
        self,
        system: str = "",
        user: str = "",
        schema: dict | None = None,
        model: str | None = None,
        max_tokens: int | None = None,
        temperature: float = 0.0,
        **kwargs,
    ) -> ExtractionResponse:
        data = {
            "claims": [
                {
                    "action": "ADD",
                    "subject": "Andrew",
                    "predicate": "lives_in",
                    "object_value": "Bangkok",
                    "confidence": 0.95,
                }
            ]
        }
        return ExtractionResponse(
            data=data,
            model=model or "mock-model",
            tokens_in=80,
            tokens_out=40,
            cost_estimate=0.0005,
            latency_ms=8.0,
        )

    async def embed(
        self,
        texts: list[str],
        model: str | None = None,
        input_type: str = "query",
        **kwargs,
    ) -> EmbedResponse:
        """Return deterministic fake embeddings based on text hash."""
        self._embed_calls += 1
        embeddings = []
        for text in texts:
            # Generate a deterministic 1536-dim embedding from text hash
            import hashlib
            h = hashlib.sha256(text.encode()).digest()
            # Use first 32 bytes to seed a simple pattern
            base = [b / 255.0 for b in h]
            # Pad to 1536 dims by repeating
            embedding = (base * 48)[:1536]
            embeddings.append(embedding)
        
        return EmbedResponse(
            embeddings=embeddings,
            model=model or "mock-embed",
            tokens_in=sum(len(t.split()) for t in texts),
        )

    def usage_summary(self) -> dict:
        return {"calls": self._complete_calls + self._embed_calls}

    def _resolve_model(self, model, default):
        return model or default


# ---------------------------------------------------------------------------
# Fixtures
# ---------------------------------------------------------------------------

@pytest.fixture
def mock_llm():
    return MockLLMProvider()


@pytest.fixture
async def plugin(mock_llm):
    """
    Create a CortexPlugin with real SQLite (:memory:) and real config,
    but mock LLM provider injected post-init.
    """
    from cortex.api.plugin import CortexPlugin
    
    p = CortexPlugin(
        config_path=None,  # uses defaults
        db_path=":memory:",
    )
    
    # Inject mock LLM
    p.llm = mock_llm
    
    # Re-wire modules that depend on LLM
    if p.extractor is not None:
        p.extractor.llm = mock_llm
    if p.entity_resolver is not None:
        p.entity_resolver.llm = mock_llm
    if p.reconciler is not None:
        p.reconciler.llm = mock_llm
    if p.retrieval is not None:
        p.retrieval.llm = mock_llm
    
    # Initialize storage (create tables)
    if p.storage is not None:
        await p.storage.initialize()
    
    yield p
    
    # Cleanup
    if p.storage is not None:
        await p.storage.close()


# ---------------------------------------------------------------------------
# Test 1: Plugin Instantiation
# ---------------------------------------------------------------------------

class TestPluginInstantiation:
    """Test 1: Can CortexPlugin() actually initialize?"""

    @pytest.mark.asyncio
    async def test_plugin_creates_with_memory_db(self, plugin):
        assert plugin is not None
        assert plugin.storage is not None

    @pytest.mark.asyncio
    async def test_plugin_has_all_core_modules(self, plugin):
        assert plugin.extractor is not None, "ExtractionPipeline missing"
        assert plugin.entity_resolver is not None, "EntityResolver missing"
        assert plugin.reconciler is not None, "ReconciliationEngine missing"
        assert plugin.cornerstone_guardian is not None, "CornerstoneGuardian missing"
        assert plugin.token_budget_manager is not None, "TokenBudgetManager missing"
        assert plugin.retrieval is not None, "RetrievalPipeline missing"

    @pytest.mark.asyncio
    async def test_storage_tables_created(self, plugin):
        """Verify SQLite schema was created via initialize()."""
        conn = await plugin.storage._get_conn()
        async with conn.execute(
            "SELECT name FROM sqlite_master WHERE type='table'"
        ) as cur:
            tables = {row[0] for row in await cur.fetchall()}
        
        # Core tables from v3 schema
        expected = {"entity", "semantic_claim", "session_record", "cornerstone_memory"}
        missing = expected - tables
        assert not missing, f"Missing tables: {missing}"


# ---------------------------------------------------------------------------
# Test 2: Remember Flow
# ---------------------------------------------------------------------------

class TestRememberFlow:
    """Test 2: Conversation → extract → resolve entities → reconcile → store."""

    @pytest.mark.asyncio
    async def test_remember_basic_conversation(self, plugin):
        conversation = [
            {"role": "user", "content": "My name is Andrew and I live in Bangkok"},
            {"role": "assistant", "content": "Nice to meet you, Andrew! Bangkok is a great city."}
        ]
        
        result = await plugin.remember(conversation)
        
        assert result is not None
        assert result.ok, f"remember() failed: {result.errors}"
        assert result.claims_extracted > 0, "No claims extracted"
        assert result.claims_stored > 0, "No claims stored"
        assert result.session_id is not None

    @pytest.mark.asyncio
    async def test_remember_stores_in_db(self, plugin):
        conversation = [
            {"role": "user", "content": "My name is Andrew and I live in Bangkok"},
            {"role": "assistant", "content": "Nice to meet you!"}
        ]
        
        result = await plugin.remember(conversation)
        assert result.ok
        
        # Verify claims exist in storage
        conn = await plugin.storage._get_conn()
        async with conn.execute("SELECT COUNT(*) FROM semantic_claim") as cur:
            row = await cur.fetchone()
            count = row[0]
        
        assert count > 0, "No claims in semantic_claim table after remember()"

    @pytest.mark.asyncio
    async def test_remember_creates_entities(self, plugin):
        conversation = [
            {"role": "user", "content": "My name is Andrew and I live in Bangkok"},
            {"role": "assistant", "content": "Nice to meet you!"}
        ]
        
        result = await plugin.remember(conversation)
        assert result.ok
        assert result.entities_resolved > 0, "No entities resolved"
        
        # Verify entities in storage
        conn = await plugin.storage._get_conn()
        async with conn.execute("SELECT COUNT(*) FROM entity") as cur:
            row = await cur.fetchone()
            count = row[0]
        
        assert count > 0, "No entities in entity table after remember()"


# ---------------------------------------------------------------------------
# Test 3: Retrieve Flow
# ---------------------------------------------------------------------------

class TestRetrieveFlow:
    """Test 3: Query → hybrid search → assemble context."""

    @pytest.mark.asyncio
    async def test_retrieve_returns_result(self, plugin):
        """Retrieve on empty DB should return a result (possibly empty items)."""
        result = await plugin.retrieve("Where does Andrew live?", token_budget=1000)
        
        # Should return a RetrievalResult, not None
        assert result is not None
        # items may be empty on a fresh DB — that's OK
        assert hasattr(result, "items")

    @pytest.mark.asyncio
    async def test_retrieve_with_stored_data(self, plugin):
        """After remember(), retrieve() should find something."""
        # Store data first
        conversation = [
            {"role": "user", "content": "My name is Andrew and I live in Bangkok"},
            {"role": "assistant", "content": "Nice to meet you, Andrew!"}
        ]
        remember_result = await plugin.remember(conversation)
        assert remember_result.ok
        
        # Now retrieve — BM25 should find matches even without real embeddings
        result = await plugin.retrieve("Andrew Bangkok", token_budget=1000)
        assert result is not None
        # With BM25 search, we should get hits if FTS is working
        # (vector search may return empty with fake embeddings, but BM25 should work)


# ---------------------------------------------------------------------------
# Test 4: Remember → Retrieve Roundtrip
# ---------------------------------------------------------------------------

class TestRememberRetrieveRoundtrip:
    """Test 4: Store something, then retrieve it."""

    @pytest.mark.asyncio
    async def test_roundtrip(self, plugin):
        # 1. Remember
        conversation = [
            {"role": "user", "content": "My name is Andrew and I live in Bangkok"},
            {"role": "assistant", "content": "Nice to meet you, Andrew! Bangkok is great."}
        ]
        remember_result = await plugin.remember(conversation)
        assert remember_result.ok, f"Remember failed: {remember_result.errors}"
        assert remember_result.claims_stored > 0
        
        # 2. Retrieve
        result = await plugin.retrieve("Andrew's location", token_budget=1000)
        assert result is not None
        
        # 3. Verify — check that stored claims are in the DB
        conn = await plugin.storage._get_conn()
        async with conn.execute(
            "SELECT subject, predicate, object_value FROM semantic_claim"
        ) as cur:
            claims = await cur.fetchall()
        
        # Should have at least one claim about Andrew or Bangkok
        subjects = [c[0] for c in claims]
        objects = [c[2] for c in claims]
        all_values = [v.lower() for v in subjects + objects]
        
        assert any("andrew" in v for v in all_values) or any("bangkok" in v for v in all_values), \
            f"Expected Andrew/Bangkok in stored claims, got: {claims}"


# ---------------------------------------------------------------------------
# Test 5: Cornerstone Operations
# ---------------------------------------------------------------------------

class TestCornerstoneOperations:
    """Test 5: Seal, verify, unseal."""

    @pytest.mark.asyncio
    async def test_seal_cornerstone(self, plugin):
        result = await plugin.seal_cornerstone(
            memory_id="test-cs-1",
            label="Eva's identity",
            content="Eva is Andrew's AI companion",
        )
        assert result is not None

    @pytest.mark.asyncio
    async def test_get_cornerstones_after_seal(self, plugin):
        await plugin.seal_cornerstone(
            memory_id="test-cs-1",
            label="Eva's identity",
            content="Eva is Andrew's AI companion",
        )
        
        cornerstones = await plugin.get_cornerstones()
        assert len(cornerstones) >= 1
        
        labels = [cs.label for cs in cornerstones]
        assert "Eva's identity" in labels

    @pytest.mark.asyncio
    async def test_unseal_cornerstone(self, plugin):
        sealed = await plugin.seal_cornerstone(
            memory_id="test-cs-unseal",
            label="Temp cornerstone",
            content="Temporary data",
        )
        assert sealed is not None
        
        # Unseal it
        result = await plugin.unseal_cornerstone(sealed.id)
        assert result is not None

    @pytest.mark.asyncio
    async def test_seal_unseal_roundtrip(self, plugin):
        # Seal
        sealed = await plugin.seal_cornerstone(
            memory_id="cs-roundtrip",
            label="Roundtrip test",
            content="Testing seal/unseal cycle",
        )
        assert sealed is not None
        cs_id = sealed.id
        
        # Verify it's in the list
        cornerstones = await plugin.get_cornerstones()
        assert any(cs.id == cs_id for cs in cornerstones)
        
        # Unseal
        unsealed = await plugin.unseal_cornerstone(cs_id)
        assert unsealed is not None


# ---------------------------------------------------------------------------
# Test 6: Forget with Cornerstone Protection
# ---------------------------------------------------------------------------

class TestForgetWithCornerstoneProtection:
    """Test 6: Try to forget a cornerstone — should fail."""

    @pytest.mark.asyncio
    async def test_forget_cornerstone_rejected(self, plugin):
        # Seal a cornerstone with identifiable content
        sealed = await plugin.seal_cornerstone(
            memory_id="protected-mem",
            label="protected-mem",
            content="Critical identity memory that must not be deleted",
        )
        assert sealed is not None
        
        # Try to forget using the label as memory_id
        # (is_cornerstone_subject checks substring match on label/content)
        result = await plugin.forget("protected-mem")
        
        assert result is not None
        assert result.deleted is False, "Cornerstone should NOT be deletable"
        assert "cornerstone" in result.reason.lower() or "unseal" in result.reason.lower()

    @pytest.mark.asyncio
    async def test_forget_non_cornerstone_succeeds(self, plugin):
        """Forgetting a non-cornerstone memory should work (soft delete)."""
        # First, store a claim directly
        claim = await plugin.storage.create_claim(
            subject="test-subject",
            predicate="test-pred",
            object_value="test-value",
            owner_id=plugin.config.owner_id,
        )
        
        # Forget it
        result = await plugin.forget(claim.id)
        assert result is not None
        assert result.deleted is True, f"Forget failed: {result.reason}"

    @pytest.mark.asyncio
    async def test_forget_after_unseal_succeeds(self, plugin):
        """After unsealing, the same memory_id should be forgettable."""
        sealed = await plugin.seal_cornerstone(
            memory_id="unseal-then-forget",
            label="unseal-then-forget",
            content="Temporary cornerstone",
        )
        assert sealed is not None
        
        # Unseal it first
        await plugin.unseal_cornerstone(sealed.id)
        
        # Now forget should succeed (no longer a cornerstone)
        result = await plugin.forget("unseal-then-forget")
        # After unseal, is_cornerstone_subject should return False
        assert result.deleted is True or "not available" not in result.reason.lower()


# ---------------------------------------------------------------------------
# Test 7: Health Check
# ---------------------------------------------------------------------------

class TestHealthCheck:
    """Test 7: health() returns a valid report."""

    @pytest.mark.asyncio
    async def test_health_returns_report(self, plugin):
        health = await plugin.health()
        assert health is not None
        assert hasattr(health, "ok")
        assert hasattr(health, "modules")
        assert hasattr(health, "storage")

    @pytest.mark.asyncio
    async def test_health_storage_ok(self, plugin):
        health = await plugin.health()
        assert health.storage == "ok", f"Storage health: {health.storage}"

    @pytest.mark.asyncio
    async def test_health_modules_present(self, plugin):
        health = await plugin.health()
        # Core modules should all be "ok" since we injected mock LLM
        for module_name in ["extraction", "entity_resolution", "reconciliation",
                           "cornerstone", "token_budget", "retrieval", "llm"]:
            assert health.modules.get(module_name) == "ok", \
                f"Module {module_name} status: {health.modules.get(module_name)}"


# ---------------------------------------------------------------------------
# Test 8: Session Lifecycle (wake/sleep)
# ---------------------------------------------------------------------------

class TestSessionLifecycle:
    """Bonus: wake() and sleep() work with real storage."""

    @pytest.mark.asyncio
    async def test_wake_creates_session(self, plugin):
        report = await plugin.wake(session_id="test-session-1")
        assert report is not None
        assert report.session_id == "test-session-1"

    @pytest.mark.asyncio
    async def test_sleep_closes_session(self, plugin):
        await plugin.wake(session_id="test-session-2")
        report = await plugin.sleep(session_id="test-session-2")
        assert report is not None
        assert report.session_id == "test-session-2"

    @pytest.mark.asyncio
    async def test_wake_sleep_roundtrip(self, plugin):
        wake_report = await plugin.wake()
        sid = wake_report.session_id
        assert sid is not None
        
        sleep_report = await plugin.sleep(session_id=sid)
        assert sleep_report.session_id == sid
