"""
Tests for Phase 23.5 — Memory Infrastructure Upgrade
Covers: Observation Schema, Memory Hook System, DAG Memory Layer,
        Compression Worker, Expandable Recall Engine, Memory Engine Interface
"""

import pytest
import time
from openmemo.core.observation import Observation, VALID_TYPES
from openmemo.core.memory_graph import MemoryGraph, MemoryNode
from openmemo.core.expandable_recall import ExpandableRecallEngine
from openmemo.core.engine import MemoryEngine
from openmemo.hooks.memory_hooks import MemoryHooks
from openmemo.workers.compression_worker import CompressionWorker
from openmemo.storage.sqlite_store import SQLiteStore


@pytest.fixture
def store(tmp_path):
    return SQLiteStore(db_path=str(tmp_path / "test.db"))


@pytest.fixture
def graph(store):
    return MemoryGraph(store)


@pytest.fixture
def engine(store):
    return MemoryEngine(store=store, compress_every=0)


# ─── Observation Schema ────────────────────────────────────────────


class TestObservationSchema:
    def test_tool_execution_factory(self):
        obs = Observation.tool_execution(
            tool="docker", action="deploy backend", result="success",
            agent_id="coder", task_id="task_1", scene="deployment",
            tags=["backend", "docker"],
        )
        assert obs.type == "tool_execution"
        assert obs.agent_id == "coder"
        assert obs.task_id == "task_1"
        assert obs.content["tool"] == "docker"
        assert obs.content["result"] == "success"
        assert obs.context["scene"] == "deployment"
        assert obs.context["tags"] == ["backend", "docker"]
        assert obs.id.startswith("obs_")

    def test_agent_step_factory(self):
        obs = Observation.agent_step(
            step="analyzing logs", reasoning="checking for errors",
            agent_id="ops", task_id="task_2", scene="ops",
        )
        assert obs.type == "agent_step"
        assert obs.content["step"] == "analyzing logs"
        assert obs.content["reasoning"] == "checking for errors"

    def test_session_event_factory(self):
        obs = Observation.session_event(
            event="session_start", session_id="sess_1",
            agent_id="coder", scene="coding",
        )
        assert obs.type == "session_event"
        assert obs.content["event"] == "session_start"
        assert obs.content["session_id"] == "sess_1"

    def test_task_event_factory(self):
        obs = Observation.task_event(
            event="task_complete", task_id="task_3",
            status="success", agent_id="planner",
        )
        assert obs.type == "task_event"
        assert obs.content["status"] == "success"

    def test_unknown_type_defaults_to_observation(self):
        obs = Observation(type="invalid_xyz")
        assert obs.type == "observation"

    def test_confidence_clamped(self):
        obs = Observation(confidence=1.5)
        assert obs.confidence == 1.0
        obs2 = Observation(confidence=-0.5)
        assert obs2.confidence == 0.0

    def test_to_text_contains_type(self):
        obs = Observation.tool_execution("git", "commit", "ok", agent_id="a")
        text = obs.to_text()
        assert "Tool Execution" in text
        assert "git" in text

    def test_to_dict_from_dict_roundtrip(self):
        obs = Observation.tool_execution("helm", "install", "done", agent_id="ops")
        d = obs.to_dict()
        obs2 = Observation.from_dict(d)
        assert obs2.id == obs.id
        assert obs2.type == obs.type
        assert obs2.content == obs.content
        assert obs2.context == obs.context
        assert obs2.confidence == obs.confidence


# ─── SQLiteStore: Observations ────────────────────────────────────


class TestObservationStorage:
    def test_put_and_get(self, store):
        obs = Observation.tool_execution("kubectl", "get pods", "running")
        obs_id = store.put_observation(obs.to_dict())
        assert obs_id == obs.id
        retrieved = store.get_observation(obs_id)
        assert retrieved is not None
        assert retrieved["type"] == "tool_execution"
        assert isinstance(retrieved["content"], dict)

    def test_list_by_agent(self, store):
        for i in range(3):
            obs = Observation.agent_step(f"step {i}", agent_id="agent_A")
            store.put_observation(obs.to_dict())
        obs_b = Observation.agent_step("step B", agent_id="agent_B")
        store.put_observation(obs_b.to_dict())

        results = store.list_observations(agent_id="agent_A")
        assert len(results) == 3

    def test_list_by_type(self, store):
        store.put_observation(Observation.tool_execution("bash", "ls", "ok").to_dict())
        store.put_observation(Observation.agent_step("think").to_dict())
        results = store.list_observations(obs_type="tool_execution")
        assert all(r["type"] == "tool_execution" for r in results)

    def test_list_by_task(self, store):
        for i in range(4):
            obs = Observation.agent_step(f"step {i}", task_id="task_99")
            store.put_observation(obs.to_dict())
        results = store.list_observations(task_id="task_99")
        assert len(results) == 4

    def test_get_missing_returns_none(self, store):
        assert store.get_observation("nonexistent") is None


# ─── DAG Memory Layer ─────────────────────────────────────────────


class TestMemoryGraph:
    def test_ingest_creates_raw_node(self, graph):
        obs = Observation.tool_execution("docker", "build", "ok", agent_id="a")
        node = graph.ingest(obs)
        assert node.type == "raw"
        assert node.obs_id == obs.id
        assert node.content == obs.to_text()

    def test_raw_node_persisted(self, graph, store):
        obs = Observation.agent_step("planning", agent_id="planner")
        node = graph.ingest(obs)
        retrieved = store.get_memory_node(node.id)
        assert retrieved is not None
        assert retrieved["type"] == "raw"

    def test_summary_node_links_children(self, graph):
        obs1 = Observation.tool_execution("git", "clone", "ok")
        obs2 = Observation.tool_execution("pip", "install", "ok")
        n1 = graph.ingest(obs1)
        n2 = graph.ingest(obs2)

        summary = graph.add_summary(
            content="Set up dev environment",
            child_ids=[n1.id, n2.id],
        )
        assert summary.type == "summary"
        assert n1.id in summary.child_ids
        assert n2.id in summary.child_ids

    def test_get_children(self, graph):
        obs1 = Observation.agent_step("step A")
        obs2 = Observation.agent_step("step B")
        n1 = graph.ingest(obs1)
        n2 = graph.ingest(obs2)
        summary = graph.add_summary("Steps A and B", [n1.id, n2.id])
        children = graph.get_children(summary.id)
        child_ids = {c.id for c in children}
        assert n1.id in child_ids
        assert n2.id in child_ids

    def test_get_parents(self, graph):
        obs = Observation.agent_step("step X")
        node = graph.ingest(obs)
        summary = graph.add_summary("Summary of X", [node.id])
        parents = graph.get_parents(node.id)
        assert any(p.id == summary.id for p in parents)

    def test_raw_never_deleted_after_summary(self, graph, store):
        obs = Observation.tool_execution("terraform", "apply", "ok")
        node = graph.ingest(obs)
        graph.add_summary("Infra applied", [node.id])
        retrieved = store.get_memory_node(node.id)
        assert retrieved is not None
        assert retrieved["type"] == "raw"

    def test_list_raw_filtered_by_agent(self, graph):
        for _ in range(3):
            graph.ingest(Observation.agent_step("do something", agent_id="alice"))
        for _ in range(2):
            graph.ingest(Observation.agent_step("do something", agent_id="bob"))
        alice_nodes = graph.list_raw(agent_id="alice")
        assert len(alice_nodes) == 3

    def test_list_summaries(self, graph):
        n1 = graph.ingest(Observation.agent_step("a"))
        n2 = graph.ingest(Observation.agent_step("b"))
        graph.add_summary("Summary AB", [n1.id, n2.id])
        summaries = graph.list_summaries()
        assert len(summaries) >= 1
        assert all(s.type == "summary" for s in summaries)

    def test_count_uncompressed_raw(self, graph):
        graph.ingest(Observation.agent_step("orphan 1"))
        graph.ingest(Observation.agent_step("orphan 2"))
        n3 = graph.ingest(Observation.agent_step("will be compressed"))
        graph.add_summary("compressed", [n3.id])
        count = graph.count_uncompressed_raw()
        assert count >= 2


# ─── Expandable Recall Engine ─────────────────────────────────────


class TestExpandableRecall:
    def _setup(self, graph):
        obs1 = Observation.tool_execution("docker", "deploy backend", "success",
                                          scene="deployment")
        obs2 = Observation.tool_execution("kubectl", "get pods", "running",
                                          scene="deployment")
        obs3 = Observation.agent_step("verify health check", scene="deployment")
        n1 = graph.ingest(obs1)
        n2 = graph.ingest(obs2)
        n3 = graph.ingest(obs3)
        summary = graph.add_summary(
            "Deployed backend and verified health",
            [n1.id, n2.id, n3.id],
            scene="deployment",
        )
        return n1, n2, n3, summary

    def test_recall_returns_summary(self, graph):
        self._setup(graph)
        recall = ExpandableRecallEngine(graph)
        result = recall.recall("deploy backend", scene="deployment")
        assert "summary" in result
        assert len(result["nodes"]) >= 1

    def test_expand_level_1_no_details(self, graph):
        self._setup(graph)
        recall = ExpandableRecallEngine(graph)
        result = recall.recall("deploy", expand_level=1)
        assert result["expand_level"] == 1

    def test_expand_level_2_includes_details(self, graph):
        self._setup(graph)
        recall = ExpandableRecallEngine(graph)
        result = recall.recall("deploy backend", expand_level=2)
        assert result["expand_level"] == 2
        assert isinstance(result["details"], list)

    def test_expand_single_node(self, graph):
        n1 = graph.ingest(Observation.agent_step("child A"))
        n2 = graph.ingest(Observation.agent_step("child B"))
        summary = graph.add_summary("parent", [n1.id, n2.id])
        recall = ExpandableRecallEngine(graph)
        result = recall.expand(summary.id)
        assert result["node"]["id"] == summary.id
        assert result["child_count"] == 2

    def test_expand_missing_node(self, graph):
        recall = ExpandableRecallEngine(graph)
        result = recall.expand("node_nonexistent")
        assert "error" in result

    def test_progressive_recall(self, graph):
        self._setup(graph)
        recall = ExpandableRecallEngine(graph)
        result = recall.progressive_recall("deploy backend", scene="deployment")
        assert "context" in result
        assert result["nodes_found"] >= 1

    def test_recall_max_expand_cap(self, graph):
        recall = ExpandableRecallEngine(graph, max_expand=2)
        result = recall.recall("test", expand_level=99)
        assert result["expand_level"] == 2


# ─── Compression Worker ────────────────────────────────────────────


class TestCompressionWorker:
    def _add_raw_nodes(self, graph, count=6, agent_id="coder", scene="ops"):
        for i in range(count):
            obs = Observation.agent_step(
                f"docker container {i} status check running",
                agent_id=agent_id, scene=scene,
            )
            graph.ingest(obs)

    def test_compression_creates_summaries(self, graph, store):
        self._add_raw_nodes(graph, count=6)
        worker = CompressionWorker(graph=graph, store=store)
        result = worker.run_once(agent_id="coder", scene="ops", min_raw=3)
        assert result["raw_processed"] >= 3
        assert result["summaries_created"] >= 1

    def test_compression_skips_if_too_few(self, graph, store):
        obs = Observation.agent_step("single step", agent_id="solo")
        graph.ingest(obs)
        worker = CompressionWorker(graph=graph, store=store)
        result = worker.run_once(agent_id="solo", min_raw=5)
        assert result["summaries_created"] == 0

    def test_raw_nodes_preserved_after_compression(self, graph, store):
        self._add_raw_nodes(graph, count=4, agent_id="bot", scene="deploy")
        raw_before = graph.list_raw(agent_id="bot", scene="deploy")
        worker = CompressionWorker(graph=graph, store=store)
        worker.run_once(agent_id="bot", scene="deploy", min_raw=3)
        raw_after = graph.list_raw(agent_id="bot", scene="deploy")
        assert len(raw_after) == len(raw_before)

    def test_compression_with_llm_fn(self, graph, store):
        self._add_raw_nodes(graph, count=5, agent_id="ai", scene="coding")
        calls = []

        def mock_llm(prompt):
            calls.append(prompt)
            return "Checked 5 docker containers, all running"

        worker = CompressionWorker(graph=graph, store=store, llm_fn=mock_llm)
        result = worker.run_once(agent_id="ai", scene="coding", min_raw=3)
        assert result["summaries_created"] >= 1
        assert len(calls) >= 1


# ─── Memory Engine ─────────────────────────────────────────────────


class TestMemoryEngine:
    def test_ingest_returns_node(self, engine):
        obs = Observation.tool_execution("npm", "install", "ok")
        node = engine.ingest(obs)
        assert node.type == "raw"
        assert node.id.startswith("node_")

    def test_ingest_persists_observation(self, engine, store):
        obs = Observation.agent_step("check logs", agent_id="dev")
        engine.ingest(obs)
        retrieved = store.get_observation(obs.id)
        assert retrieved is not None

    def test_recall_returns_dict(self, engine):
        engine.ingest(Observation.agent_step("deploy app to production"))
        result = engine.recall("deploy production")
        assert isinstance(result, dict)
        assert "summary" in result

    def test_expand_via_engine(self, engine):
        n1 = engine.ingest(Observation.agent_step("child A"))
        n2 = engine.ingest(Observation.agent_step("child B"))
        summary = engine.graph.add_summary("AB", [n1.id, n2.id])
        result = engine.expand(summary.id)
        assert result["child_count"] == 2

    def test_compress_via_engine(self, engine):
        for i in range(5):
            engine.ingest(Observation.agent_step(
                f"step {i} analyzing kubernetes cluster", agent_id="ops"
            ))
        result = engine.compress(agent_id="ops")
        assert "summaries_created" in result

    def test_convenience_tool_use_ingest(self, engine):
        node = engine.ingest_tool_use(
            tool="psql", action="SELECT count(*)", result="42",
            task_id="t1", scene="db",
        )
        assert node.type == "raw"

    def test_convenience_step_ingest(self, engine):
        node = engine.ingest_step("planning task breakdown", scene="planning")
        assert node.type == "raw"

    def test_stats(self, engine):
        engine.ingest(Observation.agent_step("a", agent_id="stat_agent"))
        stats = engine.stats(agent_id="stat_agent")
        assert "raw_nodes" in stats
        assert "observations" in stats
        assert stats["raw_nodes"] >= 1

    def test_recall_for_prompt_string(self, engine):
        engine.ingest(Observation.tool_execution("bash", "git pull", "done"))
        ctx = engine.recall_for_prompt("git pull changes")
        assert isinstance(ctx, str)


# ─── Memory Hooks ─────────────────────────────────────────────────


class TestMemoryHooks:
    def test_on_tool_use(self, engine):
        hooks = MemoryHooks(engine=engine, agent_id="hook_agent")
        node_id = hooks.on_tool_use({
            "tool": "docker",
            "input": "build image",
            "output": "built successfully",
            "task_id": "t1",
            "scene": "ci",
        })
        assert node_id is not None
        assert hooks.ingested_count == 1

    def test_on_agent_step(self, engine):
        hooks = MemoryHooks(engine=engine)
        hooks.on_agent_step({
            "step": "reviewing test results",
            "reasoning": "all tests passed",
            "agent_id": "qa_agent",
        })
        assert hooks.ingested_count == 1

    def test_on_task_complete(self, engine):
        hooks = MemoryHooks(engine=engine)
        hooks.on_task_complete({
            "task_id": "task_deploy",
            "status": "success",
            "summary": "deployed v2.0 to production",
        })
        assert hooks.ingested_count == 1

    def test_on_session_start(self, engine):
        hooks = MemoryHooks(engine=engine, agent_id="main_agent")
        hooks.on_session_start({"session_id": "sess_1", "scene": "coding"})
        assert hooks.ingested_count == 1

    def test_on_task_start(self, engine):
        hooks = MemoryHooks(engine=engine)
        hooks.on_task_start({"task_id": "t2", "scene": "deploy"})
        assert hooks.ingested_count == 1

    def test_string_shorthand(self, engine):
        hooks = MemoryHooks(engine=engine)
        hooks.on_tool_use("bash_tool_event")
        hooks.on_agent_step("reason step")
        assert hooks.ingested_count == 2

    def test_custom_listener(self, engine):
        events = []
        hooks = MemoryHooks(engine=engine)
        hooks.listen("tool_use", lambda e: events.append(e.event_type))
        hooks.on_tool_use({"tool": "curl", "input": "GET /api", "output": "200"})
        assert "tool_use" in events

    def test_no_engine_does_not_crash(self):
        hooks = MemoryHooks(engine=None, auto_ingest=True)
        result = hooks.on_tool_use({"tool": "test"})
        assert result is None
        assert hooks.ingested_count == 0

    def test_status(self, engine):
        hooks = MemoryHooks(engine=engine, agent_id="x")
        status = hooks.status()
        assert status["agent_id"] == "x"
        assert status["engine_attached"] is True


# ─── SDK Integration ───────────────────────────────────────────────


class TestSDKPhase235:
    def test_ingest_dict(self, tmp_path):
        from openmemo import OpenMemo
        memo = OpenMemo(db_path=str(tmp_path / "sdk.db"))
        result = memo.ingest({
            "type": "tool_execution",
            "agent_id": "coder",
            "content": {"tool": "git", "action": "commit", "result": "ok"},
            "context": {"scene": "coding"},
            "confidence": 0.95,
        })
        assert result["status"] == "ingested"
        assert "node_id" in result
        assert "obs_id" in result

    def test_ingest_observation_object(self, tmp_path):
        from openmemo import OpenMemo, Observation
        memo = OpenMemo(db_path=str(tmp_path / "sdk2.db"))
        obs = Observation.agent_step("thinking", agent_id="ai")
        result = memo.ingest(obs)
        assert result["status"] == "ingested"

    def test_recall_expandable(self, tmp_path):
        from openmemo import OpenMemo
        memo = OpenMemo(db_path=str(tmp_path / "sdk3.db"))
        memo.ingest({"type": "tool_execution", "content": {"tool": "docker",
                     "action": "deploy", "result": "ok"}, "context": {"scene": "ops"}})
        result = memo.recall_expandable("docker deploy", expand_level=2)
        assert "summary" in result
        assert "details" in result
        assert "expand_level" in result

    def test_compress_memory(self, tmp_path):
        from openmemo import OpenMemo
        memo = OpenMemo(db_path=str(tmp_path / "sdk4.db"))
        for i in range(5):
            obs = Observation.agent_step(f"analyzing cluster node {i} status",
                                         agent_id="ops")
            memo.ingest(obs)
        result = memo.compress_memory(agent_id="ops")
        assert "summaries_created" in result

    def test_engine_stats(self, tmp_path):
        from openmemo import OpenMemo
        memo = OpenMemo(db_path=str(tmp_path / "sdk5.db"))
        obs = Observation.agent_step("step", agent_id="stat_agent")
        memo.ingest(obs)
        stats = memo.engine_stats(agent_id="stat_agent")
        assert "raw_nodes" in stats
        assert stats["raw_nodes"] >= 1

    def test_list_observations(self, tmp_path):
        from openmemo import OpenMemo
        memo = OpenMemo(db_path=str(tmp_path / "sdk6.db"))
        obs1 = Observation.tool_execution("helm", "upgrade", "done", agent_id="ops")
        obs2 = Observation.agent_step("review", agent_id="ops")
        memo.ingest(obs1)
        memo.ingest(obs2)
        results = memo.list_observations(agent_id="ops")
        assert len(results) >= 2
