modules/llm_parser.py

"""Heuristic parser for free-form hotel search requests."""

from __future__ import annotations

import re
from copy import deepcopy
from datetime import date
from typing import Any

from .prompts import AMBIGUOUS_DATE_PHRASES, AMENITY_SYNONYMS, EXTRACTION_SCHEMA, REQUIRED_HINT_WORDS

ISO_DATE_RE = re.compile(r"\b\d{4}-\d{2}-\d{2}\b")


def _to_iso_if_valid(value: str) -> str | None:
    # Validate and normalize ISO date strings to YYYY-MM-DD.
    try:
        parsed = date.fromisoformat(value)
        return parsed.isoformat()
    except ValueError:
        return None


def _extract_location(text: str) -> str | None:
    # Capture phrase after "in" and stop near known separators.
    location_match = re.search(
        r"\bin\s+([a-zA-Z][a-zA-Z\s,\-]+?)(?=\s+(?:from|check[- ]?in|on|under|below|budget|with|for|adults|children|party|$))",
        text,
        flags=re.IGNORECASE,
    )
    if location_match:
        return location_match.group(1).strip(" ,.")
    return None


def _extract_budget(text: str) -> float | None:
    # Step 1: Prefer explicit budget cue words (under, below, max, budget).
    budget_match = re.search(r"(?:under|below|max|budget)\s*\$?\s*(\d+(?:\.\d+)?)", text, flags=re.IGNORECASE)
    if budget_match:
        return float(budget_match.group(1))
    # Step 2: Fallback to the first plain dollar amount.
    dollar_match = re.search(r"\$\s*(\d+(?:\.\d+)?)", text)
    if dollar_match:
        return float(dollar_match.group(1))
    return None


def _extract_party_fields(text: str) -> dict[str, int | None]:
    # Parse optional occupancy details to enrich search context.
    adults = None
    children = None
    party_size = None

    adults_match = re.search(r"(\d+)\s+adults?", text, flags=re.IGNORECASE)
    children_match = re.search(r"(\d+)\s+children?", text, flags=re.IGNORECASE)
    party_match = re.search(r"(?:party size|for)\s*(\d+)", text, flags=re.IGNORECASE)

    if adults_match:
        adults = int(adults_match.group(1))
    if children_match:
        children = int(children_match.group(1))
    if party_match:
        party_size = int(party_match.group(1))
    elif adults is not None or children is not None:
        party_size = (adults or 0) + (children or 0)

    return {"party_size": party_size, "adults": adults, "children": children}


def _normalize_amenity(token: str) -> str | None:
    # Canonicalize amenity wording to stabilize matching/scoring.
    clean = token.strip().lower()
    if not clean:
        return None
    return AMENITY_SYNONYMS.get(clean, clean)


def _extract_amenities(text: str) -> tuple[list[str], list[str]]:
    # Step 1: Initialize amenity buckets.
    lowered = text.lower()
    required: set[str] = set()
    preferred: set[str] = set()

    # Step 2: Choose required-mode if strong requirement hints are present.
    required_mode = any(hint in lowered for hint in REQUIRED_HINT_WORDS)
    for alias, canonical in AMENITY_SYNONYMS.items():
        if alias in lowered:
            if required_mode:
                required.add(canonical)
            else:
                preferred.add(canonical)

    # If user explicitly uses "preferred"/"nice to have", push matching amenities there.
    # Step 3: If "preferred" language appears, force matched amenities into preferred bucket.
    if "preferred" in lowered or "nice to have" in lowered:
        for alias, canonical in AMENITY_SYNONYMS.items():
            if alias in lowered:
                preferred.add(canonical)
                required.discard(canonical)

    return sorted(required), sorted(preferred)


def parse_user_request(user_request: str) -> dict[str, Any]:
    # Step 1: Start from schema defaults and prepare validation containers.
    payload: dict[str, Any] = deepcopy(EXTRACTION_SCHEMA)
    clarifications: list[str] = []
    errors: list[str] = []

    # Step 2: Short-circuit on empty input with a user-friendly clarification.
    text = (user_request or "").strip()
    if not text:
        clarifications.append("Please describe what hotel stay you want, including city and exact dates in YYYY-MM-DD.")
        return {"payload": payload, "clarifications": clarifications, "errors": errors}

    # Step 3: Flag relative/ambiguous date phrases for clarification.
    lowered = text.lower()
    if any(phrase in lowered for phrase in AMBIGUOUS_DATE_PHRASES):
        clarifications.append("Please provide exact check-in and check-out dates in YYYY-MM-DD (relative dates like 'next weekend' are ambiguous).")

    # Step 4: Extract scalar and list fields from free-form text.
    payload["location"] = _extract_location(text)
    payload["budget_max_usd"] = _extract_budget(text)
    payload.update(_extract_party_fields(text))

    required, preferred = _extract_amenities(text)
    payload["required_amenities"] = required
    payload["preferred_amenities"] = preferred

    # Step 5: Parse explicit ISO dates if present.
    found_dates = ISO_DATE_RE.findall(text)
    if len(found_dates) >= 2:
        check_in_iso = _to_iso_if_valid(found_dates[0])
        check_out_iso = _to_iso_if_valid(found_dates[1])
        payload["check_in_date"] = check_in_iso
        payload["check_out_date"] = check_out_iso
        if check_in_iso is None or check_out_iso is None:
            clarifications.append("I found date-like text, but it is not valid ISO format. Please use YYYY-MM-DD.")
    else:
        clarifications.append("Please provide check-in and check-out dates in YYYY-MM-DD.")

    # Step 6: Require destination location.
    if not payload["location"]:
        clarifications.append("Please provide the destination location (city or area).")

    # Step 7: Validate date order only when both dates exist.
    check_in = payload.get("check_in_date")
    check_out = payload.get("check_out_date")
    if check_in and check_out and check_out <= check_in:
        errors.append("Invalid date range: check_out_date must be later than check_in_date.")

    return {"payload": payload, "clarifications": clarifications, "errors": errors}