src/generator.py

"""
Generator: produces cited answers from retrieved regulatory sections.
Supports conversational history for natural follow-up responses.

The system prompt itself lives in `prompts/generator_system.md` — edit
that file (not this one) to interrogate or revise the prompt. Loaded
fresh on every call so iteration is just save-and-rerun, no restart
needed.
"""

import os
from datetime import date

import litellm
from src.config import MODEL
from src.language import get_language_name
from src.usage import _extract_usage, _empty_usage


_PROMPT_PATH = os.path.join(
    os.path.dirname(os.path.dirname(os.path.abspath(__file__))),
    "prompts",
    "generator_system.md",
)


def _build_system_prompt() -> str:
    """Build the SYSTEM_PROMPT with today's date injected for transition-aware logic.

    The ASA Therapeutic and Health Advertising Code is mid-transition (current
    code in force until 1 April 2026 for new advertising / 1 July 2026 for all
    advertising; December 2025 code applies from those dates onwards). The
    generator needs to know today's date so it can decide which code is in force
    for any given answer.

    Reads from prompts/generator_system.md every call (not cached). Slightly
    slower but means iteration on the prompt is immediate without a Python
    restart.
    """
    today = date.today().isoformat()
    with open(_PROMPT_PATH, "r", encoding="utf-8") as f:
        template = f.read()
    return template.replace("{TODAY}", today)


# NOTE: SYSTEM_PROMPT is now built per-request from prompts/generator_system.md
# rather than cached at module load. This is intentional — lets you iterate on
# the prompt without restarting the Streamlit process. Call sites below use
# _build_system_prompt() directly so today's date and any prompt edits both
# pick up on the next request.


def _build_messages(query, sections, language="en", history=None, profession=None):
    """Build the LLM messages array and metadata from sections."""
    # Build context from retrieved sections
    context_parts = []
    for i, section in enumerate(sections, 1):
        domain_label = section["domain"].replace("_", " ").title()
        url_line = f"\nSource: {section['source_url']}" if section.get("source_url") else ""
        context_parts.append(
            f"[Source {i}] Domain: {domain_label} | Section: {section['title']}{url_line}\n"
            f"{section['text'][:2000]}"  # Cap per-section length
        )

    context = "\n\n---\n\n".join(context_parts)

    # Language instruction
    if language != "en":
        lang_name = get_language_name(language) if language else language
        lang_instruction = (
            f"\nIMPORTANT: The user has selected {lang_name} as their response language. "
            f"Write your entire answer in {lang_name}, "
            f"BUT keep all regulatory citations, section numbers, Act names, and source "
            f"quotes in their original English form."
        )
    else:
        lang_instruction = ""

    # Profession context: tells the model which council/board rules bind the user.
    # The `binds:` metadata in source material disambiguates for the model; this
    # gives it the user side of the matching equation.
    if profession:
        profession_instruction = (
            f"\n\nUSER PROFESSION: {profession}.\n"
            f"Apply the {profession}'s council/board rules as authoritative for this user. "
            f"If sources include `binds:` scope tags, treat sections binding {profession} "
            f"or applying across all professions as authoritative; treat other professions' "
            f"council rules as comparative only (cite for context, not as binding). "
            f"Do NOT cite Medical Council guidance as binding for non-MD practitioners — "
            f"their own council's rules are the binding ones."
        )
    else:
        profession_instruction = ""

    # Build messages array — system prompt rebuilt fresh from
    # prompts/generator_system.md each call so prompt edits and date
    # changes both take effect on the next request without a restart.
    messages = [{"role": "system", "content": _build_system_prompt()}]

    # Add conversation history (last 4 exchanges max to stay within context)
    if history:
        for msg in history[-8:]:
            messages.append({"role": msg["role"], "content": msg["content"]})

    user_message = f"""Source material:
{context}
{lang_instruction}{profession_instruction}

Question: {query}"""

    messages.append({"role": "user", "content": user_message})

    # Build citations and source previews from sections (independent of LLM response)
    citations = []
    for i, section in enumerate(sections, 1):
        citations.append({
            "index": i,
            "domain": section["domain"].replace("_", " ").title(),
            "title": section["title"],
            "line_num": section.get("line_num"),
            "source_url": section.get("source_url"),
        })

    english_sources = []
    for section in sections:
        english_sources.append({
            "title": section["title"],
            "domain": section["domain"].replace("_", " ").title(),
            "text": section["text"][:1000],
        })

    return messages, citations, english_sources


def generate_answer(
    query: str,
    sections: list[dict],
    language: str = "en",
    history: list[dict] | None = None,
    profession: str | None = None,
) -> dict:
    """Generate an answer with citations from retrieved sections."""
    if not sections:
        return {
            "answer": "I could not find relevant regulatory information to answer this question. Please try rephrasing or specify the area more clearly (e.g. testimonial rules, supplement claims, patient privacy).",
            "citations": [],
            "english_sources": [],
            "usage": _empty_usage(),
        }

    messages, citations, english_sources = _build_messages(query, sections, language, history, profession)

    try:
        response = litellm.completion(
            model=MODEL,
            messages=messages,
            temperature=0,
            max_tokens=2000,
        )
        usage = _extract_usage(response)
        answer_text = (response.choices[0].message.content or "").strip()

        return {
            "answer": answer_text,
            "citations": citations,
            "english_sources": english_sources,
            "usage": usage,
        }

    except Exception as e:
        return {
            "answer": f"Error generating answer: {e}",
            "citations": [],
            "english_sources": [],
            "usage": _empty_usage(),
        }


def generate_answer_streaming(
    query: str,
    sections: list[dict],
    language: str = "en",
    history: list[dict] | None = None,
    profession: str | None = None,
):
    """Stream answer text chunks. Yields strings, then a final metadata dict.

    Usage:
        for chunk in generate_answer_streaming(query, sections):
            if isinstance(chunk, str):
                # display text
            else:
                # chunk is a dict with citations, english_sources, answer
    """
    if not sections:
        yield {
            "answer": "I could not find relevant regulatory information to answer this question. Please try rephrasing or specify the area more clearly (e.g. testimonial rules, supplement claims, patient privacy).",
            "citations": [],
            "english_sources": [],
        }
        return

    messages, citations, english_sources = _build_messages(query, sections, language, history, profession)

    try:
        response = litellm.completion(
            model=MODEL,
            messages=messages,
            temperature=0,
            max_tokens=2000,
            stream=True,
        )

        answer_text = ""
        for chunk in response:
            delta = chunk.choices[0].delta.content
            if delta:
                answer_text += delta
                yield delta

        # Final metadata dict after all text chunks
        yield {
            "answer": answer_text,
            "citations": citations,
            "english_sources": english_sources,
        }

    except Exception as e:
        yield f"\n\nError generating answer: {e}"
        yield {
            "answer": f"Error generating answer: {e}",
            "citations": [],
            "english_sources": [],
        }