e2e_test.py

"""
Full end-to-end smoke test for the RICS Report Genius AI system.

Flow:
  1.  Notes-expander unit check (rule-based abbreviations)
  2.  Health check
  3.  Ingest a sample DOCX via /test/ingest
  4.  Create a report linked to that document
  5.  Generate section D (About the property) with messy notes, default ai_level=3 (Balanced)
  6.  Poll until complete, print result + style profile
  7.  Re-generate section D at ai_level=1 (RAG Only) — should still succeed
  8.  Proofread section D
  9.  Enhance section D
 10.  Style-profile roundtrip: all 7 fields must be present
 11.  /extract-notes endpoint: upload the sample DOCX, verify lines returned
 12.  /upload/batch endpoint: upload the sample DOCX as a batch, verify doc row
 13.  /documents/batch-status endpoint: poll the batch-uploaded doc
 14.  /documents/tenant-chunk-summary endpoint: returns chunk count
 15.  Oversized notes-file rejection: expect 413 from /extract-notes
"""
import io
import sys
import time
from pathlib import Path

import requests

from app.generator.notes_expander import expand_notes

BASE    = "http://localhost:8000"
TENANT  = "e2e-smoke-test"
HEADERS = {"X-Tenant-ID": TENANT}
JSON_HEADERS = {**HEADERS, "Content-Type": "application/json"}

# Resolve relative to this script so the test runs on any OS and from any directory.
SAMPLE_FILE = Path(__file__).parent / "sample_survey_doc.docx"


# ─── helpers ──────────────────────────────────────────────────────────────────

def ok(msg: str) -> None:
    print(f"  ✓  {msg}")

def fail(msg: str) -> None:
    print(f"  ✗  {msg}")
    sys.exit(1)

def poll_report(report_id: str, timeout: int = 60) -> str:
    deadline = time.time() + timeout
    while time.time() < deadline:
        r = requests.get(f"{BASE}/reports/{report_id}/status", headers=HEADERS)
        r.raise_for_status()
        status = r.json()["status"]
        print(f"       status={status}")
        if status == "complete":
            return "complete"
        if status == "failed":
            fail(f"Report {report_id} failed")
        time.sleep(3)
    fail(f"Timed out waiting for report {report_id}")
    return "timeout"


# ─── Step 1: Notes expander unit check ────────────────────────────────────────

def test_notes_expander() -> None:
    print("\n=== Step 1: Notes expander (rule-based) ===")
    cases = [
        ("semi det nw3 1965ish 95sqm", "semi-detached"),
        ("roof bad needs work",         "in poor condition"),
        ("mains e mains g connected",   "electricity"),
        ("dg windows ok",               "double-glazed"),
        ("w/o insulation",              "without"),
        ("w/ cavity wall insulation",   "with"),
        ("v bad boiler",                "very poor condition"),
    ]
    for note, expected_fragment in cases:
        expanded = expand_notes([note])[0]
        if expected_fragment.lower() in expanded.lower():
            ok(f"{note!r} → ...{expected_fragment}... ✓  (full: {expanded!r})")
        else:
            fail(f"{note!r} → expected {expected_fragment!r} but got {expanded!r}")


# ─── Step 2: Health check ─────────────────────────────────────────────────────

def test_health() -> None:
    print("\n=== Step 2: Health check ===")
    r = requests.get(f"{BASE}/health")
    r.raise_for_status()
    assert r.json()["status"] == "ok"
    ok("GET /health → ok")


# ─── Step 3: Ingest sample doc ────────────────────────────────────────────────

def ingest_document() -> str:
    print("\n=== Step 3: Test ingest ===")
    payload = {"file_path": str(SAMPLE_FILE), "tenant_id": TENANT}
    r = requests.post(f"{BASE}/test/ingest", json=payload, headers=JSON_HEADERS)
    if r.status_code != 200:
        fail(f"POST /test/ingest failed {r.status_code}: {r.text}")
    doc_id = r.json()["document_id"]
    ok(f"Ingested document_id={doc_id}")
    return doc_id


# ─── Step 4: Create report ────────────────────────────────────────────────────

def create_report(doc_id: str) -> str:
    print("\n=== Step 4: Create report ===")
    r = requests.post(f"{BASE}/reports", headers=HEADERS, params={"document_id": doc_id})
    if r.status_code != 201:
        fail(f"POST /reports failed {r.status_code}: {r.text}")
    report_id = r.json()["report_id"]
    ok(f"Created report_id={report_id}")
    return report_id


# ─── Step 5: Generate section (messy notes, ai_level=3 Balanced) ──────────────
# Uses section D (About the property) — best match for general property description notes.

_TEST_SECTION = "D"
_TEST_BULLETS_D = [
    "semi det nw3 1965ish 95sqm",
    "3 bed 2 storey Victorian terrace",
    "faces NW front elevation",
    "solid brick walls, no cavity",
]
_TEST_BULLETS_E2 = [
    "roof bad needs work",
    "hip roof slate tiles, several slipped",
    "flat felt extension roof cracked",
]


def generate_section(report_id: str, ai_level: int = 3) -> dict:
    print(f"\n=== Step 5: Generate section {_TEST_SECTION} (messy notes, ai_level={ai_level}) ===")
    payload = {
        "template_id": _TEST_SECTION,
        "bullets": _TEST_BULLETS_D,
        "mode": "generate",
        "ai_level": ai_level,
    }
    r = requests.post(
        f"{BASE}/reports/{report_id}/generate",
        headers=JSON_HEADERS,
        json=payload,
    )
    if r.status_code != 200:
        fail(f"POST generate failed {r.status_code}: {r.text}")
    ok(f"Generation queued (ai_level={ai_level}): {r.json()['status']}")

    print("     Polling…")
    poll_report(report_id)

    r2 = requests.get(f"{BASE}/reports/{report_id}/sections", headers=HEADERS)
    r2.raise_for_status()
    sections = r2.json()["sections"]
    if _TEST_SECTION not in sections:
        fail(f"Section {_TEST_SECTION} not found in response")
    section = sections[_TEST_SECTION]
    text = section["text"]
    mode = section.get("mode", "?")
    sp   = section.get("style_profile") or {}
    ok(f"Section {_TEST_SECTION} generated — mode={mode}  ai_level={ai_level}  chars={len(text)}")
    ok(f"Style profile — tone={sp.get('tone','?')}  summary={sp.get('writing_style_summary','?')[:60]}…")
    print(f"\n     ── GENERATED TEXT (first 400 chars) ──")
    print(f"     {text[:400]}")
    return sections


# ─── Step 7: Generate at ai_level=1 (RAG Only) ────────────────────────────────

def generate_rag_only(report_id: str) -> None:
    print(f"\n=== Step 7: Re-generate section {_TEST_SECTION} at ai_level=1 (RAG Only) ===")
    payload = {
        "template_id": _TEST_SECTION,
        "bullets": _TEST_BULLETS_D[:2],
        "mode": "generate",
        "ai_level": 1,
        "force_regenerate": True,
    }
    r = requests.post(
        f"{BASE}/reports/{report_id}/generate",
        headers=JSON_HEADERS,
        json=payload,
    )
    if r.status_code != 200:
        fail(f"POST generate (RAG Only) failed {r.status_code}: {r.text}")
    ok("RAG-only generation queued")
    print("     Polling…")
    poll_report(report_id)
    r2 = requests.get(f"{BASE}/reports/{report_id}/sections", headers=HEADERS)
    r2.raise_for_status()
    section = r2.json()["sections"].get(_TEST_SECTION, {})
    ok(f"RAG-only section {_TEST_SECTION} returned  chars={len(section.get('text',''))}")


# ─── Step 8: Proofread mode ───────────────────────────────────────────────────

def proofread_section(report_id: str) -> None:
    print(f"\n=== Step 8: Proofread section {_TEST_SECTION} ===")
    payload = {
        "template_id": _TEST_SECTION,
        "bullets": _TEST_BULLETS_D[:2],
        "mode": "proofread",
        "ai_level": 3,
    }
    r = requests.post(
        f"{BASE}/reports/{report_id}/generate",
        headers=JSON_HEADERS,
        json=payload,
    )
    if r.status_code != 200:
        fail(f"POST proofread failed {r.status_code}: {r.text}")
    ok("Proofread queued")
    print("     Polling…")
    poll_report(report_id)
    r2 = requests.get(f"{BASE}/reports/{report_id}/sections", headers=HEADERS)
    r2.raise_for_status()
    section = r2.json()["sections"].get(_TEST_SECTION, {})
    text = section.get("text", "")
    mode = section.get("mode", "?")
    ok(f"Section {_TEST_SECTION} after proofread — mode={mode}  chars={len(text)}")
    print(f"\n     ── PROOFREAD TEXT (first 400 chars) ──")
    print(f"     {text[:400]}")


# ─── Step 9: Enhance mode ─────────────────────────────────────────────────────

def enhance_section(report_id: str) -> None:
    print(f"\n=== Step 9: Enhance section {_TEST_SECTION} ===")
    payload = {
        "template_id": _TEST_SECTION,
        "bullets": _TEST_BULLETS_D[:2],
        "mode": "enhance",
        "ai_level": 4,
    }
    r = requests.post(
        f"{BASE}/reports/{report_id}/generate",
        headers=JSON_HEADERS,
        json=payload,
    )
    if r.status_code != 200:
        fail(f"POST enhance failed {r.status_code}: {r.text}")
    ok("Enhance queued (ai_level=4)")
    print("     Polling…")
    poll_report(report_id)
    r2 = requests.get(f"{BASE}/reports/{report_id}/sections", headers=HEADERS)
    r2.raise_for_status()
    section = r2.json()["sections"].get(_TEST_SECTION, {})
    text = section.get("text", "")
    mode = section.get("mode", "?")
    ok(f"Section {_TEST_SECTION} after enhance — mode={mode}  chars={len(text)}")
    print(f"\n     ── ENHANCED TEXT (first 400 chars) ──")
    print(f"     {text[:400]}")


# ─── Step 10: Style profile roundtrip ─────────────────────────────────────────

def test_style_profile_roundtrip(report_id: str) -> None:
    print("\n=== Step 10: Style profile roundtrip ===")
    r = requests.get(f"{BASE}/reports/{report_id}/sections", headers=HEADERS)
    r.raise_for_status()
    section = r.json()["sections"].get(_TEST_SECTION, {})
    sp = section.get("style_profile") or {}
    required_keys = [
        "tone", "formality_level", "avg_sentence_complexity",
        "vocabulary_level", "common_phrases", "structural_patterns",
        "writing_style_summary",
    ]
    missing = [k for k in required_keys if k not in sp]
    if missing:
        fail(f"style_profile missing keys: {missing}")
    ok(f"All 7 style_profile fields present: {list(sp.keys())}")


# ─── Step 11: /extract-notes — valid DOCX ─────────────────────────────────────

def test_extract_notes_docx() -> None:
    print("\n=== Step 11: POST /extract-notes — DOCX upload ===")
    with open(SAMPLE_FILE, "rb") as fh:
        r = requests.post(
            f"{BASE}/extract-notes",
            headers=HEADERS,
            files={"file": ("sample_survey_doc.docx", fh, "application/vnd.openxmlformats-officedocument.wordprocessingml.document")},
        )
    if r.status_code != 200:
        fail(f"/extract-notes failed {r.status_code}: {r.text}")
    data = r.json()
    lines = data.get("lines", [])
    line_count = data.get("line_count", 0)
    if line_count == 0 or not lines:
        fail(f"/extract-notes returned no lines: {data}")
    ok(f"/extract-notes extracted {line_count} line(s) from DOCX")
    print(f"     First 3 lines: {lines[:3]}")


# ─── Step 12: /extract-notes — oversized file rejected (413) ─────────────────

def test_extract_notes_oversized() -> None:
    print("\n=== Step 12: POST /extract-notes — oversized file → 413 ===")
    # Create an in-memory fake .txt that is exactly 10 MB + 1 byte
    big_data = b"x" * (10 * 1024 * 1024 + 1)
    r = requests.post(
        f"{BASE}/extract-notes",
        headers=HEADERS,
        files={"file": ("big.txt", io.BytesIO(big_data), "text/plain")},
    )
    if r.status_code != 413:
        fail(f"Expected 413 for oversized file but got {r.status_code}: {r.text}")
    ok("Oversized notes file correctly rejected with 413")


# ─── Step 13: /upload/batch — batch upload DOCX ───────────────────────────────

def test_batch_upload() -> str:
    print("\n=== Step 13: POST /upload/batch — batch upload ===")
    with open(SAMPLE_FILE, "rb") as fh:
        r = requests.post(
            f"{BASE}/upload/batch",
            headers=HEADERS,
            data={"tenant_id": TENANT},
            files=[("files", ("sample_survey_doc.docx", fh, "application/vnd.openxmlformats-officedocument.wordprocessingml.document"))],
        )
    if r.status_code != 202:
        fail(f"POST /upload/batch failed {r.status_code}: {r.text}")
    data = r.json()
    if data.get("accepted", 0) < 1:
        fail(f"Batch upload accepted 0 files: {data}")
    doc_id = data["items"][0]["document_id"]
    ok(f"Batch upload accepted={data['accepted']}  rejected={data['rejected']}  doc_id={doc_id}")
    return doc_id


# ─── Step 14: /documents/batch-status ─────────────────────────────────────────

def test_batch_status(doc_id: str) -> None:
    print("\n=== Step 14: POST /documents/batch-status ===")
    r = requests.post(
        f"{BASE}/documents/batch-status",
        headers=JSON_HEADERS,
        json={"document_ids": [doc_id, "non-existent-id-000"]},
    )
    if r.status_code != 200:
        fail(f"/documents/batch-status failed {r.status_code}: {r.text}")
    data = r.json()
    items = data.get("items", [])
    found = [i for i in items if i["document_id"] == doc_id]
    not_found = [i for i in items if i["status"] == "not_found"]
    if not found:
        fail(f"batch-status did not return entry for doc_id={doc_id}")
    if not not_found:
        fail("batch-status did not return 'not_found' for non-existent ID")
    ok(f"batch-status: doc_id status={found[0]['status']}  non-existent→not_found ✓")
    ok(f"Aggregates: pending={data['pending']} processing={data['processing']} complete={data['complete']} failed={data['failed']}")


# ─── Step 15: /documents/tenant-chunk-summary ────────────────────────────────

def test_chunk_summary() -> None:
    print("\n=== Step 15: GET /documents/tenant-chunk-summary ===")
    r = requests.get(f"{BASE}/documents/tenant-chunk-summary", headers=HEADERS)
    if r.status_code != 200:
        fail(f"/documents/tenant-chunk-summary failed {r.status_code}: {r.text}")
    data = r.json()
    if "indexed_chunk_count" not in data:
        fail(f"Missing 'indexed_chunk_count' in response: {data}")
    ok(f"tenant-chunk-summary: indexed_chunk_count={data['indexed_chunk_count']}")


# ─── main ─────────────────────────────────────────────────────────────────────

def main() -> None:
    print("=" * 65)
    print("  RICS Report Genius AI — Full End-to-End Smoke Test")
    print("=" * 65)

    test_notes_expander()
    test_health()
    doc_id    = ingest_document()
    report_id = create_report(doc_id)
    generate_section(report_id, ai_level=3)
    generate_rag_only(report_id)
    proofread_section(report_id)
    enhance_section(report_id)
    test_style_profile_roundtrip(report_id)
    test_extract_notes_docx()
    test_extract_notes_oversized()
    batch_doc_id = test_batch_upload()
    test_batch_status(batch_doc_id)
    test_chunk_summary()

    print("\n" + "=" * 65)
    print("  ALL END-TO-END CHECKS PASSED ✓")
    print("=" * 65)


if __name__ == "__main__":
    main()