app/tests/test_generator.py

"""Tests for the LLM adapter, prompt building, and post-processor.

All OpenAI calls are mocked. These tests verify:
1. The prompt context stays within MAX_CONTEXT_TOKENS.
2. The post-processor replaces invented facts with approved missing-info phrases.
3. Verifiable facts are NOT replaced.
"""

import pytest

from app.chunking.splitter import count_tokens
from app.generator.adapter import MockLLMAdapter
from app.generator.postprocess import enforce_verify, strip_verify_tags
from app.generator.prompts import _trim_snippets, build_user_prompt

# ── Prompt building tests ─────────────────────────────────────────────────────

def test_build_user_prompt_contains_skeleton() -> None:
    """The assembled prompt should include the skeleton text."""
    prompt = build_user_prompt(
        skeleton="[Location]: [description].",
        bullets=["Semi-detached, NW3"],
        snippets=[],
        max_context_tokens=400,
    )
    assert "[Location]" in prompt
    assert "DOCUMENT-LEVEL CONTEXT" in prompt
    assert "SECTION-LEVEL CONTEXT" in prompt
    assert "PARAGRAPH-LEVEL EVIDENCE" in prompt
    assert "PROPERTY IDENTITY" in prompt


def test_build_user_prompt_contains_bullets() -> None:
    """Bullets should appear in the assembled prompt."""
    prompt = build_user_prompt(
        skeleton="skeleton",
        bullets=["Semi-detached", "95 sqm"],
        snippets=[],
        max_context_tokens=400,
    )
    assert "Semi-detached" in prompt
    assert "95 sqm" in prompt


def test_trim_snippets_respects_token_limit() -> None:
    """_trim_snippets should not exceed max_tokens."""
    long_snippets = ["word " * 100 for _ in range(10)]
    result = _trim_snippets(long_snippets, max_tokens=400)
    assert count_tokens(result) <= 400 + 20  # small tolerance for label overhead


def test_trim_snippets_empty() -> None:
    """_trim_snippets with empty list should return empty string."""
    assert _trim_snippets([], max_tokens=400) == ""


def test_context_tokens_within_limit() -> None:
    """Total context tokens in the examples block must not exceed MAX_CONTEXT_TOKENS."""
    from app.config import settings

    snippets = ["This is a retrieved snippet. " * 20 for _ in range(5)]
    result = _trim_snippets(snippets, max_tokens=settings.max_context_tokens)
    assert count_tokens(result) <= settings.max_context_tokens + 30


def test_max_output_tokens_for_survey_level_scales_with_tier() -> None:
    """Per-section token floors should reflect whole-report totals.

    Because L1 has far fewer sections overall, its per-section allowance can
    be higher than L2 while still yielding a smaller full report.
    """
    from app.generator.prompts import max_output_tokens_for_survey_level

    l1 = max_output_tokens_for_survey_level(1)
    l2 = max_output_tokens_for_survey_level(2)
    l3 = max_output_tokens_for_survey_level(3)
    # Expected profile: L1 > L2, and L3 > L2.
    assert l1 > l2
    assert l3 > l2
    assert l1 >= 900
    assert l2 >= 650
    assert l3 >= 850
    # None defaults to L3 (matches resolve_generate_system_prompt fallback).
    assert max_output_tokens_for_survey_level(None) == l3


def test_max_context_tokens_for_survey_level_scales_with_tier() -> None:
    """Same shape: L3 needs more retrieved evidence than L1. Without this,
    only ~400 tokens of RAG snippets reach the LLM, and it falls back to
    pre-training defaults — that's where invented addresses come from.
    """
    from app.generator.prompts import max_context_tokens_for_survey_level

    l1 = max_context_tokens_for_survey_level(1)
    l2 = max_context_tokens_for_survey_level(2)
    l3 = max_context_tokens_for_survey_level(3)
    assert l1 <= l2 <= l3
    assert l3 >= 1500


def test_word_target_for_survey_level_matches_full_report_totals() -> None:
    """Word targets are per-section budgets aligned to user-provided full totals."""
    from app.generator.prompts import _word_target_for_survey_level

    min_l1, max_l1 = _word_target_for_survey_level(1)
    min_l2, max_l2 = _word_target_for_survey_level(2)
    min_l3, max_l3 = _word_target_for_survey_level(3)
    # L1 (5 sections total) can be longer per section than L2 (~33 sections).
    assert min_l1 > min_l2
    assert max_l1 > max_l2
    # L3 should stay slightly above L2 per section for diagnostic depth.
    assert min_l3 >= min_l2
    assert max_l3 >= max_l2
    # Sanity windows
    assert 250 <= min_l1 <= 320 and 480 <= max_l1 <= 560
    assert 90 <= min_l2 <= 130 and 160 <= max_l2 <= 220
    assert 120 <= min_l3 <= 170 and 210 <= max_l3 <= 280


# ── Post-processor tests ──────────────────────────────────────────────────────

def test_enforce_verify_passes_through_verified_number() -> None:
    """A number present in the bullets should NOT be replaced."""
    result = enforce_verify(
        text="The property has 95 sqm floor area.",
        bullets=["Floor area is 95 sqm"],
        snippets=[],
    )
    assert "[VERIFY" not in result
    assert "Information not provided in source document." not in result
    assert "We were unable to verify this during inspection." not in result


def test_enforce_verify_wraps_invented_number() -> None:
    """An invented number should be removed from the output."""
    result = enforce_verify(
        text="The property has 120 sqm floor area.",
        bullets=["Floor area is 95 sqm"],
        snippets=[],
    )
    assert "120" not in result
    # enforce_verify collapses the whitespace left where the invented number
    # was removed (\s{2,} -> single space), so no double space survives.
    assert result == "The property has floor area."


def test_enforce_verify_wraps_invented_entity() -> None:
    """An invented named entity should be removed from the output."""
    result = enforce_verify(
        text="The property is located near Hampstead Heath.",
        bullets=["Property in NW3"],
        snippets=[],
    )
    assert "Hampstead Heath" not in result
    assert result == "The property is located near ."


def test_enforce_verify_passes_through_entity_in_snippets() -> None:
    """A named entity present in a retrieved snippet should NOT be replaced."""
    result = enforce_verify(
        text="The property is located near Hampstead Heath.",
        bullets=["Property in NW3"],
        snippets=["This area is adjacent to Hampstead Heath park."],
    )
    assert "Hampstead Heath" in result
    assert "Information not provided in source document." not in result
    assert "We were unable to verify this during inspection." not in result


def test_strip_verify_tags_extracts_values() -> None:
    """strip_verify_tags is a legacy stub — always returns empty list."""
    text = "Area [VERIFY: 120 sqm] near [VERIFY: Canary Wharf]."
    flagged = strip_verify_tags(text)
    assert flagged == []


def test_strip_verify_tags_empty_when_no_tags() -> None:
    """strip_verify_tags on clean text should return an empty list."""
    assert strip_verify_tags("Clean output with no issues.") == []


# ── MockLLMAdapter tests ──────────────────────────────────────────────────────

def test_mock_adapter_uses_override() -> None:
    """MockLLMAdapter with response_override should return exactly that string."""
    adapter = MockLLMAdapter(response_override="FIXTURE output text.")
    result = adapter.generate_section("skeleton", ["bullet"], ["snippet"])
    assert result == "FIXTURE output text."


def test_mock_adapter_default_uses_bullets() -> None:
    """MockLLMAdapter without override should include bullet content in output."""
    adapter = MockLLMAdapter()
    result = adapter.generate_section("skeleton", ["semi-detached house"], [])
    assert "semi-detached house" in result


def test_mock_adapter_includes_style_note_when_profile_given() -> None:
    """generate_section with a style_profile should embed the tone in the output."""
    from app.models.schemas import WritingStyleProfile

    profile = WritingStyleProfile(tone="technical")
    adapter = MockLLMAdapter()
    result = adapter.generate_section("skeleton", ["bullet"], [], style_profile=profile)
    assert "technical" in result


# ── build_user_prompt — style-aware vs plain ──────────────────────────────────

def test_build_user_prompt_plain_when_no_style_profile() -> None:
    """Without a style_profile the plain template must be used (no WRITING STYLE PROFILE header)."""
    prompt = build_user_prompt(
        skeleton="[Location]: [description].",
        bullets=["Semi-detached, NW3"],
        snippets=[],
        max_context_tokens=400,
    )
    assert "WRITING STYLE PROFILE" not in prompt
    assert "[Location]" in prompt


def test_build_user_prompt_includes_style_anchor() -> None:
    """Optional draft paragraph should appear in the user prompt."""
    prompt = build_user_prompt(
        skeleton="[E4]: [content].",
        bullets=["Brick cavity walls"],
        snippets=["Localised cracking to the flank elevation."],
        max_context_tokens=400,
        style_anchor="The external walls are of cavity brick construction with some localised cracking.",
    )
    assert "SURVEYOR'S DRAFT PARAGRAPH" in prompt
    assert "cavity brick" in prompt


def test_build_user_prompt_style_aware_when_profile_given() -> None:
    """With a style_profile the style-aware template must inject profile fields."""
    from app.models.schemas import WritingStyleProfile

    profile = WritingStyleProfile(
        tone="semi-formal",
        formality_level="academic",
        writing_style_summary="Concise academic prose.",
    )
    prompt = build_user_prompt(
        skeleton="[Location]: [description].",
        bullets=["Semi-detached, NW3"],
        snippets=[],
        max_context_tokens=400,
        style_profile=profile,
    )
    assert "polished RICS report section text of" in prompt  # word target injected
    assert "WRITING STYLE PROFILE" in prompt
    assert "semi-formal" in prompt
    assert "academic" in prompt
    assert "Concise academic prose." in prompt


# ── build_proofread_prompt ────────────────────────────────────────────────────

def test_build_proofread_prompt_contains_text_and_bullets() -> None:
    """The proofread prompt must contain the text to review and the bullets."""
    from app.generator.prompts import build_proofread_prompt

    prompt = build_proofread_prompt(
        text="The property is a semi-detached house.",
        bullets=["Semi-detached, NW3"],
    )
    assert "semi-detached house" in prompt
    assert "Semi-detached, NW3" in prompt


def test_build_proofread_prompt_uses_style_profile() -> None:
    """When a style_profile is supplied its tone must appear in the proofread prompt."""
    from app.generator.prompts import build_proofread_prompt
    from app.models.schemas import WritingStyleProfile

    profile = WritingStyleProfile(tone="casual", writing_style_summary="Conversational style.")
    prompt = build_proofread_prompt(
        text="Some text.",
        bullets=["fact one"],
        style_profile=profile,
    )
    assert "casual" in prompt
    assert "Conversational style." in prompt


def test_build_proofread_prompt_falls_back_to_mock_profile() -> None:
    """Without a style_profile the mock profile defaults should still appear."""
    from app.generator.prompts import build_proofread_prompt

    prompt = build_proofread_prompt(text="Some text.", bullets=[])
    assert "formal" in prompt  # mock profile default tone


# ── build_enhance_prompt ──────────────────────────────────────────────────────

def test_build_enhance_prompt_contains_current_text() -> None:
    """The enhance prompt must contain the existing section text."""
    from app.generator.prompts import build_enhance_prompt

    prompt = build_enhance_prompt(
        text="The roof is in poor condition.",
        bullets=["Roof: felt tiles, poor condition"],
        snippets=["Felt tiles typically last 15–20 years."],
        max_context_tokens=400,
    )
    assert "roof is in poor condition" in prompt
    assert "felt tiles" in prompt.lower()


def test_build_enhance_prompt_includes_evidence_snippets() -> None:
    """Retrieved evidence snippets must appear in the enhance prompt."""
    from app.generator.prompts import build_enhance_prompt

    prompt = build_enhance_prompt(
        text="Short text.",
        bullets=["bullet"],
        snippets=["UNIQUE_SNIPPET_MARKER structural survey evidence"],
        max_context_tokens=400,
    )
    assert "UNIQUE_SNIPPET_MARKER" in prompt


def test_build_enhance_prompt_uses_style_profile() -> None:
    """A supplied style_profile must inject its fields into the enhance prompt."""
    from app.generator.prompts import build_enhance_prompt
    from app.models.schemas import WritingStyleProfile

    profile = WritingStyleProfile(tone="technical", writing_style_summary="Expert surveyor voice.")
    prompt = build_enhance_prompt(
        text="Existing text.",
        bullets=["fact"],
        snippets=[],
        max_context_tokens=400,
        style_profile=profile,
    )
    assert "technical" in prompt
    assert "Expert surveyor voice." in prompt


# ── MockLLMAdapter — proofread and enhance ────────────────────────────────────

def test_mock_proofread_contains_notes_separator() -> None:
    """MockLLMAdapter.proofread must include the ---NOTES--- separator."""
    adapter = MockLLMAdapter()
    result = adapter.proofread("Original text.", bullets=["fact"])
    assert "---NOTES---" in result
    assert "Original text." in result


def test_mock_enhance_appends_no_key_notice() -> None:
    """MockLLMAdapter.enhance must append a notice about the missing API key."""
    adapter = MockLLMAdapter()
    result = adapter.enhance("Existing text.", bullets=["fact"], snippets=["snippet"])
    assert "Existing text." in result
    assert "OPENAI_API_KEY" in result


def test_mock_enhance_notes_snippet_count() -> None:
    """MockLLMAdapter.enhance should mention the number of retrieved snippets."""
    adapter = MockLLMAdapter()
    result = adapter.enhance("text", bullets=[], snippets=["s1", "s2", "s3"])
    assert "3" in result


# ── GenerateRequest schema validators ─────────────────────────────────────────

def test_generate_request_rejects_invalid_template_id() -> None:
    """template_id must be a valid RICS section code (union of all product tiers)."""
    from pydantic import ValidationError

    from app.models.schemas import GenerateRequest

    with pytest.raises(ValidationError, match="not a valid RICS section code"):
        GenerateRequest(template_id="B1", bullets=["a fact"])


def test_generate_request_accepts_valid_template_id() -> None:
    """A correct RICS section code (e.g. E4) must pass validation."""
    from app.models.schemas import GenerateRequest

    req = GenerateRequest(template_id="E4", bullets=["Solid brick walls, DPC visible"])
    assert req.template_id == "E4"


def test_generate_request_rejects_blank_template_id() -> None:
    """An empty or whitespace-only template_id must be rejected."""
    from pydantic import ValidationError

    from app.models.schemas import GenerateRequest

    with pytest.raises(ValidationError):
        GenerateRequest(template_id="  ", bullets=["a fact"])


def test_generate_request_allows_empty_bullets() -> None:
    """Empty bullets are allowed (backend may persist a blank section)."""
    from app.models.schemas import GenerateRequest

    req = GenerateRequest(template_id="D", bullets=["   ", ""])
    assert req.bullets == []