#!/usr/bin/env python3
"""
Tests for Story 2.06 (Track B): JITHydrationInterceptor — Wires Hydration Into Chain

Black-box tests:
  BB1: After pre_execute, payload has "system_injection" with "<ZERO_AMNESIA_STATE>"
  BB2: On hydration error → payload has system_injection with ERROR tag
  BB3: on_correction payload has "CORRECTION: " prefix in prompt
  BB4: on_error returns dict with error envelope

White-box tests:
  WB1: Priority == 10 (verify metadata.priority == 10)
  WB2: post_execute does NOT mutate the result dict
  WB3: post_execute logs to events.jsonl (mock file write)
  WB4: metadata.name == "jit_hydration"
  WB5: on_correction without "prompt" key doesn't crash

All mocked — no external I/O.
"""
import asyncio
import json
import sys
from pathlib import Path
from unittest.mock import AsyncMock, MagicMock, patch, mock_open, call

sys.path.insert(0, '/mnt/e/genesis-system')


# ---------------------------------------------------------------------------
# Helper: run async in test context
# ---------------------------------------------------------------------------

def _run(coro):
    """Execute a coroutine synchronously for test use."""
    return asyncio.get_event_loop().run_until_complete(coro)


# ---------------------------------------------------------------------------
# BB1: pre_execute attaches system_injection with <ZERO_AMNESIA_STATE>
# ---------------------------------------------------------------------------

def test_bb1_pre_execute_attaches_system_injection():
    """BB1: After pre_execute, payload has 'system_injection' with '<ZERO_AMNESIA_STATE>'."""
    from core.memory.jit_hydration_interceptor import JITHydrationInterceptor

    async def mock_hydration(payload: dict) -> dict:
        payload["system_injection"] = "<ZERO_AMNESIA_STATE><WORKING_CONTEXT>state</WORKING_CONTEXT></ZERO_AMNESIA_STATE>"
        return payload

    interceptor = JITHydrationInterceptor()
    payload = {"task_id": "t1", "prompt": "Do something"}

    with patch(
        "core.memory.jit_hydration_interceptor.interceptor_jit_hydration",
        side_effect=mock_hydration,
    ):
        result = _run(interceptor.pre_execute(payload))

    assert "system_injection" in result, "BB1: system_injection key must be present"
    assert "<ZERO_AMNESIA_STATE>" in result["system_injection"], \
        "BB1: system_injection must contain '<ZERO_AMNESIA_STATE>'"
    assert result["task_id"] == "t1", "BB1: original keys must be preserved"

    print("BB1 PASSED — pre_execute attaches system_injection with <ZERO_AMNESIA_STATE>")


# ---------------------------------------------------------------------------
# BB2: Hydration error → system_injection gets ERROR envelope
# ---------------------------------------------------------------------------

def test_bb2_hydration_error_produces_error_envelope():
    """BB2: When interceptor_jit_hydration raises, payload gets ERROR envelope."""
    from core.memory.jit_hydration_interceptor import JITHydrationInterceptor

    async def failing_hydration(payload: dict) -> dict:
        raise RuntimeError("Redis offline")

    interceptor = JITHydrationInterceptor()
    payload = {"task_id": "t_err", "prompt": "some task"}

    with patch(
        "core.memory.jit_hydration_interceptor.interceptor_jit_hydration",
        side_effect=failing_hydration,
    ):
        result = _run(interceptor.pre_execute(payload))

    assert "system_injection" in result, "BB2: system_injection must be set even on error"
    assert "<ERROR>" in result["system_injection"], \
        "BB2: system_injection must contain <ERROR> tag on hydration failure"
    assert "Hydration failed" in result["system_injection"], \
        "BB2: system_injection must state 'Hydration failed'"

    print("BB2 PASSED — hydration error produces ERROR envelope, execution not blocked")


# ---------------------------------------------------------------------------
# BB3: on_correction prepends "CORRECTION: " to prompt
# ---------------------------------------------------------------------------

def test_bb3_on_correction_prepends_correction_prefix():
    """BB3: on_correction should prepend 'CORRECTION: ' to existing prompt."""
    from core.memory.jit_hydration_interceptor import JITHydrationInterceptor

    async def mock_hydration(payload: dict) -> dict:
        payload["system_injection"] = "<ZERO_AMNESIA_STATE><WORKING_CONTEXT>corrected</WORKING_CONTEXT></ZERO_AMNESIA_STATE>"
        return payload

    interceptor = JITHydrationInterceptor()
    payload = {"task_id": "t_corr", "prompt": "original prompt text"}

    with patch(
        "core.memory.jit_hydration_interceptor.interceptor_jit_hydration",
        side_effect=mock_hydration,
    ):
        result = _run(interceptor.on_correction(payload))

    assert result["prompt"].startswith("CORRECTION: "), \
        f"BB3: prompt must start with 'CORRECTION: ', got: {result['prompt'][:30]}"
    assert "original prompt text" in result["prompt"], \
        "BB3: original prompt text must be preserved after prefix"
    assert "system_injection" in result, \
        "BB3: system_injection must be set after on_correction"

    print("BB3 PASSED — on_correction prepends 'CORRECTION: ' to prompt and re-hydrates")


# ---------------------------------------------------------------------------
# BB4: on_error returns dict with error envelope
# ---------------------------------------------------------------------------

def test_bb4_on_error_returns_error_envelope():
    """BB4: on_error returns dict with 'error', 'task_payload', and error system_injection."""
    from core.memory.jit_hydration_interceptor import JITHydrationInterceptor

    interceptor = JITHydrationInterceptor()
    exc = ValueError("something broke")
    payload = {"task_id": "t_onerror", "prompt": "task text"}

    result = _run(interceptor.on_error(exc, payload))

    assert isinstance(result, dict), "BB4: on_error must return a dict"
    assert "error" in result, "BB4: result must have 'error' key"
    assert "something broke" in result["error"], \
        "BB4: result['error'] must contain the exception message"
    assert "task_payload" in result, "BB4: result must have 'task_payload' key"
    assert "system_injection" in result, "BB4: result must have 'system_injection' key"
    assert "<ERROR>" in result["system_injection"], \
        "BB4: system_injection must contain <ERROR> tag"

    print("BB4 PASSED — on_error returns dict with error envelope")


# ---------------------------------------------------------------------------
# WB1: metadata.priority == 10
# ---------------------------------------------------------------------------

def test_wb1_priority_is_10():
    """WB1: JITHydrationInterceptor.metadata.priority must equal 10."""
    from core.memory.jit_hydration_interceptor import JITHydrationInterceptor

    interceptor = JITHydrationInterceptor()

    assert interceptor.metadata.priority == 10, \
        f"WB1: Expected priority=10, got priority={interceptor.metadata.priority}"

    print("WB1 PASSED — metadata.priority == 10")


# ---------------------------------------------------------------------------
# WB2: post_execute does NOT mutate result dict
# ---------------------------------------------------------------------------

def test_wb2_post_execute_does_not_mutate_result():
    """WB2: post_execute must not add, remove, or change any key in result."""
    from core.memory.jit_hydration_interceptor import JITHydrationInterceptor

    interceptor = JITHydrationInterceptor()
    result_before = {"output": "agent answer", "status": "ok"}
    result_snapshot = dict(result_before)
    task_payload = {
        "task_id": "t_wb2",
        "system_injection": "<ZERO_AMNESIA_STATE><WORKING_CONTEXT>x</WORKING_CONTEXT></ZERO_AMNESIA_STATE>",
    }

    with patch("builtins.open", mock_open()), \
         patch("core.memory.jit_hydration_interceptor.EVENTS_DIR") as mock_dir:
        mock_dir.mkdir = MagicMock()
        mock_dir.__truediv__ = lambda self, other: Path("/tmp/events.jsonl")
        _run(interceptor.post_execute(result_before, task_payload))

    assert result_before == result_snapshot, \
        f"WB2: post_execute must not mutate result. Before={result_snapshot}, After={result_before}"

    print("WB2 PASSED — post_execute does not mutate result dict")


# ---------------------------------------------------------------------------
# WB3: post_execute logs to events.jsonl
# ---------------------------------------------------------------------------

def test_wb3_post_execute_logs_to_events_jsonl():
    """WB3: post_execute writes a JSON line to events.jsonl with correct fields."""
    from core.memory.jit_hydration_interceptor import JITHydrationInterceptor

    interceptor = JITHydrationInterceptor()
    task_payload = {
        "task_id": "t_wb3",
        "system_injection": "<ZERO_AMNESIA_STATE>some content</ZERO_AMNESIA_STATE>",
    }
    result = {"output": "done"}

    written_lines = []

    m = mock_open()
    m.return_value.__enter__.return_value.write = lambda data: written_lines.append(data)

    with patch("builtins.open", m), \
         patch("core.memory.jit_hydration_interceptor.EVENTS_DIR") as mock_dir:
        mock_dir.mkdir = MagicMock()
        mock_dir.__truediv__ = lambda self, other: Path("/tmp/events.jsonl")
        _run(interceptor.post_execute(result, task_payload))

    assert len(written_lines) >= 1, "WB3: post_execute must write at least one line"
    written_json = json.loads(written_lines[0])
    assert written_json["event_type"] == "jit_hydration_complete", \
        "WB3: event_type must be 'jit_hydration_complete'"
    assert written_json["task_id"] == "t_wb3", \
        "WB3: task_id must match payload task_id"
    assert "injection_length" in written_json, \
        "WB3: event must include injection_length field"
    assert "timestamp" in written_json, \
        "WB3: event must include timestamp field"

    print("WB3 PASSED — post_execute logs correct JSON event to events.jsonl")


# ---------------------------------------------------------------------------
# WB4: metadata.name == "jit_hydration"
# ---------------------------------------------------------------------------

def test_wb4_metadata_name_is_jit_hydration():
    """WB4: metadata.name must be 'jit_hydration'."""
    from core.memory.jit_hydration_interceptor import JITHydrationInterceptor

    interceptor = JITHydrationInterceptor()

    assert interceptor.metadata.name == "jit_hydration", \
        f"WB4: Expected name='jit_hydration', got name='{interceptor.metadata.name}'"

    print("WB4 PASSED — metadata.name == 'jit_hydration'")


# ---------------------------------------------------------------------------
# WB5: on_correction without "prompt" key doesn't crash
# ---------------------------------------------------------------------------

def test_wb5_on_correction_no_prompt_key_safe():
    """WB5: on_correction with no 'prompt' key must not raise and must still hydrate."""
    from core.memory.jit_hydration_interceptor import JITHydrationInterceptor

    async def mock_hydration(payload: dict) -> dict:
        payload["system_injection"] = "<ZERO_AMNESIA_STATE><WORKING_CONTEXT>no-prompt</WORKING_CONTEXT></ZERO_AMNESIA_STATE>"
        return payload

    interceptor = JITHydrationInterceptor()
    payload = {"task_id": "t_noprompt", "description": "task without explicit prompt"}

    with patch(
        "core.memory.jit_hydration_interceptor.interceptor_jit_hydration",
        side_effect=mock_hydration,
    ):
        try:
            result = _run(interceptor.on_correction(payload))
        except Exception as e:
            raise AssertionError(
                f"WB5: on_correction must not raise when 'prompt' key absent. Got: {e}"
            )

    assert "prompt" not in result or not result.get("prompt", "").startswith("CORRECTION: "), \
        "WB5: Should not add 'CORRECTION: ' prefix when no 'prompt' key exists"
    assert "system_injection" in result, \
        "WB5: system_injection must still be set even without 'prompt' key"

    print("WB5 PASSED — on_correction without 'prompt' key safe, system_injection still set")


# ---------------------------------------------------------------------------
# Entry point
# ---------------------------------------------------------------------------

if __name__ == "__main__":
    print("=" * 65)
    print("Story 2.06 — JITHydrationInterceptor Tests")
    print("=" * 65)

    test_bb1_pre_execute_attaches_system_injection()
    test_bb2_hydration_error_produces_error_envelope()
    test_bb3_on_correction_prepends_correction_prefix()
    test_bb4_on_error_returns_error_envelope()
    test_wb1_priority_is_10()
    test_wb2_post_execute_does_not_mutate_result()
    test_wb3_post_execute_logs_to_events_jsonl()
    test_wb4_metadata_name_is_jit_hydration()
    test_wb5_on_correction_no_prompt_key_safe()

    print("=" * 65)
    print("ALL 9 TESTS PASSED — Story 2.06 (Track B)")
    print("=" * 65)