src/analyzer/context_builder.py

"""
context_builder.py — prepares text context for the LLM summarizer

Takes structured JSON grant records + optional past winners and builds a single
text block for each grant, suitable as LLM input.

Responsibilities:
- Extract key text fields from each grant JSON (title, description, sections)
- Summarize/flatten them into a readable context string
- Optionally include supporting documents (HTML sections + PDF text extracts)
- Optionally include a few relevant past-winner snippets (if any exist)

Public API
----------
build_context(grant: dict, past_winners: list[dict] | None = None, include_supporting: bool = False) -> str
build_context_with_supporting(grant: dict, k: int = 5, past_winners: list[dict] | None = None) -> str
"""
from __future__ import annotations

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

# ----------------------------- Text utilities ---------------------------------

def _clean(s: Any) -> str:
    return re.sub(r"\s+", " ", str(s or "")).strip()


def _maybe(k: str, v: Any) -> str:
    if not v:
        return ""
    return f"{k}: {_clean(v)}\n"

# ----------------------------- Context builder --------------------------------

def build_context(grant: Dict[str, Any], past_winners: Optional[List[Dict[str, Any]]] = None) -> str:
    """
    Flatten a grant JSON object + optional past winners into a readable context string.

    CRITICAL: All URLs are provided in FULL form (https://...), never partial paths.
    The model receives complete, actionable URLs directly from the snapshot JSON.
    """
    lines: List[str] = []

    # --- Basic info ---
    title = grant.get("title") or grant.get("name") or grant.get("competition_title")
    if title:
        lines.append(f"TITLE: {_clean(title)}")

    # --- FULL URL: Always provide complete https:// URLs, never relative paths ---
    url = grant.get("url") or grant.get("link")
    if url:
        lines.append(f"URL: {url}")  # Example: "https://apply-for-innovation-funding.service.gov.uk/competition/2276/overview/..."

    # --- Funding information ---
    funding = grant.get("funding") or {}
    if isinstance(funding, dict):
        max_funding = funding.get("max")
        if max_funding:
            lines.append(f"FUNDING: Up to £{max_funding:,}")
        min_funding = funding.get("min")
        if min_funding and min_funding != max_funding:
            lines.append(f"MINIMUM FUNDING: £{min_funding:,}")
    else:
        funding_text = grant.get("funding_amount") or grant.get("amount") or grant.get("funding_rate")
        if funding_text:
            lines.append(f"FUNDING: {_clean(funding_text)}")

    # --- Dates ---
    deadline = grant.get("close_date") or grant.get("deadline")
    if deadline:
        lines.append(f"DEADLINE: {_clean(deadline)}")

    open_date = grant.get("open_date")
    if open_date:
        lines.append(f"OPENS: {_clean(open_date)}")

    # --- Core text sections from snapshot ---
    sections = grant.get("sections") or {}
    if sections:
        # Process snapshot sections in order: summary, eligibility, scope, dates, how_to_apply
        section_order = [
            "summary_raw", "eligibility_raw", "scope_raw",
            "dates_raw", "how_to_apply_raw", "supporting_information_raw"
        ]
        for section_key in section_order:
            v = sections.get(section_key)
            if v and _clean(v):  # Only include non-empty sections
                section_name = section_key.replace("_raw", "").replace("_", " ").upper()
                lines.append(f"\n{section_name}:\n{_clean(v)}")
    else:
        # Fallback to common text fields
        desc = grant.get("description") or grant.get("summary") or grant.get("scope")
        if desc:
            lines.append(f"\nDESCRIPTION:\n{_clean(desc)}")

    # --- Optional extras ---
    eligibility = grant.get("eligibility")
    if eligibility:
        lines.append(f"\nELIGIBILITY:\n{_clean(eligibility)}")

    scope = grant.get("scope")
    if scope:
        lines.append(f"\nSCOPE:\n{_clean(scope)}")

    # --- Contact Information ---
    lines.append("\n--- CONTACT INFORMATION ---")
    lines.append("Email: support@iuk.ukri.org")
    lines.append("Phone: 0300 321 4357")
    lines.append("Hours: 9am-12pm, 2pm-5pm, Monday-Friday (excluding bank holidays)")

    # --- Past winners summary ---
    if past_winners:
        lines.append("\n--- RELATED PAST WINNERS ---")
        for w in past_winners[:5]:  # limit to top 5 to avoid overloading tokens
            snippet_parts: List[str] = []
            snippet_parts.append(_maybe("Project", w.get("project_title")))
            snippet_parts.append(_maybe("Organisation", w.get("lead_org")))
            snippet_parts.append(_maybe("Award", w.get("award_amount")))
            snippet_parts.append(_maybe("Competition", w.get("competition")))
            abs_ = _clean(w.get("abstract"))
            if abs_:
                snippet_parts.append(f"Abstract: {abs_[:400]}{'…' if len(abs_)>400 else ''}\n")
            lines.append("".join(snippet_parts))

    # --- Return ---
    context_text = "\n".join(lines).strip()
    return context_text


# ----------------------------- Supporting documents loader --------------------------------

def _load_supporting_docs_jsonl(path: Optional[str] = None) -> Dict[str, List[Dict[str, str]]]:
    """
    Load the supporting documents JSONL file and index by grant_id.

    Returns: dict[grant_id] -> list of supporting docs
    """
    if path is None:
        path = "data/supporting_jsonl/docs.jsonl"

    try:
        p = Path(path)
        if not p.exists():
            return {}

        indexed = {}
        with open(p, "r", encoding="utf-8") as f:
            for line in f:
                if not line.strip():
                    continue
                try:
                    doc = json.loads(line)
                    gid = doc.get("grant_id", "").replace("competition-", "")
                    if gid:
                        if gid not in indexed:
                            indexed[gid] = []
                        indexed[gid].append(doc)
                except json.JSONDecodeError:
                    continue
        return indexed
    except Exception:
        return {}


def get_supporting_docs_for_grant(grant_id: str, k: int = 5, doc_types: Optional[List[str]] = None) -> List[Dict[str, str]]:
    """
    Retrieve supporting documents for a grant.

    Args:
        grant_id: Grant ID (with or without "competition-" prefix)
        k: Number of documents to return
        doc_types: Filter by document type ("supporting_html", "supporting_pdf", etc.)

    Returns: List of documents with extracted text
    """
    # Load cache on first use (could be cached module-level)
    cache = _load_supporting_docs_jsonl()

    # Normalize grant ID
    gid = str(grant_id).replace("competition-", "").strip()
    docs = cache.get(gid, [])

    # Filter by type if requested
    if doc_types:
        docs = [d for d in docs if d.get("doc_type") in doc_types]

    # Return top k
    return docs[:k]


def build_context_with_supporting(
    grant: Dict[str, Any],
    k: int = 5,
    pdf_only: bool = False,
    past_winners: Optional[List[Dict[str, Any]]] = None
) -> str:
    """
    Build context including supporting documents (PDFs + HTML sections).

    Args:
        grant: Grant dict with id/competition_id field (or will extract from URL)
        k: Number of supporting docs to include
        pdf_only: If True, only include PDF documents
        past_winners: Optional past winners for comparison

    Returns: Context string with supporting materials embedded
    """
    # Start with base context
    lines = [build_context(grant, past_winners)]

    # Extract grant ID from multiple sources
    gid = grant.get("id") or grant.get("competition_id")

    # If not found, try extracting from URL
    if not gid:
        url = grant.get("url") or ""
        match = re.search(r"/competition/(\d+)", url)
        if match:
            gid = match.group(1)

    if not gid:
        return lines[0]

    # Filter document types
    doc_types = ["supporting_pdf"] if pdf_only else ["supporting_pdf", "supporting_html"]

    # Load supporting docs
    supporting = get_supporting_docs_for_grant(str(gid), k=k, doc_types=doc_types)

    if supporting:
        lines.append("\n" + "="*80)
        lines.append("SUPPORTING MATERIALS & PDF CONTENT:")
        lines.append("="*80)

        for i, doc in enumerate(supporting, 1):
            doc_type = doc.get("doc_type", "unknown")
            section = doc.get("section", "Supporting Info")
            text = doc.get("text", "")

            # Truncate long text but keep it substantial
            if len(text) > 2000:
                text = text[:2000] + "\n[... truncated ...]"

            lines.append(f"\n[{i}] {section.upper()} ({doc_type})")
            lines.append("-" * 60)
            lines.append(text)

    return "\n".join(lines)


# Self-test
if __name__ == "__main__":
    fake_grant = {
        "title": "AI Battery Research Program",
        "funding_amount": "up to £1M",
        "deadline": "2025-12-17",
        "sections": {
            "summary_raw": "Funding for early-stage AI-driven battery optimization.",
            "scope_raw": "Projects must demonstrate significant improvement in energy density.",
        },
    }
    fake_winners = [
        {
            "project_title": "BatteryX AI",
            "lead_org": "EnergyAI Ltd",
            "award_amount": "£500,000",
            "competition": "Battery Innovation 2023",
            "abstract": "Developed machine learning models for lithium-ion battery efficiency.",
        }
    ]
    print(build_context(fake_grant, fake_winners))