app/generator/notes_expander.py

"""Notes expander — transforms raw inspection field notes into professional statements.

Users typically enter short, informal, incomplete bullet notes such as:
    "roof bad maybe 10yr felt tiles needs work"
    "semi det nw3 1965ish 95sqm"
    "gas elec mains ok heating old boiler"

This module pre-processes those notes into clearer, structured fact statements
*before* they reach the main generation LLM, giving the generator higher-quality
input to work with.

Rules observed in every path:
- Exact numbers, dates, addresses, and postcodes are never altered.
- Vague condition words are mapped to professional RICS equivalents.
- Common abbreviations are expanded.
- Any content that is inferred rather than stated is tagged [INFERRED] so the
  downstream enforce_verify pass can flag it for the user to confirm.
- When no OpenAI key is configured, a fast rule-based mock path is used so the
  pipeline stays fully functional offline.
"""

import asyncio
import json
import logging
import re

logger = logging.getLogger(__name__)

_EXPAND_BATCH_SIZE = 28


def _expand_tokens_for_batch(batch_len: int) -> int:
    return min(4096, max(400, batch_len * 72))


async def _expand_notes_batch_async(
    bullets: list[str],
    *,
    section_code: str,
    chat_model: str,
) -> list[str]:
    """Expand one batch; rule-based fallback per bullet on failure."""
    from app.config import settings
    from app.llm.openai_chat import chat_completions_create

    system_prompt = (
        "You are a UK RICS surveyor assistant. Expand raw inspector field notes into "
        "professional fact statements suitable for a Home Survey report.\n"
        "Rules:\n"
        "1. Preserve every number, measurement, date, postcode, and proper noun exactly.\n"
        "2. Expand abbreviations (e.g. det → detached, sqm → square metres).\n"
        "3. Map vague condition words to professional equivalents.\n"
        "4. Tag inferred content with [INFERRED].\n"
        "5. Do NOT invent facts not implied by the notes.\n"
        "6. Return a JSON array of strings — one expanded statement per input bullet, same order.\n"
        "7. Output ONLY the JSON array.\n"
        "8. STRICT BRITISH ENGLISH ONLY."
    )
    section_context = f" for RICS section {section_code}" if section_code else ""
    user_prompt = (
        f"Expand these raw inspector field notes{section_context} into "
        f"professional RICS fact statements:\n\n"
        + "\n".join(f"- {b}" for b in bullets)
    )
    try:
        raw = await chat_completions_create(
            messages=[
                {"role": "system", "content": system_prompt},
                {"role": "user", "content": user_prompt},
            ],
            model=chat_model or settings.chat_model,
            max_tokens=_expand_tokens_for_batch(len(bullets)),
            temperature=0.1,
            phase="notes_expand",
            section_id=section_code or None,
        )
        expanded: list[str] = json.loads(raw or "[]")
        if isinstance(expanded, list) and len(expanded) == len(bullets):
            return [str(e) for e in expanded]
    except json.JSONDecodeError as exc:
        logger.warning("Notes expander batch non-JSON (%s) — rule-based fallback", exc)
    except Exception as exc:
        logger.warning("Notes expander batch failed (%s) — rule-based fallback", exc)
    return [_rule_based_expand(b) for b in bullets]

# ── Abbreviation dictionary ────────────────────────────────────────────────────

# Multi-word service phrases MUST come before single-letter direction
# abbreviations so "mains e" is never corrupted to "mains east" by the
# single-letter "e" rule firing first.  Longer / more-specific patterns
# are always listed before their shorter substrings.
_ABBREV: dict[str, str] = {
    # ── services (multi-word first — must precede single-letter directions) ──
    "mains g": "mains gas",
    "mains e": "mains electricity",
    "mains w": "mains water",
    "elec.": "electricity",
    "elec":  "electricity",
    "c/h":   "central heating",
    "ch":    "central heating",
    "u/flr": "underfloor",
    "u/f":   "underfloor",
    # ── property types ────────────────────────────────────────────────────────
    "semi det":  "semi-detached",
    "semi-det":  "semi-detached",
    "det":       "detached",
    "terr":      "terraced",
    "ter":       "terraced",
    "flat":      "flat",
    "maisonette": "maisonette",
    "bungalow":  "bungalow",
    # ── construction ─────────────────────────────────────────────────────────
    "rcc":  "reinforced concrete",
    "rc":   "reinforced concrete",
    "br":   "brick",
    "bk":   "brick",
    "upvc": "uPVC",
    "t/g":  "triple-glazed",
    "d/g":  "double-glazed",
    "s/g":  "single-glazed",
    "tg":   "triple-glazed",
    "dg":   "double-glazed",
    "sg":   "single-glazed",
    "c/w":  "complete with",
    "w/o":  "without",
    "w/":   "with",
    # ── compass directions (single letters last — after all multi-word rules) ─
    "ne": "north-east", "nw": "north-west",
    "se": "south-east", "sw": "south-west",
    "n":  "north", "s":  "south", "e":  "east", "w":  "west",
    # ── area / measurement ────────────────────────────────────────────────────
    "sq.m": "sq m", "sqm": "sq m", "m2": "sq m",
    "sqft": "sq ft",
    # ── approximate age markers ───────────────────────────────────────────────
    "ish": "", "approx": "approximately", "c.": "circa", "ca.": "circa",
    # ── condition shorthand ───────────────────────────────────────────────────
    # "poor" MUST come first: "v bad", "vbad", and "pf" all expand to phrases
    # containing the word "poor".  If "poor" were processed after any of those,
    # the _ABBREV loop (which is inherently cascading) would re-expand "poor"
    # inside the already-expanded text and produce garbled output like
    # "in very in poor condition condition and requiring urgent attention".
    "poor": "in poor condition",
    "v gd": "in very good condition",
    "vgd":  "in very good condition",
    "v bad": "in very poor condition and requiring urgent attention",
    "vbad":  "in very poor condition and requiring urgent attention",
    "ok":   "in satisfactory condition",
    "gd":   "in good condition",
    "fair": "in fair condition",
    "pf":   "in poor/fair condition",
    "bad":  "in poor condition and requiring attention",
    "nr":   "not readily visible",
    "nv":   "not visible",
    "ni":   "not inspected",
    "nm":   "not measured",
    # ── time / age ────────────────────────────────────────────────────────────
    "yrs": "years old", "yr": "years old",
    # ── general ───────────────────────────────────────────────────────────────
    "re:": "regarding",
    "fyi": "",
    "poss": "possibly",
    "prob": "probably",
    "nec":  "necessary",
    "reqs": "requires",
    "req":  "required",
    "recs": "recommends",
    "rec":  "recommended",
    "asap": "urgently",
    "immed": "immediately",
}

# ── Condition-word mapping ─────────────────────────────────────────────────────

_CONDITION_WORDS: dict[str, str] = {
    "bad":       "in poor condition, requiring attention",
    "terrible":  "in very poor condition, requiring urgent attention",
    "awful":     "in very poor condition and beyond reasonable repair without major work",
    "horrible":  "in very poor condition",
    "cracked":   "exhibiting cracking",
    "cracking":  "exhibiting cracking which should be monitored",
    "damp":      "showing evidence of dampness",
    "wet":       "showing evidence of water ingress",
    "leaking":   "exhibiting signs of water ingress",
    "rotten":    "exhibiting decay/rot requiring repair",
    "rotting":   "exhibiting active decay requiring repair",
    "loose":     "loose and requiring re-fixing",
    "missing":   "with some elements missing, requiring reinstatement",
    "rusty":     "exhibiting corrosion/rust",
    "stained":   "exhibiting staining, the cause of which should be investigated",
    "peeling":   "with peeling/flaking finish requiring redecoration",
    "tired":     "in aged condition requiring general maintenance",
    "needs work": "requires attention and further investigation",
    "needs repair": "requires repair",
    # NOTE: "urgent" is intentionally absent here.  The word appears inside
    # several _ABBREV expansions (e.g. "v bad" → "…requiring urgent attention").
    # Including it here would cause the condition-words pass to re-expand the
    # already-expanded output, producing garbled repetition like
    # "requiring requiring urgent attention attention".
}


def _rule_based_expand(bullet: str) -> str:
    """Apply a fast rule-based expansion to a single bullet string.

    Expands common abbreviations and maps condition words to professional
    equivalents without making any external API calls.

    Args:
        bullet: Raw single-line inspector note.

    Returns:
        Expanded string (may still contain gaps — the LLM will fill those).
    """
    text = bullet.strip()

    # Expand abbreviations.  The lookbehind uses (?<![a-zA-Z]) rather than
    # (?<!\w) so that unit/age suffixes like "95sqm" or "10yr" are still caught
    # (the digit before the abbrev is not a letter, just not a word boundary).
    # When the match is immediately preceded by a digit a space is inserted so
    # that "95sqm" → "95 sq m" and "10yr" → "10 years old", not "95sq m" / "10years old".
    for abbrev, expansion in _ABBREV.items():
        pattern = r"(?<![a-zA-Z])" + re.escape(abbrev) + r"(?!\w)"
        src = text

        def _abbrev_repl(m: re.Match[str], exp: str = expansion, hay: str = src) -> str:
            prefix = " " if m.start() > 0 and hay[m.start() - 1].isdigit() else ""
            return prefix + exp

        text = re.sub(pattern, _abbrev_repl, text, flags=re.IGNORECASE)

    # Normalise multiple spaces
    text = re.sub(r"  +", " ", text).strip()

    # Expand condition words in a single non-cascading pass: build a combined
    # pattern and replace all matches at once using a dispatch dict.  This
    # prevents "cracked" → "exhibiting cracking" from then re-triggering the
    # "cracking" rule, and stops "terrible" → "...urgent..." from re-firing
    # the "urgent" rule on its own output.
    sorted_conditions = sorted(_CONDITION_WORDS.keys(), key=len, reverse=True)
    combined = re.compile(
        r"(?<!\w)(" + "|".join(re.escape(k) for k in sorted_conditions) + r")(?!\w)",
        re.IGNORECASE,
    )
    text = combined.sub(lambda m: _CONDITION_WORDS[m.group(0).lower()], text)

    return text


def _llm_expand_raw(
    *,
    system_prompt: str,
    user_prompt: str,
    openai_api_key: str,
    chat_model: str,
    max_tokens: int = 600,
) -> str:
    from openai import OpenAI

    client = OpenAI(api_key=openai_api_key)
    response = client.chat.completions.create(
        model=chat_model,
        messages=[
            {"role": "system", "content": system_prompt},
            {"role": "user", "content": user_prompt},
        ],
        max_tokens=max_tokens,
        temperature=0.1,
    )
    return (response.choices[0].message.content or "").strip()


async def expand_notes_async(
    bullets: list[str],
    section_code: str = "",
    openai_api_key: str = "",
    chat_model: str = "gpt-4o-mini",
) -> list[str]:
    """Async notes expansion using the throttled OpenAI chat helper."""
    if not bullets:
        return []
    if not (openai_api_key or "").strip():
        return [_rule_based_expand(b) for b in bullets]

    from app.config import settings

    model = chat_model or settings.chat_model
    if len(bullets) <= _EXPAND_BATCH_SIZE:
        return await _expand_notes_batch_async(
            bullets, section_code=section_code, chat_model=model
        )
    # Multiple batches are independent — run them concurrently (bounded by the
    # global LLM throttle) instead of awaiting each in series. Order is
    # preserved by gathering in batch order.
    batches = [
        bullets[i : i + _EXPAND_BATCH_SIZE]
        for i in range(0, len(bullets), _EXPAND_BATCH_SIZE)
    ]
    batch_results = await asyncio.gather(
        *(
            _expand_notes_batch_async(
                batch, section_code=section_code, chat_model=model
            )
            for batch in batches
        )
    )
    out: list[str] = []
    for expanded in batch_results:
        out.extend(expanded)
    return out


def expand_notes(
    bullets: list[str],
    section_code: str = "",
    openai_api_key: str = "",
    chat_model: str = "gpt-4o-mini",
) -> list[str]:
    """Expand raw inspector notes into professional, structured fact statements.

    When ``openai_api_key`` is provided the LLM path is used; otherwise a fast
    rule-based mock path is applied so the pipeline functions offline.

    The expanded bullets are passed to the generation LLM as richer input.
    Exact numbers, postcodes, and dates are never modified by this function.

    Args:
        bullets: Raw bullet strings as entered by the user (may be informal).
        section_code: RICS section code (e.g. ``"E2"`` for Roof coverings), used for context.
        openai_api_key: OpenAI API key; empty string triggers mock fallback.
        chat_model: Chat model to use for LLM-based expansion.

    Returns:
        List of expanded bullet strings, one per original bullet.
        Returns the original list unchanged on error or empty input.
    """
    if not bullets:
        return bullets

    if not openai_api_key:
        logger.debug("No OpenAI key — using rule-based note expansion")
        return [_rule_based_expand(b) for b in bullets]

    system_prompt = (
        "You are an expert RICS surveyor acting as an inspection notes interpreter. "
        "Your task is to expand informal, shorthand inspection notes into clear, "
        "professional fact statements ready for a formal RICS survey report. "
        "\n\nRules you MUST follow:\n"
        "1. Preserve ALL exact numbers, measurements, dates, postcodes, and addresses exactly as given.\n"
        "2. Expand abbreviations (e.g. 'semi det' → 'semi-detached', 'dg' → 'double-glazed').\n"
        "3. Map informal condition words to professional RICS equivalents "
        "   (e.g. 'bad' → 'in poor condition requiring attention').\n"
        "4. Fill in standard professional context for clearly incomplete observations "
        "   but tag anything inferred with [INFERRED].\n"
        "5. Do NOT invent measurements, costs, or specific facts not implied by the notes.\n"
        "6. Return a JSON array of strings — one expanded statement per input bullet.\n"
        "7. Maintain the same order as the input.\n"
        "8. Output ONLY the JSON array, no commentary, no markdown fences.\n"
        "9. STRICT BRITISH ENGLISH ONLY: use British spellings at all times — "
        "   colour (not color), centre (not center), storey/storeys (not story/stories for floors), "
        "   metre/metres (not meter/meters), aluminium (not aluminum), mould (not mold), "
        "   analyse (not analyze), organise (not organize), grey (not gray), "
        "   draught (not draft for air), kerb (not curb), -ise/-isation suffixes (not -ize/-ization). "
        "   Never use American English spellings."
    )

    section_context = f" for RICS section {section_code}" if section_code else ""

    if len(bullets) <= _EXPAND_BATCH_SIZE:
        batches = [bullets]
    else:
        batches = [
            bullets[i : i + _EXPAND_BATCH_SIZE]
            for i in range(0, len(bullets), _EXPAND_BATCH_SIZE)
        ]
    out: list[str] = []
    for batch in batches:
        batch_prompt = (
            f"Expand these raw inspector field notes{section_context} into "
            f"professional RICS fact statements:\n\n"
            + "\n".join(f"- {b}" for b in batch)
        )
        try:
            raw = _llm_expand_raw(
                system_prompt=system_prompt,
                user_prompt=batch_prompt,
                openai_api_key=openai_api_key,
                chat_model=chat_model,
                max_tokens=_expand_tokens_for_batch(len(batch)),
            )
            expanded: list[str] = json.loads(raw)
            if isinstance(expanded, list) and len(expanded) == len(batch):
                out.extend(str(e) for e in expanded)
                continue
        except json.JSONDecodeError as exc:
            logger.warning("Notes expander batch non-JSON (%s) — rule-based fallback", exc)
        except Exception as exc:
            logger.warning("Notes expander batch failed (%s) — rule-based fallback", exc)
        out.extend(_rule_based_expand(b) for b in batch)
    logger.info(
        "Notes expansion complete: %d bullets expanded for section=%s",
        len(bullets),
        section_code,
    )
    return out