helpers.py

"""Helper utility functions for text processing and data manipulation."""

from __future__ import annotations

import json
import re
import textwrap
from pathlib import Path
from typing import Any, Dict, List, Tuple

try:
    from PIL import Image as PILImage
except Exception:  # pragma: no cover
    PILImage = None

from config import BASE_DIR
from models import StateInfo


def _slugify(text: str) -> str:
    """Convert text to a URL-friendly slug."""
    return re.sub(r"[^a-z0-9]+", "_", text.lower()).strip("_")


def _label_from_slug(slug: str) -> str:
    """Convert a slug back to a human-readable label."""
    return " ".join(part.capitalize() for part in slug.replace("-", "_").split("_"))


def _first_heading(markdown: str) -> str:
    """Extract the first heading from markdown text."""
    for line in markdown.splitlines():
        line = line.strip()
        if line.startswith("#"):
            return line.lstrip("#").strip()
    return "Untitled Report"


def _extract_play_headings(markdown: str) -> List[str]:
    """Extract all play headings from markdown text."""
    matches = re.findall(r"^###\s*(Play\s+\d+:[^\n]+)", markdown, flags=re.MULTILINE)
    return [m.strip() for m in matches]


def load_state_catalog(catalog_path: Path) -> List[StateInfo]:
    """Load state information from catalog JSON file."""
    if not catalog_path.exists():
        return []

    payload = json.loads(catalog_path.read_text(encoding="utf-8"))
    states = []
    for item in payload.get("states", []):
        states.append(
            StateInfo(
                slug=item["slug"],
                label=item.get("label", _label_from_slug(item["slug"])),
                source=item.get("source", ""),
                title=item.get("title", "Untitled Report"),
                plays=item.get("plays", []),
            )
        )
    return states


def _state_slug_from_img_url(img_url: str) -> str:
    """Extract state slug from image URL path."""
    parts = Path(img_url).parts
    if len(parts) >= 2 and parts[0] == "charts":
        return _slugify(parts[1])
    return ""


def load_question_bank(path: Path) -> Tuple[List[Dict[str, Any]], str | None]:
    """Load question bank from JSON file."""
    if not path.exists():
        return [], f"Question bank not found: {path}"

    try:
        payload = json.loads(path.read_text(encoding="utf-8"))
    except Exception as exc:
        return [], f"Failed reading question bank: {exc}"

    if not isinstance(payload, list):
        return [], "Question bank must be a JSON list."

    entries: List[Dict[str, Any]] = []
    for idx, item in enumerate(payload, start=1):
        if not isinstance(item, dict):
            continue
        question = str(item.get("question", "")).strip()
        img_url = str(item.get("img_url", "")).strip()
        if not question or not img_url:
            continue
        entries.append(
            {
                "qid": idx,
                "question": question,
                "img_url": img_url,
                "img_path": str(BASE_DIR / img_url),
                "state_slug": _state_slug_from_img_url(img_url),
            }
        )

    if not entries:
        return [], "No valid question entries were found in question bank."
    return entries, None


def _clip(text: str, max_chars: int = 900) -> str:
    """Clip text to maximum character length."""
    text = " ".join(text.split())
    if len(text) <= max_chars:
        return text
    return text[: max_chars - 3] + "..."


def _chunk_section_or_preview(text: str) -> str:
    """Extract section title or preview text from a chunk."""
    for line in text.splitlines():
        stripped = line.strip()
        if not stripped:
            continue
        heading_match = re.match(r"^#{1,6}\s+(.+)$", stripped)
        if heading_match:
            return _clip(heading_match.group(1), max_chars=140)
        break

    collapsed = " ".join(text.split())
    if not collapsed:
        return "No section title available"

    sentence_match = re.search(r"(.+?[.!?])(?:\s|$)", collapsed)
    if sentence_match:
        return _clip(sentence_match.group(1), max_chars=180)
    return _clip(collapsed, max_chars=180)


def _normalize_answer_for_chart_matching(assistant_text: str) -> str:
    """Normalize assistant text for chart matching."""
    lines = []
    for line in assistant_text.splitlines():
        if line.strip().lower().startswith("follow-up question:"):
            continue
        lines.append(line)
    cleaned = " ".join(lines)
    cleaned = re.sub(r"\[(?:S|WEB)\d+\]", " ", cleaned)
    return " ".join(cleaned.split())


def _wrap_caption(text: str, width: int = 64) -> str:
    """Wrap caption text to specified width."""
    return textwrap.fill(" ".join(text.split()), width=width, break_long_words=False, break_on_hyphens=False)


def _question_candidates_for_states(
    state_labels: List[str],
    question_bank: List[Dict[str, Any]],
    states_by_label: Dict[str, StateInfo],
) -> List[Dict[str, Any]]:
    """Filter question candidates based on selected states."""
    if not question_bank:
        return []
    state_slugs = {states_by_label[label].slug for label in state_labels if label in states_by_label}
    if not state_slugs:
        return question_bank
    filtered = [entry for entry in question_bank if entry.get("state_slug") in state_slugs]
    return filtered if filtered else question_bank


def _parse_related_question_numbers(raw_text: str, candidate_count: int, max_items: int) -> List[int]:
    """Parse question numbers from LLM JSON response."""
    parsed: Dict[str, Any] = {}
    try:
        parsed = json.loads(raw_text)
    except Exception:
        match = re.search(r"\{[\s\S]*\}", raw_text)
        if match:
            try:
                parsed = json.loads(match.group(0))
            except Exception:
                parsed = {}

    numbers = parsed.get("question_numbers", []) if isinstance(parsed, dict) else []
    if not isinstance(numbers, list):
        numbers = []

    selected: List[int] = []
    for value in numbers:
        if not isinstance(value, int):
            continue
        if value < 1 or value > candidate_count or value in selected:
            continue
        selected.append(value)
        if len(selected) >= max_items:
            break
    return selected


def _keyword_fallback_questions(
    assistant_text: str,
    candidates: List[Dict[str, Any]],
    max_items: int,
) -> List[Dict[str, Any]]:
    """Find related questions using keyword matching as fallback."""
    query_terms = set(re.findall(r"[a-z]{4,}", assistant_text.lower()))
    scored: List[Tuple[int, int]] = []
    for idx, item in enumerate(candidates):
        question_terms = set(re.findall(r"[a-z]{4,}", str(item.get("question", "")).lower()))
        score = len(query_terms & question_terms)
        if score > 0:
            scored.append((score, idx))
    scored.sort(reverse=True)
    return [candidates[idx] for _, idx in scored[:max_items]]


def extract_play_section(state: StateInfo, play_title: str) -> str:
    """Extract a specific play section from a state's source document."""
    if not state.source:
        return ""

    path = BASE_DIR / state.source
    if not path.exists():
        return ""

    text = path.read_text(encoding="utf-8")
    pattern = rf"^###\s*{re.escape(play_title)}\s*$"
    start_match = re.search(pattern, text, flags=re.MULTILINE)

    if not start_match:
        play_num_match = re.search(r"Play\s+(\d+)", play_title)
        if play_num_match:
            play_num = play_num_match.group(1)
            fallback_pattern = rf"^###\s*Play\s+{play_num}:[^\n]+$"
            start_match = re.search(fallback_pattern, text, flags=re.MULTILINE)

    if not start_match:
        return ""

    remainder = text[start_match.end() :]
    next_play = re.search(r"^###\s*Play\s+\d+:[^\n]+$", remainder, flags=re.MULTILINE)
    end_idx = start_match.end() + next_play.start() if next_play else len(text)
    section = text[start_match.start() : end_idx].strip()
    return section


def _state_scope_labels(state_slugs: List[str], states_by_slug: Dict[str, StateInfo]) -> str:
    """Convert state slugs to comma-separated labels."""
    labels = [states_by_slug[slug].label for slug in state_slugs if slug in states_by_slug]
    return ", ".join(labels)


def related_question_gallery_items(
    assistant_text: str,
    selected_questions: List[Dict[str, Any]],
    max_items: int = 10,
) -> Tuple[List[Tuple[Any, str]], str]:
    """Generate gallery items for related questions with charts."""
    gallery: List[Tuple[Any, str]] = []
    debug_qids: List[int] = []
    debug_images: List[str] = []
    seen_paths = set()

    for item in selected_questions:
        img_path = str(item.get("img_path", ""))
        if not img_path or img_path in seen_paths:
            continue
        if not Path(img_path).exists():
            continue
        seen_paths.add(img_path)
        caption = _wrap_caption(_clip(str(item.get("question", "")), max_chars=240), width=64)
        if PILImage is not None:
            try:
                with PILImage.open(img_path) as img:
                    gallery.append((img.copy(), caption))
            except Exception:
                gallery.append((img_path, caption))
        else:
            gallery.append((img_path, caption))
        qid = item.get("qid")
        if isinstance(qid, int):
            debug_qids.append(qid)
        debug_images.append(Path(img_path).name)

    if debug_qids:
        debug_line = (
            f"`Chart debug (latest response)` qids={debug_qids} "
            f"| images={debug_images} | shown={len(gallery)}"
        )
    else:
        debug_line = "`Chart debug (latest response)` qids=[] | images=[] | shown=0"

    return gallery, debug_line