"""
Story 9.03 — Test Suite
========================
NightlyEpochRunner: Gemini Distillation

File under test: core/epoch/nightly_epoch_runner.py

BB Tests (4):
  BB1: 5 mock conversations → mock Gemini returns 7 axioms → all 7 in result
  BB2: Empty conversations list → {"axioms": [], "week_summary": "No conversations this week"}
  BB3: Each axiom has all 4 required keys (id, content, category, confidence)
  BB4: Axiom missing a required key → filtered out of result

WB Tests (4):
  WB1: Gemini Pro model used — generate() called with model="gemini-pro"
  WB2: conversations_json in prompt includes all fields from conversations
  WB3: JSON parsed with json.loads() — invalid JSON → "Distillation failed"
  WB4: gemini_client is None → {"axioms": [], "week_summary": "No Gemini client configured"}

All tests use MagicMock / AsyncMock — zero live API calls.
"""
from __future__ import annotations

import asyncio
import json
import sys
from unittest.mock import AsyncMock, MagicMock

import pytest

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

from core.epoch.nightly_epoch_runner import (
    DISTILLATION_PROMPT,
    NightlyEpochRunner,
    _AXIOM_REQUIRED_KEYS,
)


# ---------------------------------------------------------------------------
# Helpers
# ---------------------------------------------------------------------------


def _arun(coro):
    """Run a coroutine synchronously."""
    return asyncio.run(coro)


def _make_conversations(n: int) -> list[dict]:
    """Build *n* fake conversation dicts."""
    return [
        {
            "conversation_id": f"conv-{i:04d}",
            "started_at": f"2026-02-{i + 1:02d}",
            "transcript_raw": f"User: message {i}. AIVA: reply {i}.",
            "enriched_entities": {"entity": i},
            "decisions_made": [f"decision-{i}"],
            "action_items": [f"action-{i}"],
            "key_facts": [f"fact-{i}"],
            "kinan_directives": [f"directive-{i}"],
        }
        for i in range(n)
    ]


def _make_axioms(n: int) -> list[dict]:
    """Build *n* valid axiom dicts with all required keys."""
    categories = ["preference", "fact", "strategy", "directive"]
    return [
        {
            "id": f"epoch_2026_02_25_{i:03d}",
            "content": f"Axiom content number {i}",
            "category": categories[i % len(categories)],
            "confidence": round(0.7 + (i % 3) * 0.1, 1),
        }
        for i in range(n)
    ]


def _make_gemini_client(axioms: list[dict], week_summary: str = "Good week.") -> AsyncMock:
    """
    Return an AsyncMock Gemini client whose generate() returns a valid JSON
    string containing the supplied axioms and week_summary.
    """
    payload = json.dumps({"axioms": axioms, "week_summary": week_summary})
    client = MagicMock()
    client.generate = AsyncMock(return_value=payload)
    return client


# ---------------------------------------------------------------------------
# BB1 — 5 conversations → mock Gemini returns 7 axioms → all 7 in result
# ---------------------------------------------------------------------------


class TestBB1_SevenAxiomsReturnedFromGemini:
    """BB1: distill() must return all valid axioms from the Gemini response."""

    def test_seven_axioms_all_present_in_result(self):
        """BB1: Gemini returns 7 axioms → result['axioms'] has length 7."""
        conversations = _make_conversations(5)
        axioms = _make_axioms(7)
        gemini = _make_gemini_client(axioms, week_summary="Productive week.")
        runner = NightlyEpochRunner(gemini_client=gemini)

        result = _arun(runner.distill(conversations))

        assert len(result["axioms"]) == 7, (
            f"Expected 7 axioms, got {len(result['axioms'])}"
        )

    def test_axioms_ids_match_gemini_response(self):
        """BB1: IDs from Gemini response are preserved in the result."""
        conversations = _make_conversations(5)
        axioms = _make_axioms(7)
        gemini = _make_gemini_client(axioms)
        runner = NightlyEpochRunner(gemini_client=gemini)

        result = _arun(runner.distill(conversations))

        returned_ids = {ax["id"] for ax in result["axioms"]}
        expected_ids = {ax["id"] for ax in axioms}
        assert returned_ids == expected_ids, (
            f"Axiom IDs mismatch: got {returned_ids}, expected {expected_ids}"
        )

    def test_week_summary_present_in_result(self):
        """BB1: distill() must return a week_summary string."""
        conversations = _make_conversations(5)
        axioms = _make_axioms(7)
        gemini = _make_gemini_client(axioms, week_summary="Great week overall.")
        runner = NightlyEpochRunner(gemini_client=gemini)

        result = _arun(runner.distill(conversations))

        assert result.get("week_summary") == "Great week overall.", (
            f"Expected week_summary='Great week overall.', got {result.get('week_summary')!r}"
        )

    def test_result_is_dict_with_two_keys(self):
        """BB1: Return value is a dict with exactly 'axioms' and 'week_summary'."""
        conversations = _make_conversations(3)
        axioms = _make_axioms(5)
        gemini = _make_gemini_client(axioms)
        runner = NightlyEpochRunner(gemini_client=gemini)

        result = _arun(runner.distill(conversations))

        assert isinstance(result, dict), f"Expected dict, got {type(result)}"
        assert "axioms" in result, "Result must contain 'axioms' key"
        assert "week_summary" in result, "Result must contain 'week_summary' key"


# ---------------------------------------------------------------------------
# BB2 — Empty conversations → graceful degradation
# ---------------------------------------------------------------------------


class TestBB2_EmptyConversationsGracefulDegradation:
    """BB2: Empty conversations list must return early with safe defaults."""

    def test_empty_list_returns_no_conversations_message(self):
        """BB2: distill([]) → week_summary = 'No conversations this week'."""
        gemini = _make_gemini_client(_make_axioms(3))
        runner = NightlyEpochRunner(gemini_client=gemini)

        result = _arun(runner.distill([]))

        assert result == {
            "axioms": [],
            "week_summary": "No conversations this week",
        }, f"Unexpected result for empty conversations: {result!r}"

    def test_empty_list_does_not_call_gemini(self):
        """BB2: Gemini generate() must NOT be called when conversations is empty."""
        gemini = _make_gemini_client(_make_axioms(3))
        runner = NightlyEpochRunner(gemini_client=gemini)

        _arun(runner.distill([]))

        gemini.generate.assert_not_called()

    def test_empty_axioms_in_result_is_a_list(self):
        """BB2: result['axioms'] must be an empty list (not None) on empty input."""
        runner = NightlyEpochRunner(gemini_client=MagicMock())

        result = _arun(runner.distill([]))

        assert isinstance(result["axioms"], list), (
            f"Expected list for axioms, got {type(result['axioms'])}"
        )
        assert result["axioms"] == []


# ---------------------------------------------------------------------------
# BB3 — Each axiom has all 4 required keys
# ---------------------------------------------------------------------------


class TestBB3_AxiomHasAllRequiredKeys:
    """BB3: Every axiom in the result must have id, content, category, confidence."""

    def test_all_axioms_have_four_required_keys(self):
        """BB3: Each returned axiom dict has exactly the 4 required keys."""
        conversations = _make_conversations(3)
        axioms = _make_axioms(5)
        gemini = _make_gemini_client(axioms)
        runner = NightlyEpochRunner(gemini_client=gemini)

        result = _arun(runner.distill(conversations))

        for i, ax in enumerate(result["axioms"]):
            missing = _AXIOM_REQUIRED_KEYS - ax.keys()
            assert not missing, (
                f"Axiom at index {i} is missing keys: {missing}. Got: {ax}"
            )

    def test_axiom_id_is_string(self):
        """BB3: Each axiom['id'] must be a string."""
        conversations = _make_conversations(2)
        axioms = _make_axioms(3)
        gemini = _make_gemini_client(axioms)
        runner = NightlyEpochRunner(gemini_client=gemini)

        result = _arun(runner.distill(conversations))

        for ax in result["axioms"]:
            assert isinstance(ax["id"], str), (
                f"Expected axiom id to be str, got {type(ax['id'])}"
            )

    def test_axiom_confidence_is_numeric(self):
        """BB3: Each axiom['confidence'] must be a numeric value."""
        conversations = _make_conversations(2)
        axioms = _make_axioms(3)
        gemini = _make_gemini_client(axioms)
        runner = NightlyEpochRunner(gemini_client=gemini)

        result = _arun(runner.distill(conversations))

        for ax in result["axioms"]:
            assert isinstance(ax["confidence"], (int, float)), (
                f"Expected numeric confidence, got {type(ax['confidence'])}: {ax['confidence']!r}"
            )

    def test_axiom_category_is_string(self):
        """BB3: Each axiom['category'] must be a string."""
        conversations = _make_conversations(2)
        axioms = _make_axioms(4)
        gemini = _make_gemini_client(axioms)
        runner = NightlyEpochRunner(gemini_client=gemini)

        result = _arun(runner.distill(conversations))

        for ax in result["axioms"]:
            assert isinstance(ax["category"], str), (
                f"Expected string category, got {type(ax['category'])}: {ax['category']!r}"
            )


# ---------------------------------------------------------------------------
# BB4 — Axiom missing a required key → filtered out
# ---------------------------------------------------------------------------


class TestBB4_InvalidAxiomsFiltered:
    """BB4: Axioms missing any required key must be silently filtered out."""

    def test_axiom_missing_id_is_filtered(self):
        """BB4: Axiom without 'id' key must not appear in result."""
        conversations = _make_conversations(2)
        # Two valid axioms + one missing 'id'
        axioms = _make_axioms(2)
        bad_axiom = {"content": "some content", "category": "fact", "confidence": 0.8}
        all_axioms = axioms + [bad_axiom]
        payload = json.dumps({"axioms": all_axioms, "week_summary": "Test."})
        gemini = MagicMock()
        gemini.generate = AsyncMock(return_value=payload)
        runner = NightlyEpochRunner(gemini_client=gemini)

        result = _arun(runner.distill(conversations))

        assert len(result["axioms"]) == 2, (
            f"Expected 2 valid axioms (bad one filtered), got {len(result['axioms'])}"
        )

    def test_axiom_missing_confidence_is_filtered(self):
        """BB4: Axiom without 'confidence' key must not appear in result."""
        conversations = _make_conversations(2)
        axioms = _make_axioms(3)
        # Replace the last axiom with one missing 'confidence'
        bad_axiom = {"id": "epoch_xxx_999", "content": "no confidence", "category": "fact"}
        payload = json.dumps({"axioms": axioms[:2] + [bad_axiom], "week_summary": "Test."})
        gemini = MagicMock()
        gemini.generate = AsyncMock(return_value=payload)
        runner = NightlyEpochRunner(gemini_client=gemini)

        result = _arun(runner.distill(conversations))

        assert len(result["axioms"]) == 2, (
            f"Expected 2 valid axioms after filtering, got {len(result['axioms'])}"
        )

    def test_all_invalid_axioms_returns_empty_list(self):
        """BB4: If all axioms are malformed → result['axioms'] is empty list."""
        conversations = _make_conversations(2)
        # All axioms missing required keys
        bad_axioms = [
            {"content": "no id", "category": "fact"},
            {"id": "x", "category": "strategy"},
            {"content": "no id or category"},
        ]
        payload = json.dumps({"axioms": bad_axioms, "week_summary": "All bad."})
        gemini = MagicMock()
        gemini.generate = AsyncMock(return_value=payload)
        runner = NightlyEpochRunner(gemini_client=gemini)

        result = _arun(runner.distill(conversations))

        assert result["axioms"] == [], (
            f"Expected empty list when all axioms are invalid, got {result['axioms']!r}"
        )

    def test_mixed_valid_invalid_only_valid_returned(self):
        """BB4: Mix of valid and invalid axioms → only valid ones in result."""
        conversations = _make_conversations(2)
        valid = _make_axioms(3)
        invalid = [
            {"id": "bad1"},                        # missing content, category, confidence
            {"content": "x", "confidence": 0.5},   # missing id, category
        ]
        payload = json.dumps({"axioms": valid + invalid, "week_summary": "Mixed."})
        gemini = MagicMock()
        gemini.generate = AsyncMock(return_value=payload)
        runner = NightlyEpochRunner(gemini_client=gemini)

        result = _arun(runner.distill(conversations))

        assert len(result["axioms"]) == 3, (
            f"Expected exactly 3 valid axioms, got {len(result['axioms'])}"
        )
        returned_ids = {ax["id"] for ax in result["axioms"]}
        expected_ids = {ax["id"] for ax in valid}
        assert returned_ids == expected_ids


# ---------------------------------------------------------------------------
# WB1 — Gemini Pro model used (model="gemini-pro" in generate call)
# ---------------------------------------------------------------------------


class TestWB1_GeminiProModelUsed:
    """WB1: distill() must use gemini-pro (not Flash) — quality required."""

    def test_generate_called_with_gemini_pro_model(self):
        """WB1: generate() must be called with model='gemini-pro'."""
        conversations = _make_conversations(2)
        axioms = _make_axioms(3)
        gemini = _make_gemini_client(axioms)
        runner = NightlyEpochRunner(gemini_client=gemini)

        _arun(runner.distill(conversations))

        gemini.generate.assert_called_once()
        _, kwargs = gemini.generate.call_args
        assert kwargs.get("model") == "gemini-pro", (
            f"Expected model='gemini-pro', got {kwargs.get('model')!r}"
        )

    def test_generate_not_called_with_flash_model(self):
        """WB1: Flash model must NOT be used for distillation."""
        conversations = _make_conversations(2)
        axioms = _make_axioms(3)
        gemini = _make_gemini_client(axioms)
        runner = NightlyEpochRunner(gemini_client=gemini)

        _arun(runner.distill(conversations))

        _, kwargs = gemini.generate.call_args
        model_arg = kwargs.get("model", "")
        assert "flash" not in model_arg.lower(), (
            f"distill() must use gemini-pro, not Flash. Got model={model_arg!r}"
        )

    def test_generate_called_exactly_once(self):
        """WB1: generate() must be called exactly once per distill() invocation."""
        conversations = _make_conversations(3)
        axioms = _make_axioms(4)
        gemini = _make_gemini_client(axioms)
        runner = NightlyEpochRunner(gemini_client=gemini)

        _arun(runner.distill(conversations))

        assert gemini.generate.call_count == 1, (
            f"Expected generate() called once, got {gemini.generate.call_count}"
        )


# ---------------------------------------------------------------------------
# WB2 — conversations_json in prompt includes all fields
# ---------------------------------------------------------------------------


class TestWB2_ConversationsJsonInPrompt:
    """WB2: The prompt sent to Gemini must contain the serialised conversations."""

    def test_prompt_contains_conversation_id(self):
        """WB2: Serialised conversations must appear in the Gemini prompt."""
        conversations = _make_conversations(3)
        gemini = _make_gemini_client(_make_axioms(3))
        runner = NightlyEpochRunner(gemini_client=gemini)

        _arun(runner.distill(conversations))

        prompt_arg = gemini.generate.call_args[0][0]
        # Each conversation_id must appear in the prompt
        for conv in conversations:
            assert conv["conversation_id"] in prompt_arg, (
                f"conversation_id {conv['conversation_id']!r} not found in prompt"
            )

    def test_prompt_contains_transcript_raw(self):
        """WB2: transcript_raw values from conversations must appear in the prompt."""
        conversations = _make_conversations(2)
        gemini = _make_gemini_client(_make_axioms(2))
        runner = NightlyEpochRunner(gemini_client=gemini)

        _arun(runner.distill(conversations))

        prompt_arg = gemini.generate.call_args[0][0]
        for conv in conversations:
            # transcript_raw is a unique string per conversation
            snippet = f"message {conversations.index(conv)}"
            assert snippet in prompt_arg, (
                f"Expected transcript snippet {snippet!r} in prompt"
            )

    def test_prompt_is_first_positional_arg_to_generate(self):
        """WB2: Prompt must be passed as the first positional argument to generate()."""
        conversations = _make_conversations(2)
        gemini = _make_gemini_client(_make_axioms(2))
        runner = NightlyEpochRunner(gemini_client=gemini)

        _arun(runner.distill(conversations))

        args, _ = gemini.generate.call_args
        assert len(args) >= 1, "generate() must be called with at least one positional arg"
        assert isinstance(args[0], str), (
            f"First positional arg to generate() must be str (the prompt), got {type(args[0])}"
        )
        # Confirm it looks like the distillation prompt (contains "Genesis Memory Distiller")
        assert "Genesis Memory Distiller" in args[0], (
            "Prompt does not appear to use DISTILLATION_PROMPT template"
        )


# ---------------------------------------------------------------------------
# WB3 — JSON parsed with json.loads; invalid JSON → "Distillation failed"
# ---------------------------------------------------------------------------


class TestWB3_JsonParsingBehavior:
    """WB3: distill() must parse Gemini output with json.loads, not eval."""

    def test_invalid_json_returns_distillation_failed(self):
        """WB3: Gemini returns non-JSON string → week_summary = 'Distillation failed'."""
        conversations = _make_conversations(2)
        gemini = MagicMock()
        gemini.generate = AsyncMock(return_value="This is NOT valid JSON!!")
        runner = NightlyEpochRunner(gemini_client=gemini)

        result = _arun(runner.distill(conversations))

        assert result == {"axioms": [], "week_summary": "Distillation failed"}, (
            f"Expected Distillation failed for invalid JSON, got {result!r}"
        )

    def test_partial_json_returns_distillation_failed(self):
        """WB3: Partial/truncated JSON → 'Distillation failed'."""
        conversations = _make_conversations(2)
        gemini = MagicMock()
        gemini.generate = AsyncMock(return_value='{"axioms": [{"id": "truncated"')
        runner = NightlyEpochRunner(gemini_client=gemini)

        result = _arun(runner.distill(conversations))

        assert result["week_summary"] == "Distillation failed", (
            f"Expected 'Distillation failed' for partial JSON, got {result['week_summary']!r}"
        )

    def test_empty_string_response_returns_distillation_failed(self):
        """WB3: Empty string from Gemini → 'Distillation failed'."""
        conversations = _make_conversations(2)
        gemini = MagicMock()
        gemini.generate = AsyncMock(return_value="")
        runner = NightlyEpochRunner(gemini_client=gemini)

        result = _arun(runner.distill(conversations))

        assert result["week_summary"] == "Distillation failed"

    def test_generate_exception_returns_distillation_failed(self):
        """WB3: generate() raising an exception → 'Distillation failed', not raised."""
        conversations = _make_conversations(2)
        gemini = MagicMock()
        gemini.generate = AsyncMock(side_effect=RuntimeError("Gemini API down"))
        runner = NightlyEpochRunner(gemini_client=gemini)

        # Must not raise
        try:
            result = _arun(runner.distill(conversations))
        except Exception as exc:
            pytest.fail(f"distill() must not propagate exceptions, but raised: {exc}")

        assert result == {"axioms": [], "week_summary": "Distillation failed"}, (
            f"Expected Distillation failed on exception, got {result!r}"
        )


# ---------------------------------------------------------------------------
# WB4 — gemini_client is None → safe empty result
# ---------------------------------------------------------------------------


class TestWB4_NoneGeminiClientGracefulDegradation:
    """WB4: gemini_client=None must return safe empty result without raising."""

    def test_none_client_returns_no_gemini_client_message(self):
        """WB4: distill() with gemini_client=None → 'No Gemini client configured'."""
        conversations = _make_conversations(3)
        runner = NightlyEpochRunner(gemini_client=None)

        result = _arun(runner.distill(conversations))

        assert result == {
            "axioms": [],
            "week_summary": "No Gemini client configured",
        }, f"Unexpected result for None gemini_client: {result!r}"

    def test_none_client_does_not_raise(self):
        """WB4: distill() with None client must not raise any exception."""
        conversations = _make_conversations(2)
        runner = NightlyEpochRunner(gemini_client=None)

        try:
            _arun(runner.distill(conversations))
        except Exception as exc:
            pytest.fail(
                f"distill() with gemini_client=None must not raise, got: {exc}"
            )

    def test_none_client_returns_empty_axioms_list(self):
        """WB4: result['axioms'] must be an empty list (not None) when client is None."""
        runner = NightlyEpochRunner(gemini_client=None)

        result = _arun(runner.distill(_make_conversations(1)))

        assert isinstance(result["axioms"], list), (
            f"Expected list for axioms with None client, got {type(result['axioms'])}"
        )
        assert result["axioms"] == []

    def test_runner_gemini_attribute_stores_client(self):
        """WB4: NightlyEpochRunner stores gemini_client as self.gemini."""
        gemini = _make_gemini_client(_make_axioms(2))
        runner = NightlyEpochRunner(gemini_client=gemini)
        assert runner.gemini is gemini, (
            "gemini_client must be stored as self.gemini on the runner"
        )

    def test_runner_none_gemini_attribute_is_none(self):
        """WB4: When gemini_client=None, runner.gemini must be None."""
        runner = NightlyEpochRunner(gemini_client=None)
        assert runner.gemini is None, (
            f"Expected runner.gemini to be None, got {runner.gemini!r}"
        )


# ---------------------------------------------------------------------------
# Regression — aggregate_week() still works (Story 9.02 unchanged)
# ---------------------------------------------------------------------------


class TestRegression_AggregateWeekUnchanged:
    """Regression: Story 9.02 aggregate_week() must still pass after 9.03 edits."""

    def test_aggregate_week_no_pg_conn_returns_empty_list(self):
        """Regression: aggregate_week() with no pg_conn still returns []."""
        runner = NightlyEpochRunner()  # no args
        result = _arun(runner.aggregate_week())
        assert result == []

    def test_runner_pg_attribute_present(self):
        """Regression: runner.pg attribute must still exist."""
        runner = NightlyEpochRunner(pg_conn=None)
        assert hasattr(runner, "pg"), "runner.pg attribute must exist"
        assert runner.pg is None


# ---------------------------------------------------------------------------
# Run summary
# ---------------------------------------------------------------------------

if __name__ == "__main__":
    import sys as _sys
    result = pytest.main([__file__, "-v", "--tb=short"])
    _sys.exit(result)