FabledScribe/src/fabledassistant/services/research.py

"""Web research pipeline: sub-queries → SearXNG → fetch → synthesize → note."""

import asyncio
import json
import logging
import re

import httpx

from fabledassistant.config import Config
from fabledassistant.services.llm import fetch_url_content, generate_completion, stream_chat
from fabledassistant.services.notes import create_note, update_note
from fabledassistant.models.note import Note

logger = logging.getLogger(__name__)

SEARXNG_QUERIES = 5      # sub-queries to generate
RESULTS_PER_QUERY = 3    # results fetched from SearXNG per query
PAGES_PER_QUERY = 3      # pages actually read per sub-query (top N results)
MAX_SYNTHESIS_SOURCES = 12  # deduplicated sources passed to synthesis LLM
CHARS_PER_SOURCE = 2000     # content chars per source sent to synthesis


def _build_sources_block(sources: list[dict]) -> str:
    """Format fetched sources into a text block for LLM prompts."""
    parts = []
    for i, s in enumerate(sources, 1):
        content = (s.get("content") or s.get("snippet") or "")[:CHARS_PER_SOURCE]
        parts.append(
            f"[Source {i}] {s['title']}\nURL: {s['url']}\nSearch query: {s['query']}\n\n{content}"
        )
    return "\n\n" + ("─" * 60) + "\n\n".join(parts)


async def _generate_outline(topic: str, sources: list[dict], model: str) -> list[dict]:
    """Generate a topic outline from fetched research sources.

    Returns a list of {"title": str, "focus": str} dicts (2–8 entries).
    Retries once on failure before returning [] (callers fall back to single-note).
    """
    import json as _json

    sources_block = _build_sources_block(sources) if sources else "(no sources)"
    messages = [
        {
            "role": "system",
            "content": (
                "You are a research organizer. Given research sources on a topic, produce a JSON array "
                "of section objects that together cover the topic comprehensively from distinct angles.\n\n"
                "Rules:\n"
                "- Return exactly 3–7 sections\n"
                "- Each section must cover a unique angle — no overlap between sections\n"
                "- Titles must work as standalone note titles (specific, not generic like 'Overview')\n"
                "- focus: one sentence describing exactly what this section covers\n"
                "- Respond with ONLY a JSON array, no other text\n\n"
                'Example: [{"title": "CRISPR: Molecular Mechanisms", "focus": "How Cas9 identifies and cuts DNA at guide-RNA-specified sites"}]'
            ),
        },
        {
            "role": "user",
            "content": f"Topic: {topic}\n\nSources:\n{sources_block}",
        },
    ]
    for attempt in range(2):
        try:
            # Pin num_ctx explicitly. The prompt carries up to 12 sources at
            # 2000 chars each (~6k tokens of source material alone) plus the
            # system prompt — well over Ollama's default model window on
            # qwen3. Without this, Ollama silently truncates the prompt, the
            # model can't see most of the sources, JSON parsing fails twice,
            # and the pipeline falls back to a single monolith note
            # (`research.py:251`). Do not remove even if `generate_completion`
            # appears to default this — see the comment there.
            raw = await generate_completion(
                messages, model, max_tokens=400, num_ctx=16384
            )
            raw = raw.strip()
            raw = re.sub(r"^```(?:json)?\s*", "", raw)
            raw = re.sub(r"\s*```$", "", raw)
            idx = raw.find("[")
            if idx >= 0:
                parsed, _ = _json.JSONDecoder().raw_decode(raw[idx:])
                if isinstance(parsed, list):
                    sections = [
                        s for s in parsed
                        if isinstance(s, dict) and s.get("title") and s.get("focus")
                    ]
                    if len(sections) >= 2:
                        return sections[:8]
        except Exception:
            logger.warning(
                "Outline generation attempt %d failed for topic '%s'",
                attempt + 1, topic, exc_info=True,
            )
    return []


async def _synthesize_section(
    section_title: str,
    section_focus: str,
    sources: list[dict],
    model: str,
) -> tuple[str, str]:
    """Synthesize one focused note section.

    Returns (section_title, body_markdown). Does not stream.
    """
    sources_block = _build_sources_block(sources) if sources else "(no sources provided)"
    messages = [
        {
            "role": "system",
            "content": (
                "You are a focused research writer. Write a single well-structured note section "
                "on the specific topic provided.\n\n"
                "Requirements:\n"
                f"- Focus strictly on: {section_focus}\n"
                "- 300–600 words of substantive prose\n"
                "- Use ### for subsections only when they genuinely aid clarity\n"
                "- Do NOT include a top-level # heading — the title is set separately\n"
                "- Write in detailed prose paragraphs — not bullet points\n"
                "- End with a '## Sources' section listing relevant source URLs as markdown hyperlinks\n"
                "- Ignore source material that falls outside your assigned focus"
            ),
        },
        {
            "role": "user",
            "content": (
                f"Section title: {section_title}\n"
                f"Focus: {section_focus}\n\n"
                f"Sources:\n{sources_block}"
            ),
        },
    ]
    raw = await generate_completion(messages, model, max_tokens=2048, num_ctx=16384)
    return section_title, raw.strip()


async def _generate_executive_summary(
    topic: str,
    section_bodies: list[tuple[str, str]],
    model: str,
) -> str:
    """Generate a 2-3 paragraph executive summary from completed section notes.

    Args:
        section_bodies: list of (title, body) pairs from the section notes.

    Returns summary markdown (no heading — caller adds structure).
    """
    sections_block = "\n\n".join(
        f"### {title}\n{body[:1500]}" for title, body in section_bodies
    )
    messages = [
        {
            "role": "system",
            "content": (
                "You are a research summarizer. Given several completed research sections on a topic, "
                "write a concise executive summary.\n\n"
                "Requirements:\n"
                "- 2–3 paragraphs of substantive prose (150–300 words total)\n"
                "- Cover the key findings and insights across all sections\n"
                "- Highlight the most important or surprising takeaways\n"
                "- Write so someone can decide which sections to read in detail\n"
                "- Do NOT include headings, bullet points, or source citations\n"
                "- Do NOT start with 'This research' or 'This document' — jump straight into the content"
            ),
        },
        {
            "role": "user",
            "content": f"Topic: {topic}\n\nSections:\n{sections_block}",
        },
    ]
    try:
        # Pin num_ctx explicitly — see `_generate_outline` comment for the
        # rationale. This prompt carries N sections × 1500 chars of section
        # prose, which can easily exceed the default model window. Don't
        # trust the `generate_completion` default to stick.
        raw = await generate_completion(
            messages, model, max_tokens=600, num_ctx=16384
        )
        return raw.strip()
    except Exception:
        logger.warning("Executive summary generation failed for '%s'", topic, exc_info=True)
        return ""


async def run_research_pipeline(
    topic: str,
    user_id: int,
    model: str,
    buf=None,
    project_id: int | None = None,
) -> Note:
    """Full research pipeline: search → fetch → outline → section notes → index note.

    Emits status events via buf throughout (when buf is provided).
    Returns the index note (or a single fallback note on outline failure).
    """
    def _status(msg: str) -> None:
        if buf is not None:
            buf.append_event("status", {"status": msg})

    # Step 1: Generate sub-queries
    _status("Generating search queries...")
    queries = await _generate_sub_queries(topic, model)
    logger.info("Research: generated %d sub-queries for topic '%s'", len(queries), topic)

    # Step 2: Search all queries in parallel (200 ms stagger to avoid hammering SearXNG)
    async def _search_with_stagger(i: int, query: str) -> tuple[str, list[dict]]:
        if i > 0:
            await asyncio.sleep(0.2 * i)
        _status(f"Searching: {query}...")
        results = await _search_searxng(query)
        logger.info("Research: query '%s' → %d results", query, len(results))
        return query, results

    search_results = await asyncio.gather(
        *[_search_with_stagger(i, q) for i, q in enumerate(queries)]
    )

    # Deduplicate URLs across all queries
    seen_urls: set[str] = set()
    url_tasks: list[tuple[str, dict, str]] = []  # (url, result_dict, query)
    for query, results in search_results:
        for result in results[:PAGES_PER_QUERY]:
            url = result.get("url", "")
            if url and url not in seen_urls:
                seen_urls.add(url)
                url_tasks.append((url, result, query))

    # Fetch all unique URLs in parallel
    async def _fetch_source(url: str, result: dict, query: str) -> dict:
        title = result.get("title", url)
        _status(f"Reading: {title[:60]}...")
        content = await fetch_url_content(url)
        return {
            "url": url,
            "title": title,
            "query": query,
            "snippet": result.get("snippet", ""),
            "content": content,
        }

    all_sources: list[dict] = list(await asyncio.gather(
        *[_fetch_source(url, result, query) for url, result, query in url_tasks]
    ))

    if not all_sources:
        raise ValueError(f"No results found for '{topic}'")

    good_sources = [s for s in all_sources if not s["content"].startswith("[Failed to fetch")]

    if not good_sources:
        raise ValueError(f"Could not read any sources for '{topic}'")

    synthesis_sources = good_sources[:MAX_SYNTHESIS_SOURCES]
    logger.info(
        "Research: %d/%d sources successfully fetched, using %d for synthesis",
        len(good_sources), len(all_sources), len(synthesis_sources),
    )

    # Step 3: Generate topic outline
    _status("Generating outline...")
    outline = await _generate_outline(topic, synthesis_sources, model)

    # Fallback: outline failed or too short → single monolithic note
    if not outline:
        logger.warning("Research outline empty, falling back to single note for '%s'", topic)
        _status("Synthesizing report...")
        title, body = await _synthesize_note(topic, synthesis_sources, model, buf=None)
        note = await create_note(
            user_id=user_id, title=title, body=body, tags=["research"], project_id=project_id,
        )
        logger.info("Research (fallback): created note id=%d title='%s'", note.id, note.title)
        return note

    # Step 4: Synthesize each section in parallel
    for section in outline:
        _status(f"Writing: {section['title']}...")

    raw_results = await asyncio.gather(
        *[_synthesize_section(s["title"], s["focus"], synthesis_sources, model) for s in outline],
        return_exceptions=True,
    )

    # Collect successful results for index + summary generation
    section_results: list[tuple[dict, str, str]] = []  # (outline_entry, title, body)
    for section, result in zip(outline, raw_results):
        if isinstance(result, Exception):
            logger.warning("Section synthesis failed for '%s': %s", section["title"], result)
            continue
        sec_title, sec_body = result
        section_results.append((section, sec_title, sec_body))

    # All sections failed — fall back to single note
    if not section_results:
        logger.warning("All section syntheses failed, falling back to single note for '%s'", topic)
        _status("Synthesizing report (fallback)...")
        title, body = await _synthesize_note(topic, synthesis_sources, model, buf=None)
        note = await create_note(
            user_id=user_id, title=title, body=body, tags=["research"], project_id=project_id,
        )
        return note

    # Step 5: Generate executive summary from section content
    _status("Writing summary...")
    executive_summary = await _generate_executive_summary(
        topic, [(t, b) for _, t, b in section_results], model,
    )

    # Step 6: Create index note first (so section notes can reference it via parent_id)
    from datetime import date as _date
    index_lines = [
        f"Research overview for **{topic}** — {_date.today().isoformat()}",
        "",
        f"Generated from {len(synthesis_sources)} web sources across {len(section_results)} sections.",
        "",
    ]
    if executive_summary:
        index_lines += ["## Summary", "", executive_summary, ""]
    index_lines += ["## Sections", ""]
    # Placeholder — will be updated with real links after section notes are created
    index_note = await create_note(
        user_id=user_id,
        title=f"Research: {topic}",
        body="\n".join(index_lines),
        tags=["research", "research-index"],
        project_id=project_id,
    )

    # Step 7: Create section notes with parent_id pointing to index
    _status(f"Saving {len(section_results)} notes...")
    section_note_pairs: list[tuple[dict, Note]] = []
    for section, sec_title, sec_body in section_results:
        try:
            note = await create_note(
                user_id=user_id,
                title=sec_title,
                body=sec_body,
                tags=["research"],
                project_id=project_id,
                parent_id=index_note.id,
            )
            section_note_pairs.append((section, note))
        except Exception:
            logger.warning("Failed to save section note '%s'", sec_title, exc_info=True)

    # Step 8: Update index note body with real links to section notes
    for section, note in section_note_pairs:
        index_lines.append(f"- [{note.title}](/notes/{note.id}) — {section['focus']}")

    await update_note(
        user_id=user_id,
        note_id=index_note.id,
        body="\n".join(index_lines),
    )

    logger.info(
        "Research: created %d section notes + index id=%d for topic '%s'",
        len(section_note_pairs), index_note.id, topic,
    )
    return index_note


async def _generate_sub_queries(topic: str, model: str) -> list[str]:
    """Ask the model for focused search queries for the topic."""
    messages = [
        {
            "role": "system",
            "content": (
                f"You are a research assistant. Given a research topic, generate exactly {SEARXNG_QUERIES} "
                "focused web search queries that together would provide comprehensive coverage of the topic. "
                "Vary the angle of each query: include overview, implementation details, best practices, "
                "common problems, and real-world examples. "
                "Respond with ONLY a JSON array of strings, no other text. "
                'Example: ["query one", "query two", "query three"]'
            ),
        },
        {"role": "user", "content": f"Topic: {topic}"},
    ]
    try:
        raw = await generate_completion(messages, model, max_tokens=200)
        raw = raw.strip()
        raw = re.sub(r"^```(?:json)?\s*", "", raw)
        raw = re.sub(r"\s*```$", "", raw)
        idx = raw.find("[")
        if idx >= 0:
            parsed, _ = json.JSONDecoder().raw_decode(raw[idx:])
            if isinstance(parsed, list) and parsed:
                queries = [str(q).strip() for q in parsed if str(q).strip()]
                if queries:
                    return queries[:SEARXNG_QUERIES]
    except Exception:
        logger.warning("Sub-query generation failed, falling back to topic", exc_info=True)
    return [topic]


async def _search_searxng(query: str) -> list[dict]:
    """Search SearXNG and return top results as [{url, title, snippet}]."""
    url = Config.SEARXNG_URL.rstrip("/") + "/search"
    params = {"q": query, "format": "json", "categories": "general"}
    for attempt in range(3):
        try:
            async with httpx.AsyncClient(timeout=10.0) as client:
                resp = await client.get(url, params=params)
                if resp.status_code == 429:
                    retry_after = int(resp.headers.get("Retry-After", "5"))
                    wait = min(retry_after, 10) * (attempt + 1)
                    logger.warning(
                        "SearXNG 429 for query '%s' (attempt %d/3), waiting %ds",
                        query, attempt + 1, wait,
                    )
                    await asyncio.sleep(wait)
                    continue
                resp.raise_for_status()
                data = resp.json()
                results = data.get("results", [])
                out = []
                for r in results[:RESULTS_PER_QUERY]:
                    out.append({
                        "url": r.get("url", ""),
                        "title": r.get("title", ""),
                        "snippet": r.get("content", ""),
                    })
                return out
        except httpx.HTTPStatusError:
            logger.warning("SearXNG search failed for query '%s'", query, exc_info=True)
            return []
        except Exception:
            logger.warning("SearXNG search failed for query '%s'", query, exc_info=True)
            return []
    logger.warning("SearXNG search gave up after 3 attempts for query '%s'", query)
    return []


async def _search_searxng_images(query: str) -> list[dict]:
    """Search SearXNG image category and return [{img_src, page_url, title, source_domain}]."""
    url = Config.SEARXNG_URL.rstrip("/") + "/search"
    params = {"q": query, "format": "json", "categories": "images"}
    for attempt in range(3):
        try:
            async with httpx.AsyncClient(timeout=10.0) as client:
                resp = await client.get(url, params=params)
                if resp.status_code == 429:
                    retry_after = int(resp.headers.get("Retry-After", "5"))
                    wait = min(retry_after, 10) * (attempt + 1)
                    logger.warning(
                        "SearXNG image 429 for '%s' (attempt %d/3), waiting %ds",
                        query, attempt + 1, wait,
                    )
                    await asyncio.sleep(wait)
                    continue
                resp.raise_for_status()
                data = resp.json()
                out = []
                for r in data.get("results", []):
                    img_src = r.get("img_src") or r.get("thumbnail_src", "")
                    if not img_src:
                        continue
                    try:
                        from urllib.parse import urlparse
                        source_domain = urlparse(r.get("url", "")).netloc or ""
                    except Exception:
                        source_domain = ""
                    out.append({
                        "img_src": img_src,
                        "page_url": r.get("url", ""),
                        "title": r.get("title", ""),
                        "source_domain": source_domain,
                    })
                return out
        except httpx.HTTPStatusError:
            logger.warning("SearXNG image search failed for '%s'", query, exc_info=True)
            return []
        except Exception:
            logger.warning("SearXNG image search failed for '%s'", query, exc_info=True)
            return []
    logger.warning("SearXNG image search gave up after 3 attempts for '%s'", query)
    return []


async def _synthesize_note(
    topic: str,
    sources: list[dict],
    model: str,
    buf=None,
) -> tuple[str, str]:
    """Synthesize a comprehensive markdown research document from fetched sources.

    Returns (title, body_markdown).
    When buf is provided, tokens are streamed into the chat buffer in real time
    so the user can see the note being written. Uses an extended context window.
    """
    sources_block = _build_sources_block(sources)

    messages = [
        {
            "role": "system",
            "content": (
                "You are a thorough researcher and writer. "
                "Your task is to write an exhaustive, well-structured document on the given topic — "
                "not a brief summary or intro paragraph.\n\n"
                "Requirements:\n"
                "- Write at least 2500 words of substantive content (excluding the Sources section)\n"
                "- Choose sections (##) that make sense for the topic — let the subject matter determine the structure. "
                "A technical topic might need implementation, configuration, and troubleshooting sections. "
                "A comparison topic might need dedicated sections per subject being compared plus a summary. "
                "A scientific topic might need background, mechanisms, research findings, and implications. "
                "Use your judgment — minimum 6 major sections.\n"
                "- Use ### for subsections where they add clarity\n"
                "- Write in detailed prose paragraphs — do not reduce sections to bullet-point lists\n"
                "- Include specific details, examples, data points, comparisons, and nuance from the sources\n"
                "- Do not pad with vague generalities — every paragraph should say something concrete\n"
                "- The first line must be the document title starting with '# '\n"
                "- End with a '## Sources' section listing every source as a markdown hyperlink\n\n"
                "The reader wants to finish this document with a thorough understanding of the topic, "
                "not just an overview."
            ),
        },
        {
            "role": "user",
            "content": (
                f"Write a comprehensive reference document on: {topic}\n\n"
                f"Sources ({len(sources)} pages fetched):\n{sources_block}"
            ),
        },
    ]

    if buf is not None:
        # Stream tokens into the chat buffer so the user sees the note being written
        raw_parts: list[str] = []
        async for token in stream_chat(
            messages, model, options={"num_ctx": 16384, "num_predict": 8192}
        ):
            raw_parts.append(token)
            buf.append_event("chunk", {"chunk": token})
            buf.content_so_far += token
        raw = "".join(raw_parts).strip()
    else:
        raw = await generate_completion(
            messages,
            model,
            max_tokens=8192,
            num_ctx=16384,
        )
        raw = raw.strip()

    # Extract title from first # heading
    lines = raw.splitlines()
    title = f"Research: {topic}"
    body_lines = lines
    if lines and lines[0].startswith("# "):
        title = lines[0][2:].strip()
        body_lines = lines[1:]

    body = "\n".join(body_lines).strip()
    return title, body