"""
Tests for Context Memory Layer — PRD v1.0

Coverage:
    §5   Data models (SessionNode, TimelineEvent, RawMessage)
    §6   Storage (ContextStore)
    §7   Summary generation (ContextSummary)
    §8   Timeline events (ContextTimeline)
    §9   Progressive recall (ContextRecall)
    §10  Token budget (ContextBudget)
    §11  Prompt injection order
    §12  ContextMemory facade + adapter integration
"""

import os
import sys
import time
import pytest
import tempfile

OPENMEMO_ROOT = os.path.join(os.path.dirname(__file__), "..")
sys.path.insert(0, OPENMEMO_ROOT)

from openmemo.context_memory.context_store import (
    ContextStore, SessionNode, TimelineEvent, RawMessage
)
from openmemo.context_memory.context_summary import ContextSummary
from openmemo.context_memory.context_timeline import ContextTimeline
from openmemo.context_memory.context_recall import ContextRecall
from openmemo.context_memory.context_budget import ContextBudget, estimate_tokens
from openmemo.context_memory import ContextMemory


@pytest.fixture
def tmp_db(tmp_path):
    return str(tmp_path / "test_context.db")


@pytest.fixture
def store(tmp_db):
    return ContextStore(db_path=tmp_db)


@pytest.fixture
def cm(tmp_db):
    return ContextMemory(db_path=tmp_db, max_tokens=800, summarize_every_n=5)


# ══════════════════════════════════════════════════════════════════════════════
# §5 — Data Models
# ══════════════════════════════════════════════════════════════════════════════

class TestDataModels:
    def test_session_node_defaults(self):
        s = SessionNode(session_id="s1")
        assert s.session_id == "s1"
        assert s.summary == ""
        assert s.message_count == 0
        assert s.created_at > 0

    def test_session_node_to_dict(self):
        s = SessionNode(session_id="s1", goal="deploy API")
        d = s.to_dict()
        assert d["session_id"] == "s1"
        assert d["goal"] == "deploy API"

    def test_timeline_event_defaults(self):
        ev = TimelineEvent(session_id="s1", type="action", content="step done")
        assert ev.type == "action"
        assert ev.content == "step done"
        assert ev.event_id != ""
        assert ev.timestamp > 0

    def test_timeline_event_to_dict(self):
        ev = TimelineEvent(session_id="s1", type="decision", content="use docker")
        d = ev.to_dict()
        assert d["type"] == "decision"
        assert d["content"] == "use docker"

    def test_raw_message_defaults(self):
        m = RawMessage(session_id="s1", role="user", content="help me deploy")
        assert m.role == "user"
        assert m.message_id != ""


# ══════════════════════════════════════════════════════════════════════════════
# §6 — ContextStore
# ══════════════════════════════════════════════════════════════════════════════

class TestContextStore:
    def test_save_and_get_session(self, store):
        s = SessionNode(session_id="s1", user_id="alice",
                         summary="Deploying backend")
        store.save_session(s)
        got = store.get_session("s1")
        assert got is not None
        assert got.user_id == "alice"
        assert got.summary == "Deploying backend"

    def test_get_or_create_session(self, store):
        s = store.get_or_create_session("s2", user_id="bob")
        assert s.session_id == "s2"
        assert s.user_id == "bob"

    def test_get_nonexistent_session_returns_none(self, store):
        assert store.get_session("nonexistent") is None

    def test_session_upsert(self, store):
        s = SessionNode(session_id="s1", summary="v1")
        store.save_session(s)
        s.summary = "v2"
        store.save_session(s)
        got = store.get_session("s1")
        assert got.summary == "v2"

    def test_list_sessions_by_user(self, store):
        for i in range(3):
            store.save_session(SessionNode(session_id=f"s{i}", user_id="alice"))
        store.save_session(SessionNode(session_id="s99", user_id="bob"))
        result = store.list_sessions(user_id="alice")
        assert len(result) == 3
        assert all(s.user_id == "alice" for s in result)

    def test_increment_message_count(self, store):
        store.get_or_create_session("s1")
        count = store.increment_message_count("s1")
        assert count == 1
        count = store.increment_message_count("s1")
        assert count == 2

    def test_add_and_get_events(self, store):
        store.get_or_create_session("s1")
        ev = TimelineEvent(session_id="s1", type="action", content="step A")
        eid = store.add_event(ev)
        events = store.get_timeline("s1")
        assert len(events) == 1
        assert events[0].content == "step A"

    def test_get_recent_events(self, store):
        store.get_or_create_session("s1")
        for i in range(10):
            store.add_event(TimelineEvent(session_id="s1", type="action",
                                           content=f"step {i}"))
        recent = store.get_recent_events("s1", limit=3)
        assert len(recent) == 3
        # should be in chronological order
        assert recent[-1].content == "step 9"

    def test_count_events(self, store):
        store.get_or_create_session("s1")
        for _ in range(5):
            store.add_event(TimelineEvent(session_id="s1", type="result",
                                           content="ok"))
        assert store.count_events("s1") == 5

    def test_add_and_get_raw_messages(self, store):
        store.get_or_create_session("s1")
        store.add_raw_message(RawMessage(session_id="s1", role="user",
                                          content="hello"))
        store.add_raw_message(RawMessage(session_id="s1", role="assistant",
                                          content="hi there"))
        msgs = store.get_raw_messages("s1")
        assert len(msgs) == 2
        assert msgs[0].role == "user"

    def test_delete_session(self, store):
        store.get_or_create_session("s1")
        store.add_event(TimelineEvent(session_id="s1", type="action",
                                       content="test"))
        store.delete_session("s1")
        assert store.get_session("s1") is None
        assert store.get_timeline("s1") == []


# ══════════════════════════════════════════════════════════════════════════════
# §7 — ContextSummary
# ══════════════════════════════════════════════════════════════════════════════

class TestContextSummary:
    def test_summarize_basic(self, store):
        store.get_or_create_session("s1")
        store.add_raw_message(RawMessage(session_id="s1", role="user",
                                          content="I want to deploy my FastAPI backend"))
        store.add_raw_message(RawMessage(session_id="s1", role="assistant",
                                          content="I decided to use docker-compose"))
        summarizer = ContextSummary(store)
        session = summarizer.summarize("s1")
        # Should extract something meaningful
        assert session.session_id == "s1"
        # Goal or decisions should be extracted
        assert session.summary != "" or session.goal != "" or session.decisions != ""

    def test_summarize_empty_session(self, store):
        store.get_or_create_session("s1")
        summarizer = ContextSummary(store)
        session = summarizer.summarize("s1")
        assert session.session_id == "s1"

    def test_should_summarize_threshold(self, store):
        store.get_or_create_session("s1")
        summarizer = ContextSummary(store, summarize_every_n=5)
        # Before threshold
        assert not summarizer.should_summarize("s1")
        # Simulate 5 messages
        for _ in range(5):
            store.increment_message_count("s1")
        assert summarizer.should_summarize("s1")

    def test_summarize_on_complete(self, store):
        store.get_or_create_session("s1")
        summarizer = ContextSummary(store)
        session = summarizer.summarize_on_complete(
            "s1", task="Deploy FastAPI", result="Docker image built"
        )
        assert session.progress == "Docker image built"

    def test_format_for_prompt_empty(self, store):
        store.get_or_create_session("s1")
        summarizer = ContextSummary(store)
        assert summarizer.format_for_prompt("s1") == ""

    def test_format_for_prompt_with_data(self, store):
        s = SessionNode(session_id="s1", goal="deploy API",
                         progress="50% done", decisions="use docker")
        store.save_session(s)
        summarizer = ContextSummary(store)
        block = summarizer.format_for_prompt("s1")
        assert "CONTEXT SUMMARY" in block
        assert "deploy API" in block
        assert "50% done" in block

    def test_rule_extracts_goal(self, store):
        store.get_or_create_session("s1")
        store.add_raw_message(RawMessage(session_id="s1", role="user",
                                          content="I need to build a REST API"))
        summarizer = ContextSummary(store)
        session = summarizer.summarize("s1")
        assert "api" in (session.goal + session.summary).lower()

    def test_rule_extracts_decision(self, store):
        store.get_or_create_session("s1")
        store.add_raw_message(RawMessage(session_id="s1", role="user",
                                          content="I decided to use PostgreSQL"))
        summarizer = ContextSummary(store)
        session = summarizer.summarize("s1")
        assert "postgres" in (session.decisions + session.summary).lower()


# ══════════════════════════════════════════════════════════════════════════════
# §8 — ContextTimeline
# ══════════════════════════════════════════════════════════════════════════════

class TestContextTimeline:
    def test_add_event(self, store):
        store.get_or_create_session("s1")
        tl = ContextTimeline(store)
        eid = tl.add("s1", "action", "Docker image built")
        assert eid is not None

    def test_event_types(self, store):
        store.get_or_create_session("s1")
        tl = ContextTimeline(store)
        tl.on_agent_step("s1", "Analyzing the task")
        tl.on_tool_call("s1", "bash", "exit code 0")
        tl.on_decision("s1", "Chose PostgreSQL")
        tl.on_result("s1", "Deployment successful")
        tl.on_error("s1", "Port 5000 in use")
        events = store.get_timeline("s1")
        assert len(events) == 5

    def test_invalid_event_type_defaults_to_action(self, store):
        store.get_or_create_session("s1")
        tl = ContextTimeline(store)
        eid = tl.add("s1", "invalid_type", "some content")
        events = store.get_timeline("s1")
        assert events[0].type == "action"

    def test_empty_content_returns_none(self, store):
        store.get_or_create_session("s1")
        tl = ContextTimeline(store)
        eid = tl.add("s1", "action", "")
        assert eid is None

    def test_get_top(self, store):
        store.get_or_create_session("s1")
        tl = ContextTimeline(store)
        for i in range(10):
            tl.add("s1", "action", f"step {i}")
        top = tl.get_top("s1", limit=3)
        assert len(top) == 3
        assert top[-1].content == "step 9"

    def test_format_for_prompt(self, store):
        store.get_or_create_session("s1")
        tl = ContextTimeline(store)
        tl.on_decision("s1", "Use FastAPI")
        tl.on_result("s1", "Server started")
        block = tl.format_for_prompt("s1")
        assert "Timeline:" in block
        assert "Use FastAPI" in block
        assert "Server started" in block

    def test_format_empty_session(self, store):
        store.get_or_create_session("s1")
        tl = ContextTimeline(store)
        assert tl.format_for_prompt("s1") == ""

    def test_count(self, store):
        store.get_or_create_session("s1")
        tl = ContextTimeline(store)
        for _ in range(7):
            tl.add("s1", "action", "step")
        assert tl.count("s1") == 7


# ══════════════════════════════════════════════════════════════════════════════
# §10 — ContextBudget
# ══════════════════════════════════════════════════════════════════════════════

class TestContextBudget:
    def test_estimate_tokens(self):
        assert estimate_tokens("") == 0
        assert estimate_tokens("hello world") > 0
        assert estimate_tokens("a " * 100) > estimate_tokens("a " * 10)

    def test_fit_all_within_budget(self):
        budget = ContextBudget(max_tokens=800)
        result = budget.fit(
            user_memory_block="USER PROFILE\n  - prefers Python",
            summary_block="CONTEXT SUMMARY\n  Goal: deploy API",
            timeline_block="Timeline:\n  [action] step 1",
        )
        assert "user_memory" in result
        assert "summary" in result
        assert "timeline" in result

    def test_fit_drops_low_priority_when_over_budget(self):
        budget = ContextBudget(max_tokens=30)  # very tight budget
        long_user = "USER PROFILE\n" + " user prefers " * 20
        long_summary = "CONTEXT SUMMARY\n" + " deploying backend " * 20
        result = budget.fit(
            user_memory_block=long_user,
            summary_block=long_summary,
            timeline_block="Timeline:\n  [action] some step",
            raw_block="Raw: some old messages here",
        )
        # Timeline and raw may be dropped or truncated
        # But user_memory should always be attempted first
        assert "user_memory" in result

    def test_combine_ordering(self):
        budget = ContextBudget(max_tokens=800)
        combined = budget.combine(
            user_memory_block="USER PROFILE",
            summary_block="CONTEXT SUMMARY",
            timeline_block="Timeline:",
        )
        # Context summary should appear before user memory
        assert combined.index("CONTEXT SUMMARY") < combined.index("USER PROFILE")

    def test_truncate_to_tokens(self):
        budget = ContextBudget(max_tokens=100)
        long_text = "word " * 200
        result = budget._truncate_to_tokens(long_text, 20)
        assert len(result.split()) < 30  # roughly truncated
        assert result.endswith("...")

    def test_remaining_tokens(self):
        budget = ContextBudget(max_tokens=800)
        rem = budget.remaining_tokens("hello world")
        assert rem < 800
        assert rem > 0


# ══════════════════════════════════════════════════════════════════════════════
# §9 — ContextRecall
# ══════════════════════════════════════════════════════════════════════════════

class TestContextRecall:
    def test_recall_level1_summary_only(self, store):
        s = SessionNode(session_id="s1", goal="deploy API",
                         summary="Deploying FastAPI")
        store.save_session(s)
        recall = ContextRecall(store)
        block = recall.recall("s1", level=1)
        assert "CONTEXT SUMMARY" in block or "deploy" in block.lower()

    def test_recall_level2_includes_timeline(self, store):
        s = SessionNode(session_id="s1", goal="deploy API",
                         summary="Deploying FastAPI")
        store.save_session(s)
        store.add_event(TimelineEvent(session_id="s1", type="action",
                                       content="Docker image built"))
        recall = ContextRecall(store)
        block = recall.recall("s1", level=2)
        assert "Docker image built" in block or block != ""

    def test_recall_level3_includes_raw(self, store):
        store.get_or_create_session("s1")
        store.add_raw_message(RawMessage(session_id="s1", role="user",
                                          content="Raw message content here"))
        recall = ContextRecall(store)
        block = recall.recall("s1", level=3)
        # With level 3, raw messages should be attempted
        assert block is not None

    def test_recall_as_dict(self, store):
        s = SessionNode(session_id="s1", goal="build service",
                         summary="Building a FastAPI service", progress="50%")
        store.save_session(s)
        store.add_event(TimelineEvent(session_id="s1", type="decision",
                                       content="Use SQLite"))
        recall = ContextRecall(store)
        result = recall.recall_as_dict("s1")
        assert result["session_id"] == "s1"
        assert result["goal"] == "build service"
        assert len(result["timeline"]) == 1

    def test_auto_level_empty_session(self, store):
        store.get_or_create_session("s1")
        recall = ContextRecall(store)
        block = recall.auto_level("s1")
        assert isinstance(block, str)

    def test_auto_level_with_user_memory(self, store):
        s = SessionNode(session_id="s1", goal="goal", summary="summary text")
        store.save_session(s)
        recall = ContextRecall(store)
        block = recall.auto_level("s1", user_memory_block="USER PROFILE\n- prefers Python")
        assert isinstance(block, str)


# ══════════════════════════════════════════════════════════════════════════════
# §12 — ContextMemory Facade
# ══════════════════════════════════════════════════════════════════════════════

class TestContextMemory:
    def test_add_message_creates_session(self, cm):
        mid = cm.add_message("s1", "user", "Deploy my backend", user_id="alice")
        assert mid != ""
        session = cm.get_session("s1")
        assert session.session_id == "s1"
        assert session.message_count == 1

    def test_add_messages_increments_count(self, cm):
        for i in range(3):
            cm.add_message("s1", "user" if i % 2 == 0 else "assistant",
                           f"message {i}")
        session = cm.get_session("s1")
        assert session.message_count == 3

    def test_auto_summary_triggered_at_n(self, cm):
        # summarize_every_n=5
        for i in range(5):
            cm.add_message("s1", "user", f"I need to build service {i}")
        session = cm.get_session("s1")
        assert session.message_count == 5
        # Summary should have been generated
        assert session.summary != "" or session.goal != ""

    def test_add_event(self, cm):
        eid = cm.add_event("s1", "decision", "Use PostgreSQL")
        assert eid != ""

    def test_on_agent_step(self, cm):
        cm.on_agent_step("s1", "Analyzing requirements", user_id="alice")
        # Should have created session and timeline event
        session = cm.get_session("s1")
        assert session.session_id == "s1"

    def test_on_tool_call(self, cm):
        cm.add_message("s1", "user", "test")
        cm.on_tool_call("s1", "bash", "exit 0")

    def test_on_task_complete(self, cm):
        cm.add_message("s1", "user", "Deploy my API")
        session = cm.on_task_complete(
            "s1", task="Deploy FastAPI", result="Deployed successfully"
        )
        assert session.progress == "Deployed successfully"

    def test_recall_level1(self, cm):
        cm.add_message("s1", "user", "I want to build a REST API")
        cm.on_task_complete("s1", task="Build REST API", result="Done")
        block = cm.recall("s1", level=1)
        assert isinstance(block, str)

    def test_recall_level2(self, cm):
        cm.add_message("s1", "user", "Deploy backend")
        cm.add_event("s1", "action", "Docker image built")
        block = cm.recall("s1", level=2)
        assert isinstance(block, str)

    def test_inject_returns_string(self, cm):
        cm.add_message("s1", "user", "Deploy my FastAPI backend")
        block = cm.inject("s1")
        assert isinstance(block, str)

    def test_inject_with_user_memory(self, cm):
        cm.add_message("s1", "user", "help me deploy")
        cm.add_event("s1", "decision", "use docker")
        block = cm.inject("s1", user_memory_block="USER PROFILE\n- prefers Python")
        assert isinstance(block, str)

    def test_format_summary(self, cm):
        cm.on_task_complete("s1", task="Deploy API", result="Done")
        block = cm.format_summary("s1")
        assert isinstance(block, str)

    def test_format_timeline(self, cm):
        cm.add_event("s1", "action", "step 1")
        block = cm.format_timeline("s1")
        assert isinstance(block, str)

    def test_to_dict(self, cm):
        cm.add_message("s1", "user", "hello")
        cm.add_event("s1", "result", "done")
        d = cm.to_dict("s1")
        assert d["session_id"] == "s1"
        assert "timeline" in d
        assert "summary" in d

    def test_list_sessions(self, cm):
        cm.add_message("s1", "user", "hello", user_id="alice")
        cm.add_message("s2", "user", "world", user_id="alice")
        sessions = cm.list_sessions(user_id="alice")
        assert len(sessions) >= 2

    def test_delete_session(self, cm):
        cm.add_message("s1", "user", "hello")
        cm.delete_session("s1")
        session = cm.get_session("s1")
        assert session.session_id == "s1"  # get_or_create returns new
        assert session.message_count == 0  # clean session


# ══════════════════════════════════════════════════════════════════════════════
# §11 — Injection Order (Integration)
# ══════════════════════════════════════════════════════════════════════════════

class TestInjectionOrder:
    def test_context_before_user_memory(self, cm):
        """PRD §11.1: Context summary must appear before user memory."""
        cm.on_task_complete("s1", task="Deploy API",
                             result="Docker image built")
        cm.add_event("s1", "decision", "Use docker-compose")
        block = cm.inject(
            "s1",
            user_memory_block="USER PROFILE\n  - prefers FastAPI"
        )
        if "CONTEXT SUMMARY" in block and "USER PROFILE" in block:
            assert block.index("CONTEXT SUMMARY") < block.index("USER PROFILE")

    def test_empty_session_inject_returns_empty(self, cm):
        """Empty session should return empty string."""
        block = cm.inject("brand_new_session_xyz")
        assert block == "" or isinstance(block, str)

    def test_token_budget_respected(self):
        """PRD §10: Token budget is enforced."""
        budget = ContextBudget(max_tokens=20)
        user_block = "USER PROFILE\n" + "very long content " * 50
        summary_block = "CONTEXT SUMMARY\n" + "more long content " * 50
        result = budget.fit(user_memory_block=user_block,
                             summary_block=summary_block)
        # Result must fit within budget
        combined = " ".join(result.values())
        token_count = estimate_tokens(combined)
        assert token_count <= 25  # small margin for truncation


# ══════════════════════════════════════════════════════════════════════════════
# Adapter Integration
# ══════════════════════════════════════════════════════════════════════════════

class TestAdapterContextIntegration:
    def test_openclaw_has_context_memory(self, tmp_db):
        from openmemo.adapters.openclaw import OpenClawMemoryBackend
        backend = OpenClawMemoryBackend(db_path=tmp_db, agent_id="agent1",
                                        conversation_id="session_abc")
        assert backend._context_memory is not None

    def test_openclaw_on_agent_step_records_event(self, tmp_db):
        from openmemo.adapters.openclaw import OpenClawMemoryBackend
        backend = OpenClawMemoryBackend(db_path=tmp_db, agent_id="agent1",
                                        conversation_id="sess1")
        backend.on_agent_step("Analyzing requirements", session_id="sess1")
        # No exception = pass

    def test_openclaw_on_task_complete_triggers_summary(self, tmp_db):
        from openmemo.adapters.openclaw import OpenClawMemoryBackend
        backend = OpenClawMemoryBackend(db_path=tmp_db, agent_id="agent1",
                                        conversation_id="sess2")
        backend.on_task_complete("Deploy API", "Done", session_id="sess2")
        # No exception = pass

    def test_openclaw_on_message_records(self, tmp_db):
        from openmemo.adapters.openclaw import OpenClawMemoryBackend
        backend = OpenClawMemoryBackend(db_path=tmp_db, agent_id="agent1",
                                        conversation_id="sess3")
        backend.on_message("user", "Help me deploy", session_id="sess3")
        # No exception = pass

    def test_base_adapter_context_methods(self, tmp_db):
        from openmemo.adapters.base_adapter import BaseMemoryAdapter
        adapter = BaseMemoryAdapter(db_path=tmp_db, agent_id="agent1",
                                     conversation_id="sess4")
        assert adapter._context_memory is not None
        mid = adapter.context_on_message("sess4", "user", "hello")
        adapter.context_on_agent_step("sess4", "thinking...")
        adapter.context_on_task_complete("sess4", task="test", result="ok")
        block = adapter.context_recall("sess4", level=2)
        assert isinstance(block, str)

    def test_inject_context_with_session(self, tmp_db):
        from openmemo.adapters.base_adapter import BaseMemoryAdapter
        adapter = BaseMemoryAdapter(db_path=tmp_db, agent_id="agent1")
        # Add some context
        adapter.context_on_message("sess5", "user", "Deploy my backend")
        adapter.context_on_task_complete("sess5", task="Deploy",
                                          result="Docker image built")
        # inject_context should prepend context summary
        result = adapter.inject_context(
            "What should I do next?", session_id="sess5"
        )
        assert isinstance(result, str)
        assert "What should I do next?" in result

    def test_inject_context_ordering(self, tmp_db):
        """PRD §11.1: Context summary appears before user memory in injection."""
        from openmemo.adapters.base_adapter import BaseMemoryAdapter
        adapter = BaseMemoryAdapter(db_path=tmp_db, agent_id="agent1",
                                     user_id="alice")
        # Build context
        adapter.context_on_task_complete("sess6", task="Deploy",
                                          result="Done")
        # Inject with user memory enabled
        result = adapter.inject_context(
            "What next?", session_id="sess6",
            include_user_memory=False
        )
        assert isinstance(result, str)
