src/analyzer/qa_service.py

"""
QA service layer for Grant Analyst.

Provides streaming and non-streaming QA responses with prompt injection hardening.
API and UI should use this service instead of calling LLM/search directly.
"""

from __future__ import annotations

import json
import logging
from typing import Iterable, List, Optional

from .models import QARequest, QAChunk, ChunkType, Grant, CitationInfo
from .search.service import search_grants
from .llm_client import LLMClient
from .config import get_settings
from .cache.memo import cache_get, cache_put

logger = logging.getLogger(__name__)

# Shared LLM client singleton
_llm_client: LLMClient | None = None


def get_llm_client() -> LLMClient:
    """Get or create shared LLM client instance."""
    global _llm_client
    if _llm_client is None:
        _llm_client = LLMClient()
    return _llm_client


# Context limits for LLM calls
MAX_GRANTS = 5
MAX_DESC_CHARS = 1200

# Concise system prompt - reduced token usage
SYSTEM_PROMPT = """You are a specialised grant assistant focused on Innovate UK and related UK/EU funding schemes.

Use ONLY the retrieved grant records and their metadata as your primary evidence.
If required information is not present in the retrieved context, say you are unsure rather than guessing.

When you answer:
- Be concise and structured.
- Prioritise: grant name, funder, key eligibility, funding amount/range, deadline, and URL.
- Highlight constraints or caveats clearly.
- If multiple grants are relevant, list them clearly instead of writing long paragraphs.
"""


def _sanitize_grant_text(text: str) -> str:
    """
    Sanitize grant text to remove potential prompt injection attempts.

    Strips lines that look like instructions or system prompts.

    Args:
        text: Raw text from grant data

    Returns:
        Sanitized text safe for inclusion in prompts
    """
    if not text:
        return ""

    # Keywords that indicate potential injection
    injection_keywords = [
        "ignore previous",
        "ignore all previous",
        "system prompt",
        "you are now",
        "forget everything",
        "new instructions",
        "disregard",
        "override",
        "act as",
    ]

    lines = text.split("\n")
    safe_lines = []

    for line in lines:
        line_lower = line.lower().strip()

        # Skip lines that look like injection attempts
        if any(keyword in line_lower for keyword in injection_keywords):
            logger.warning(f"Filtered potential injection: {line[:50]}...")
            continue

        safe_lines.append(line)

    return "\n".join(safe_lines)


def format_grant_context(grants: List[Grant]) -> str:
    """
    Format top search results into a compact context string for the LLM.

    Limits both the number of grants and description length to control token usage.

    Args:
        grants: List of Grant objects

    Returns:
        Formatted, compact context string
    """
    lines: list[str] = []

    for i, grant in enumerate(grants[:MAX_GRANTS]):
        # Extract key attributes
        title = grant.title or "Untitled"
        funder = grant.source or grant.programme or ""
        ref_id = grant.id or ""
        url = grant.url or ""
        deadline = str(grant.close_date) if grant.close_date else ""
        status = grant.status or ""

        # Extract and truncate description/summary
        desc = grant.summary or grant.scope or ""
        if desc and len(desc) > MAX_DESC_CHARS:
            # Sanitize and truncate
            desc = _sanitize_grant_text(desc[:MAX_DESC_CHARS]).rstrip() + "..."
        elif desc:
            desc = _sanitize_grant_text(desc)

        # Format funding amount
        amount = ""
        if grant.funding:
            funding_parts = []
            if grant.funding.min is not None:
                funding_parts.append(f"£{grant.funding.min:,.0f}")
            if grant.funding.max is not None:
                funding_parts.append(f"£{grant.funding.max:,.0f}")
            if funding_parts:
                amount = " - ".join(funding_parts)

        # Build compact block
        block = [
            f"Grant {i+1}: {title}",
            f"  Reference: {ref_id}" if ref_id else "",
            f"  Funder: {funder}" if funder else "",
            f"  Status: {status}" if status else "",
            f"  Funding: {amount}" if amount else "",
            f"  Deadline: {deadline}" if deadline else "",
            f"  URL: {url}" if url else "",
            f"  Summary: {desc}" if desc else "",
        ]
        lines.append("\n".join([ln for ln in block if ln]))

    return "\n\n".join(lines)


def stream_answer(req: QARequest) -> Iterable[QAChunk]:
    """
    Generate streaming QA response.

    Yields QAChunk objects that can be serialized to NDJSON or SSE.

    Args:
        req: Validated QA request

    Yields:
        QAChunk objects of various types (metadata, token, citations, done, error)
    """
    settings = get_settings()

    try:
        # Initialize LLM client
        llm = get_llm_client()

        # Send metadata
        yield QAChunk(type=ChunkType.METADATA, session_id=req.session_id, query=req.query)

        # Search for relevant grants
        logger.debug(f"Query: {req.query[:100]}")
        hits = search_grants(req.query, req.filters, limit=10)
        logger.info(f"Search found {len(hits)} hits")

        if not hits:
            # No grants found
            yield QAChunk(
                type=ChunkType.TOKEN,
                content="I couldn't find any grants matching your query. Try different keywords or broader terms.",
            )
            yield QAChunk(type=ChunkType.DONE, latency_ms=0)
            return

        # Extract grants and build compact context
        grants = [hit.grant for hit in hits]
        context = format_grant_context(grants)

        # Build prompt
        user_prompt = f"""User query: {req.query}

Relevant grant opportunities:

{context}

Based on the grants above, answer the user's query concisely and accurately.
Cite specific grants by ID and title. If none of the grants are truly relevant, say so."""

        messages = [
            {"role": "system", "content": SYSTEM_PROMPT},
            {"role": "user", "content": user_prompt},
        ]

        # Stream LLM response (log moved to debug to reduce overhead)
        logger.debug(f"Streaming LLM for query: {req.query[:50]}")

        # Use streaming chat
        for token in llm.chat(
            messages,
            stream=True,
            max_tokens=1200,
            model_type="analyzer",  # Use analyzer model for QA
        ):
            yield QAChunk(type=ChunkType.TOKEN, content=token)

        # Send citations
        citations = [
            {"grant_id": grant.id, "title": grant.title, "url": grant.url, "score": hit.score}
            for grant, hit in zip(grants[:5], hits[:5])
        ]

        yield QAChunk(type=ChunkType.CITATIONS, citations=citations)

        # Send completion
        yield QAChunk(type=ChunkType.DONE)

    except Exception as e:
        logger.error(f"Error in stream_answer: {e}", exc_info=True)
        yield QAChunk(type=ChunkType.ERROR, error=str(e))


def answer_question(req: QARequest) -> dict:
    """
    Generate non-streaming QA response with memoization.

    Collects all chunks from stream_answer and returns a complete response.
    Caches results to avoid redundant LLM calls for identical queries.

    Args:
        req: Validated QA request

    Returns:
        Dict with answer, citations, and metadata
    """
    import time

    start_time = time.time()

    # Build cache key from normalized query and filters
    normalized_query = req.query.strip()
    filters_dict = req.filters.dict() if req.filters else {}
    cache_key = f"qa:{normalized_query}|filters:{json.dumps(filters_dict, sort_keys=True)}"

    # Try cache first
    cached = cache_get(cache_key)
    if cached:
        logger.debug(f"Cache hit for query: {normalized_query[:50]}")
        return cached

    answer_parts = []
    citations = []
    error = None
    success = True

    try:
        for chunk in stream_answer(req):
            if chunk.type == ChunkType.TOKEN:
                if chunk.content:
                    answer_parts.append(chunk.content)
            elif chunk.type == ChunkType.CITATIONS:
                citations = chunk.citations or []
            elif chunk.type == ChunkType.ERROR:
                error = chunk.error
                success = False
                break

    except Exception as e:
        error = str(e)
        success = False
        logger.error(f"Error in answer_question: {e}", exc_info=True)

    latency_ms = int((time.time() - start_time) * 1000)

    result = {
        "session_id": req.session_id or "unknown",
        "query": req.query,
        "answer": "".join(answer_parts),
        "citations": citations,
        "latency_ms": latency_ms,
        "success": success,
        "error": error,
    }

    # Cache successful responses
    if success and not error:
        try:
            cache_put(cache_key, result)
            logger.debug(f"Cached response for: {normalized_query[:50]}")
        except Exception as cache_err:
            logger.debug(f"Cache storage failed: {cache_err}")

    return result