#!/usr/bin/env python3
"""Integration test for neo-cortex evolution features.

Runs the FULL system — real embeddings (Jina), real classification (Groq/Gemini),
real ChromaDB, real SQLite. Zero mocks.

Uses REAL conversation data from cortex_db/conversation_log.db as source.
Builds a fresh isolated DB, ingests real turns, then verifies all features.

Usage:
    cd ~/neo-ram/neo-cortex
    CORTEX_DATA_DIR=~/neo-ram uv run python tests/integration/run.py
"""

import asyncio
import json
import os
import shutil
import sqlite3
import sys
import tempfile
import time

# Use development code, not installed
sys.path.insert(0, os.path.join(os.path.dirname(__file__), "..", "..", "src"))


def load_keys_from_mcp():
    """Load API keys from .mcp.json if not already in environment."""
    mcp_path = os.path.expanduser("~/neo-ram/.mcp.json")
    if not os.path.exists(mcp_path):
        return
    with open(mcp_path) as f:
        data = json.load(f)
    env = data.get("mcpServers", {}).get("neo-cortex", {}).get("env", {})
    for key in ("JINA_API_KEY", "GROQ_API_KEY", "GEMINI_API_KEY"):
        if key not in os.environ and key in env:
            os.environ[key] = env[key]


def load_real_turns(conv_db_path: str, n_sessions: int = 5) -> list[dict]:
    """Load real conversation turns from the conversation_log DB.

    Picks N diverse sessions and extracts user/assistant pairs.
    Returns list of {session_id, question, answer, model, turn_number}.
    """
    conn = sqlite3.connect(conv_db_path)
    conn.row_factory = sqlite3.Row

    # Get distinct sessions with enough content
    sessions = conn.execute("""
        SELECT session_id, COUNT(*) as cnt
        FROM conversation_log
        WHERE role IN ('user', 'assistant')
        AND length(content) > 50
        GROUP BY session_id
        HAVING cnt >= 4
        ORDER BY MAX(timestamp) DESC
        LIMIT ?
    """, (n_sessions,)).fetchall()

    turns = []
    for sess in sessions:
        sid = sess["session_id"]
        rows = conn.execute("""
            SELECT role, content, model, timestamp
            FROM conversation_log
            WHERE session_id = ? AND role IN ('user', 'assistant')
            AND length(content) > 50
            ORDER BY timestamp
        """, (sid,)).fetchall()

        # Pair user+assistant into turns
        turn_num = 0
        i = 0
        while i < len(rows) - 1:
            if rows[i]["role"] == "user" and rows[i + 1]["role"] == "assistant":
                question = rows[i]["content"][:1000]
                answer = rows[i + 1]["content"][:2000]
                model = rows[i + 1]["model"] or "opus"

                turns.append({
                    "session_id": sid,
                    "question": question,
                    "answer": answer,
                    "model": model,
                    "turn_number": turn_num,
                    "timestamp": rows[i]["timestamp"],
                })
                turn_num += 1
                i += 2
            else:
                i += 1

            # Max 3 turns per session
            if turn_num >= 3:
                break

    conn.close()
    return turns


class IntegrationTest:
    """Full system integration test with real APIs and real data."""

    def __init__(self):
        self.tmpdir = tempfile.mkdtemp(prefix="cortex_integ_")
        self.db_path = os.path.join(self.tmpdir, "chroma")
        self.index_path = os.path.join(self.tmpdir, "index.db")
        self.results: list[tuple[str, bool, str]] = []
        self.cortex = None
        self.index = None
        self.store = None
        self.graph = None
        self.embedder = None
        self.memory_ids: list[str] = []

    async def setup(self):
        """Create a full CortexService with real components."""
        from neo_cortex.classifier import Classifier
        from neo_cortex.cortex import CortexService
        from neo_cortex.dedup import DedupStore
        from neo_cortex.embedder import JinaEmbedder
        from neo_cortex.graph import ConceptGraph
        from neo_cortex.llm import create_client
        from neo_cortex.memory_index import MemoryIndex
        from neo_cortex.store import MemoryStore

        print(f"  temp dir: {self.tmpdir}")
        self.store = MemoryStore(self.db_path, "integration_test")
        print(f"  ChromaDB OK")

        jina_key = os.environ["JINA_API_KEY"]
        self.embedder = JinaEmbedder(jina_key)
        print(f"  JinaEmbedder OK")

        llm = create_client()
        self.classifier = Classifier(llm)
        print(f"  Classifier OK ({llm.name})")

        self.index = MemoryIndex(self.index_path)
        rd = self.index.register_runtime_day()
        print(f"  MemoryIndex OK, runtime_day={rd}")

        dedup = DedupStore(self.index.conn)
        self.graph = ConceptGraph(self.index.conn)
        print(f"  DedupStore + ConceptGraph OK")

        self.cortex = CortexService(
            self.store, self.embedder, self.classifier,
            index=self.index, dedup=dedup, graph=self.graph,
        )
        print(f"  CortexService READY")

    async def ingest_real_data(self, turns: list[dict]):
        """Ingest real conversation turns into the test system."""
        from neo_cortex.models import IngestRequest

        for t in turns:
            req = IngestRequest(
                session_id=t["session_id"],
                question=t["question"],
                answer=t["answer"],
                model=t["model"],
                turn_number=t["turn_number"],
            )
            try:
                mid = await self.cortex.ingest(req)
                if mid:
                    self.memory_ids.append(mid)
                    print(f"  + {mid}")
                else:
                    print(f"  - skipped (noise/dedup)")
            except Exception as e:
                print(f"  ! error: {e}")

    def manipulate_time(self):
        """Set different ages on memories to test time-dependent features.

        Uses BOTH calendar timestamps (for display) and runtime days (for decay).
        - First 1/3: 10 cal days old, rd_offset=10 (spontaneous candidates)
        - Middle 1/3: 30 cal days old, rd_offset=25 (heavy decay with slowdown=10)
        - Last 1/3: 1 cal day old, rd_offset=0 (recent, minimal decay)

        Injects synthetic runtime_days rows to simulate 30 past execution days.
        """
        import datetime
        now = time.time()

        # Inject 30 synthetic runtime days so rd_offset math works
        today = datetime.date.today()
        for i in range(30):
            d = (today - datetime.timedelta(days=i)).isoformat()
            self.index._conn.execute(
                "INSERT OR IGNORE INTO runtime_days (day) VALUES (?)", (d,)
            )
        self.index._conn.commit()
        current_rd = self.index.get_runtime_day_count()
        print(f"  runtime days after injection: {current_rd}")

        n = len(self.memory_ids)
        third = max(1, n // 3)

        groups = [
            (self.memory_ids[:third], 10, 10, "10d old, rd_offset=10 — spontaneous"),
            (self.memory_ids[third:third*2], 30, 25, "30d old, rd_offset=25 — heavy decay"),
            (self.memory_ids[third*2:], 1, 0, "1d old, rd_offset=0 — recent"),
        ]

        for ids, cal_days, rd_offset, label in groups:
            ts = now - 86400 * cal_days
            created_rd = max(1, current_rd - rd_offset)
            for mid in ids:
                self.index._conn.execute(
                    "UPDATE memories SET timestamp = ?, last_accessed = NULL, access_count = 0, "
                    "created_at_rd = ?, last_accessed_rd = ? WHERE id = ?",
                    (ts, created_rd, created_rd, mid),
                )
                self.store._collection.update(ids=[mid], metadatas=[{"timestamp": ts}])
            print(f"  {len(ids)} memories: {label} (created_rd={created_rd})")
        self.index._conn.commit()

    def record(self, name: str, passed: bool, detail: str = ""):
        """Record a test result."""
        status = "\033[32mPASS\033[0m" if passed else "\033[31mFAIL\033[0m"
        self.results.append((name, passed, detail))
        print(f"  [{status}] {name}")
        if detail:
            print(f"         {detail}")

    # ==================== TESTS ====================

    async def test_recall_finds_relevant(self):
        """Recall returns semantically relevant memories."""
        print()
        result = await self.cortex.recall("memory system architecture", n=3, smart=True)

        self.record(
            "recall returns results",
            result.count > 0,
            f"count={result.count}",
        )
        if result.count > 0:
            best = result.memories[0]
            self.record(
                "top result has good similarity",
                best.similarity > 0.2,
                f"sim={best.similarity:.3f}, topic={best.record.topic}",
            )
            self.record(
                "MBEL output generated",
                bool(result.mbel),
                f"mbel length={len(result.mbel) if result.mbel else 0} chars",
            )

    async def test_stability_increases(self):
        """Recalling a memory increases its stability."""
        print()
        if not self.memory_ids:
            self.record("stability: has data", False, "no memories")
            return

        mid = self.memory_ids[0]
        stability_before = self.index.get_stability(mid)

        # Do a recall that will match this memory
        record = self.index.get_by_id(mid)
        if record:
            await self.cortex.recall(record.question[:50], n=1, smart=False)

        stability_after = self.index.get_stability(mid)
        self.record(
            "stability increases after recall",
            stability_after > stability_before,
            f"before={stability_before:.2f}, after={stability_after:.2f}",
        )

    async def test_rrf_hybrid_search(self):
        """RRF merges vector and FTS5 results."""
        print()
        # Pick a keyword from a known memory to trigger FTS match
        if not self.memory_ids:
            self.record("RRF: has data", False, "no memories")
            return

        record = self.index.get_by_id(self.memory_ids[-1])
        if not record:
            self.record("RRF: record found", False, "record not in index")
            return

        # Use a few words from the question to trigger FTS
        keywords = record.question.split()[:3]
        query = " ".join(keywords)

        result = await self.cortex.recall(query, n=5, smart=True)
        self.record(
            "RRF returns results for keyword query",
            result.count > 0,
            f"query='{query}', count={result.count}",
        )

        # Verify the target memory appears in results
        target_found = any(m.record.id == record.id for m in result.memories)
        self.record(
            "keyword-queried memory found in results",
            target_found,
            f"target={record.id}, found in top {result.count}",
        )

    async def test_dream_cycle(self):
        """Dream applies runtime-day decay and returns surprisal IDs."""
        print()
        current_rd = self.index.get_runtime_day_count()
        dream = await self.cortex.dream()

        self.record(
            "dream completes",
            dream.status == "completed",
            f"status={dream.status}, dream_count={dream.dream_count}",
        )
        self.record(
            "dream returns surprisal field",
            hasattr(dream, "surprisal") and isinstance(dream.surprisal, list),
            f"surprisal={dream.surprisal[:5]}",
        )
        self.record(
            "surprisal identifies novel memories",
            len(dream.surprisal) > 0,
            f"count={len(dream.surprisal)}",
        )

        # Check runtime-day-based decay
        all_data = self.index.get_all_for_dream()

        # rd_offset=25 memories: execution_days=25, slowdown=10
        # exp(-25/10/1) = exp(-2.5) ≈ 0.08 → near floor
        heavy_decay = [d for d in all_data
                       if (current_rd - d.get("last_accessed_rd", 0)) >= 20]
        if heavy_decay:
            low_energy = [d for d in heavy_decay if d["energy"] < 0.2]
            self.record(
                "heavy-decay memories (rd_offset=25) lost energy",
                len(low_energy) > 0,
                f"{len(low_energy)}/{len(heavy_decay)} have energy<0.2",
            )

        # rd_offset=0 memories: execution_days=0 → energy unchanged (1.0)
        recent = [d for d in all_data
                  if (current_rd - d.get("last_accessed_rd", 0)) == 0]
        if recent:
            high_energy = [d for d in recent if d["energy"] > 0.9]
            self.record(
                "recent memories (rd_offset=0) kept energy",
                len(high_energy) == len(recent),
                f"{len(high_energy)}/{len(recent)} have energy>0.9",
            )

    async def test_spontaneous_recall(self):
        """Spontaneous surfaces stale high-energy memories."""
        print()
        # Boost energy of old memories so they qualify
        all_data = self.index.get_all_for_dream()
        old_stale = [d for d in all_data
                     if (time.time() - d["timestamp"]) / 86400 > 7
                     and (d.get("access_count") or 0) == 0]

        for d in old_stale:
            self.store.update_energy(d["id"], 0.9)
            self.index.update_energy(d["id"], 0.9)

        spontaneous = await self.cortex.spontaneous(n=3)
        self.record(
            "spontaneous returns memories",
            len(spontaneous) > 0,
            f"count={len(spontaneous)}, ids={[m.id for m in spontaneous[:3]]}",
        )

    async def test_observation_levels(self):
        """Observation levels persist in SQLite and ChromaDB."""
        print()
        from neo_cortex.models import MemoryRecord, ObservationLevel

        record = MemoryRecord(
            id="integ_obs_test", session_id="integ-obs",
            question="Is this a reported fact?",
            answer_preview="According to external source, yes.",
            document="Q: Is this reported?\nA: According to external source, yes.",
            observation_level=ObservationLevel.REPORTED,
        )

        embedding = await self.embedder.embed_passage(record.document)
        self.store.insert(record, embedding)
        self.index.insert(record)

        # Check SQLite
        rows = self.index.get_by_ids(["integ_obs_test"])
        self.record(
            "observation_level in SQLite",
            len(rows) == 1,
            f"stored and retrievable",
        )

        # Check ChromaDB
        chroma = self.store._collection.get(ids=["integ_obs_test"], include=["metadatas"])
        meta = chroma["metadatas"][0] if chroma["metadatas"] else {}
        self.record(
            "observation_level in ChromaDB",
            meta.get("observation_level") == "reported",
            f"value={meta.get('observation_level')}",
        )

    async def test_runtime_day_decay(self):
        """Runtime day tracking: decay uses execution days, not calendar days."""
        print()
        current_rd = self.index.get_runtime_day_count()
        self.record(
            "runtime day registered",
            current_rd >= 1,
            f"current_rd={current_rd}",
        )

        # Verify memories have runtime day columns populated
        all_data = self.index.get_all_for_dream()
        has_rd = [d for d in all_data if d.get("created_at_rd", 0) > 0]
        self.record(
            "memories have runtime day metadata",
            len(has_rd) > 0,
            f"{len(has_rd)}/{len(all_data)} have created_at_rd > 0",
        )

        # Verify decay formula uses execution days with slowdown
        from neo_cortex.decay import ebbinghaus_weight
        # 10 execution days, stability=1, slowdown=10 → exp(-1) ≈ 0.37
        w = ebbinghaus_weight(10, stability=1.0, slowdown=10)
        self.record(
            "decay formula with slowdown=10",
            0.35 < w < 0.40,
            f"ebbinghaus(exec=10, stab=1, slow=10) = {w:.4f} (expected ~0.37)",
        )

        # 0 execution days → no decay
        w0 = ebbinghaus_weight(0, stability=1.0, slowdown=10)
        self.record(
            "zero execution days = no decay",
            w0 == 1.0,
            f"ebbinghaus(exec=0) = {w0:.4f}",
        )

    async def test_graph_boost_in_recall(self):
        """Graph boost improves basic recall results."""
        print()
        if len(self.memory_ids) < 2:
            self.record("graph boost: has data", False, "need at least 2 memories")
            return

        # Add concepts to graph
        self.graph.add_memory(self.memory_ids[0], ["test-concept-alpha", "memory-system"])
        self.graph.add_memory(self.memory_ids[1], ["memory-system", "test-concept-beta"])

        self.record(
            "graph has nodes",
            self.graph.node_count() > 0,
            f"nodes={self.graph.node_count()}",
        )

        # Basic recall (smart=False) uses graph boost now
        result = await self.cortex.recall("memory system", n=3, smart=False)
        self.record(
            "basic recall with graph returns results",
            result.count > 0,
            f"count={result.count}",
        )

    async def test_metacognitive_reflection(self):
        """meta_reflect generates real metacognitive insights from real memories."""
        print()
        from neo_cortex.reflection import meta_reflect
        from neo_cortex.llm import create_client

        llm = create_client()

        # Needs memories in index to reflect on
        if len(self.memory_ids) < 3:
            self.record("meta_reflect: has data", False, "need at least 3 memories")
            return

        count = await meta_reflect(self.cortex, llm)
        self.record(
            "meta_reflect creates insights",
            count >= 0,
            f"created {count} metacognitive memories",
        )

        # Verify metacognitive memories exist in SQLite
        rows = self.index._conn.execute(
            "SELECT id, summary, observation_level, evidence_ids "
            "FROM memories WHERE observation_level = 'metacognitive'"
        ).fetchall()
        self.record(
            "metacognitive memories stored in SQLite",
            len(rows) >= count,
            f"found {len(rows)} metacognitive records",
        )

        if rows:
            # Verify evidence_ids are populated
            has_evidence = [r for r in rows if r["evidence_ids"]]
            self.record(
                "metacognitive memories have evidence_ids",
                len(has_evidence) > 0,
                f"{len(has_evidence)}/{len(rows)} have evidence_ids",
            )

            # Verify they exist in ChromaDB too
            meta_ids = [r["id"] for r in rows]
            chroma = self.store._collection.get(ids=meta_ids[:3], include=["metadatas"])
            chroma_meta = [m.get("observation_level") for m in (chroma["metadatas"] or [])]
            self.record(
                "metacognitive memories in ChromaDB",
                all(ol == "metacognitive" for ol in chroma_meta),
                f"obs_levels={chroma_meta}",
            )

    async def test_self_model_in_timeline(self):
        """Timeline CLI builds Self-Model section from metacognitive memories."""
        print()
        from neo_cortex.timeline_cli import _get_metacognitive_memories, _build_self_model_section

        meta_mems = _get_metacognitive_memories(self.index._conn, limit=10)
        section = _build_self_model_section(meta_mems)

        self.record(
            "timeline fetches metacognitive memories",
            isinstance(meta_mems, list),
            f"found {len(meta_mems)} metacognitive memories",
        )

        if meta_mems:
            self.record(
                "Self-Model section generated",
                "## [Self-Model]" in section,
                f"section length={len(section)} chars",
            )
            self.record(
                "Self-Model has @learned:: lines",
                "@learned::" in section,
                f"lines={len([l for l in section.splitlines() if '@learned::' in l])}",
            )
        else:
            self.record(
                "Self-Model empty (no metacognitive yet)",
                section == "",
                "expected empty — meta_reflect may have produced 0 insights",
            )

    async def test_identity_review_mcp_tool(self):
        """memory_identity_review MCP tool runs without crash."""
        print()
        from neo_cortex import mcp as mcp_module

        original = mcp_module._cortex
        try:
            mcp_module._cortex = self.cortex
            result = await mcp_module.memory_identity_review.fn()

            is_dict = isinstance(result, dict)
            has_status = is_dict and "status" in result
            self.record(
                "memory_identity_review returns valid response",
                has_status,
                f"status={result.get('status') if is_dict else type(result).__name__}",
            )

            if is_dict and result.get("status") == "proposals_ready":
                self.record(
                    "identity review has proposals",
                    result.get("count", 0) > 0,
                    f"count={result.get('count')}, proposals={len(result.get('proposals', []))}",
                )
        finally:
            mcp_module._cortex = original

    async def test_spontaneous_mcp_tool(self):
        """The memory_spontaneous MCP tool works end-to-end."""
        print()
        from neo_cortex import mcp as mcp_module

        original = mcp_module._cortex
        try:
            mcp_module._cortex = self.cortex
            result = await mcp_module.memory_spontaneous.fn(n=3, detail="compact")

            is_str = isinstance(result, str)
            is_empty_dict = isinstance(result, dict) and result.get("count") == 0

            self.record(
                "memory_spontaneous tool callable",
                is_str or is_empty_dict,
                f"type={type(result).__name__}, "
                + (f"len={len(result)}" if is_str else f"result={result}"),
            )
        finally:
            mcp_module._cortex = original

    # ==================== RUNNER ====================

    async def run_all(self, conv_db_path: str):
        """Run all integration tests."""
        print("=" * 60)
        print("  NEO-CORTEX INTEGRATION TEST")
        print("  Full system — real APIs, real data, zero mocks")
        print("=" * 60)

        # Phase 1: Setup
        print("\n--- SETUP ---")
        await self.setup()

        # Phase 2: Load real data
        print("\n--- LOAD REAL CONVERSATIONS ---")
        turns = load_real_turns(conv_db_path, n_sessions=5)
        print(f"  loaded {len(turns)} turns from {conv_db_path}")

        if not turns:
            print("  ABORT: no conversation data found.")
            return False

        # Phase 3: Ingest
        print(f"\n--- INGEST ({len(turns)} turns) ---")
        await self.ingest_real_data(turns)
        print(f"  stored: {len(self.memory_ids)} memories")

        if len(self.memory_ids) < 3:
            print("  ABORT: too few memories ingested.")
            return False

        # Phase 4: Manipulate time
        print(f"\n--- TIME MANIPULATION ---")
        self.manipulate_time()

        # Phase 5: Run tests
        print("\n--- TEST: Basic Recall ---")
        await self.test_recall_finds_relevant()

        print("\n--- TEST: Stability Reinforcement ---")
        await self.test_stability_increases()

        print("\n--- TEST: RRF Hybrid Search ---")
        await self.test_rrf_hybrid_search()

        print("\n--- TEST: Runtime Day Decay ---")
        await self.test_runtime_day_decay()

        print("\n--- TEST: Dream + Surprisal ---")
        await self.test_dream_cycle()

        print("\n--- TEST: Spontaneous Recall ---")
        await self.test_spontaneous_recall()

        print("\n--- TEST: Observation Levels ---")
        await self.test_observation_levels()

        print("\n--- TEST: Graph Boost ---")
        await self.test_graph_boost_in_recall()

        print("\n--- TEST: Meta-Cognitive Reflection ---")
        await self.test_metacognitive_reflection()

        print("\n--- TEST: Self-Model in Timeline ---")
        await self.test_self_model_in_timeline()

        print("\n--- TEST: Identity Review MCP ---")
        await self.test_identity_review_mcp_tool()

        print("\n--- TEST: MCP Tool ---")
        await self.test_spontaneous_mcp_tool()

        # Summary
        print("\n" + "=" * 60)
        passed = sum(1 for _, p, _ in self.results if p)
        failed = sum(1 for _, p, _ in self.results if not p)
        total = len(self.results)

        if failed == 0:
            print(f"  \033[32mALL PASSED: {passed}/{total}\033[0m")
        else:
            print(f"  \033[31m{failed} FAILED\033[0m, {passed} passed (total {total})")
            print("\n  Failed tests:")
            for name, p, detail in self.results:
                if not p:
                    print(f"    - {name}: {detail}")

        print("=" * 60)
        return failed == 0

    def cleanup(self):
        """Remove temp directory."""
        if os.path.exists(self.tmpdir):
            shutil.rmtree(self.tmpdir, ignore_errors=True)
            print(f"\n[cleanup] removed {self.tmpdir}")


async def main():
    load_keys_from_mcp()

    # Verify keys
    for key in ("JINA_API_KEY", "GEMINI_API_KEY"):
        if key not in os.environ:
            print(f"ERROR: {key} not set. Set CORTEX_DATA_DIR=~/neo-ram or export keys manually.")
            sys.exit(1)

    # Find conversation_log.db
    data_dir = os.environ.get("CORTEX_DATA_DIR", os.path.expanduser("~/neo-ram"))
    conv_db = os.path.join(data_dir, "cortex_db", "conversation_log.db")
    if not os.path.exists(conv_db):
        print(f"ERROR: conversation_log.db not found at {conv_db}")
        sys.exit(1)

    test = IntegrationTest()
    try:
        success = await test.run_all(conv_db)
        sys.exit(0 if success else 1)
    except KeyboardInterrupt:
        print("\nInterrupted.")
        sys.exit(130)
    except Exception as e:
        print(f"\nFATAL: {e}")
        import traceback
        traceback.print_exc()
        sys.exit(2)
    finally:
        test.cleanup()


if __name__ == "__main__":
    asyncio.run(main())
