"""Retrieval and citation handling for AI Mode.

The safety property, stated plainly: **retrieval decides what may be cited, and
the server enforces it.** The model is handed a numbered shortlist and told to
cite from it by id. Afterwards every citation it wrote is checked against that
shortlist and anything else is deleted. A citation the model invented cannot
survive, because surviving is not a thing the model gets to decide.

This is the same discipline as the article page not printing answers: a property
the system holds, not one the model is trusted to respect.

Two kinds of link exist and are treated differently. *Curated* links are rows an
educator created — a question tied to an article section, a card to a question.
They are assertions, so they are trusted, and a retrieved row that carries one to
another retrieved row is boosted: two things an educator already tied together
answering the same query is evidence, not coincidence. *Retrieved* links are
ranked guesses computed per query and stored nowhere.
"""
import logging
import re

from sqlalchemy import text as sa_text
from sqlalchemy.orm import Session

from app.models.article import Article, ArticleSectionIndex, QuestionArticleLink
from app.models.flashcard import Flashcard, FlashcardDeck
from app.models.question import Question
from app.models.user import User
from app.services.quiz_builder import bank_query, exam_scope_predicate
from app.services.search_service import hybrid_ids

logger = logging.getLogger(__name__)

# How many of each kind retrieval offers the model. Small on purpose: a
# shortlist the model can hold is worth more than a corpus it skims.
PER_KIND = 6
MAX_SOURCES = 14
# An excerpt long enough to answer from, short enough that fourteen of them fit.
EXCERPT_CHARS = 700
# A curated tie between two retrieved rows is evidence an educator left behind.
CURATED_BOOST = 0.5

CITATION_RE = re.compile(r"\[\[(article|section|question|card):([A-Za-z0-9#_-]+)\]\]")


def _clean(value: str | None, limit: int = EXCERPT_CHARS) -> str:
    text = re.sub(r"!\[[^\]]*\]\([^)]*\)", " ", value or "")
    text = re.sub(r"\[([^\]]*)\]\([^)]*\)", r"\1", text)
    text = re.sub(r"<[^>]+>", " ", text)
    text = re.sub(r"\s+", " ", text).strip()
    return text[:limit] + ("…" if len(text) > limit else "")


def _articles(db: Session, user: User, query: str) -> list[dict]:
    ranked, _ = hybrid_ids(db, query, "article", limit=PER_KIND * 3)
    if not ranked:
        return []
    rows = db.query(Article).filter(Article.id.in_(ranked)).all()
    if not user.is_moderator:
        rows = [a for a in rows if a.status == "published" or a.user_id == user.id]
    order = {rid: i for i, rid in enumerate(ranked)}
    rows.sort(key=lambda a: order.get(a.id, len(order)))
    return [{
        "kind": "article", "ref": str(a.id), "id": a.id,
        "title": a.title,
        "text": _clean(a.summary or a.content),
        "score": 1.0 / (1 + order.get(a.id, 0)),
    } for a in rows[:PER_KIND]]


def _sections(db: Session, user: User, query: str) -> list[dict]:
    ranked, _ = hybrid_ids(db, query, "article_section", limit=PER_KIND * 3)
    if not ranked:
        return []
    rows = db.query(ArticleSectionIndex).filter(ArticleSectionIndex.id.in_(ranked)).all()
    if not rows:
        return []
    articles = {a.id: a for a in db.query(Article).filter(
        Article.id.in_({r.article_id for r in rows})).all()}
    order = {rid: i for i, rid in enumerate(ranked)}
    rows.sort(key=lambda r: order.get(r.id, len(order)))
    out = []
    for row in rows[:PER_KIND]:
        article = articles.get(row.article_id)
        if not article:
            continue
        if article.status != "published" and not user.is_moderator and article.user_id != user.id:
            continue
        out.append({
            "kind": "section", "ref": f"{article.id}#{row.section_id}",
            "id": article.id, "section_id": row.section_id,
            "title": f"{article.title} › {row.title or 'section'}",
            "text": _clean(row.content),
            "score": 1.0 / (1 + order.get(row.id, 0)),
        })
    return out


def _questions(db: Session, user: User, query: str) -> list[dict]:
    ranked, _ = hybrid_ids(db, query, "question", limit=PER_KIND * 3)
    if not ranked:
        return []
    q = bank_query(db, user).filter(Question.id.in_(ranked))
    scope = exam_scope_predicate(db, user)
    if scope is not None:
        q = q.filter(scope)
    order = {rid: i for i, rid in enumerate(ranked)}
    rows = sorted(q.all(), key=lambda r: order.get(r.id, len(order)))
    return [{
        "kind": "question", "ref": str(row.id), "id": row.id,
        "title": f"Question #{row.id}",
        # The stem only. An answer belongs to the quiz runner, and a chat that
        # printed it would hand away the practice it is meant to prepare for.
        "text": _clean(row.question_text, 320),
        "score": 1.0 / (1 + order.get(row.id, 0)),
    } for row in rows[:PER_KIND]]


def _cards(db: Session, user: User, query: str) -> list[dict]:
    ranked, _ = hybrid_ids(db, query, "flashcard", limit=PER_KIND * 3)
    if not ranked:
        return []
    own = [d.id for d in db.query(FlashcardDeck.id).filter(
        FlashcardDeck.user_id == user.id, FlashcardDeck.deleted_at.is_(None)).all()]
    if not own:
        return []
    rows = db.query(Flashcard).filter(
        Flashcard.id.in_(ranked), Flashcard.deck_id.in_(own)).all()
    order = {rid: i for i, rid in enumerate(ranked)}
    rows.sort(key=lambda r: order.get(r.id, len(order)))
    return [{
        "kind": "card", "ref": str(row.id), "id": row.id,
        "title": _clean(row.front, 90),
        "text": _clean(row.back, 260),
        "score": 1.0 / (1 + order.get(row.id, 0)),
    } for row in rows[:PER_KIND]]


def _apply_curated_boost(db: Session, sources: list[dict]) -> None:
    """Lift a source that an educator tied to another source in this shortlist.

    Two rows an educator already linked, both surfacing for the same question, is
    a stronger signal than either ranking alone — that is what a curated link is
    for. Nothing is stored; the boost lives only in this ordering.
    """
    article_ids = {s["id"] for s in sources if s["kind"] in ("article", "section")}
    question_ids = {s["id"] for s in sources if s["kind"] == "question"}
    if not article_ids or not question_ids:
        return
    pairs = db.query(QuestionArticleLink).filter(
        QuestionArticleLink.article_id.in_(article_ids),
        QuestionArticleLink.question_id.in_(question_ids)).all()
    if not pairs:
        return
    linked_articles = {p.article_id for p in pairs}
    linked_questions = {p.question_id for p in pairs}
    for source in sources:
        if source["kind"] in ("article", "section") and source["id"] in linked_articles:
            source["score"] += CURATED_BOOST
            source["curated"] = True
        elif source["kind"] == "question" and source["id"] in linked_questions:
            source["score"] += CURATED_BOOST
            source["curated"] = True


def retrieve(db: Session, user: User, query: str) -> list[dict]:
    """The only things the model will be allowed to cite for this message."""
    query = (query or "").strip()
    if len(query) < 2:
        return []
    sources: list[dict] = []
    for finder in (_sections, _articles, _questions, _cards):
        try:
            sources.extend(finder(db, user, query))
        except Exception:
            # One corpus failing narrows the answer; it does not end it.
            logger.warning("AI Mode retrieval failed for %s", finder.__name__, exc_info=True)
    _apply_curated_boost(db, sources)
    sources.sort(key=lambda s: -s["score"])
    return sources[:MAX_SOURCES]


def sources_block(sources: list[dict]) -> str:
    """The shortlist, as the model sees it."""
    lines = []
    for source in sources:
        marker = f"[[{source['kind']}:{source['ref']}]]"
        curated = " (an educator linked this to another source here)" if source.get("curated") else ""
        lines.append(f"{marker} {source['title']}{curated}\n{source['text']}")
    return "\n\n".join(lines)


def build_prompt(sources: list[dict]) -> str:
    if not sources:
        return (
            "You are a study assistant for a pediatrics learning platform.\n"
            "Nothing in this learner's library matches their question. Say so plainly "
            "in one or two sentences and suggest what they might search for instead. "
            "Do not answer from your own knowledge, and do not cite anything."
        )
    return (
        "You are a study assistant for a pediatrics learning platform.\n\n"
        "Answer only from the sources below. They are the learner's own library — "
        "if they do not contain the answer, say so rather than filling the gap from "
        "your own knowledge, which the learner cannot check against anything.\n\n"
        "Cite with the exact marker shown, for example [[article:7]] or "
        "[[section:7#abc123]], placed at the end of the sentence it supports. Never "
        "write a URL and never cite a marker that is not listed here.\n\n"
        "Never reveal the answer to a practice question. You may say what a question "
        "is about so the learner can go and attempt it.\n\n"
        "Be brief: a few sentences or a short list.\n\n"
        f"SOURCES\n\n{sources_block(sources)}"
    )


def enforce_citations(reply: str, sources: list[dict]) -> tuple[str, list[dict]]:
    """Drop every citation that was not retrieved, and report the ones that stand.

    This is the step that makes a hallucinated citation impossible rather than
    unlikely. The model can write whatever marker it likes; only markers in the
    shortlist survive contact with this function.
    """
    allowed = {f"{s['kind']}:{s['ref']}": s for s in sources}
    used: dict[str, dict] = {}

    def replace(match: re.Match) -> str:
        key = f"{match.group(1)}:{match.group(2)}"
        source = allowed.get(key)
        if source is None:
            return ""   # Invented, or pointing at something this learner may not see.
        used[key] = source
        return match.group(0)

    cleaned = CITATION_RE.sub(replace, reply)
    # Deleting a marker can leave a double space or a space before a full stop.
    cleaned = re.sub(r"[ \t]{2,}", " ", cleaned)
    cleaned = re.sub(r"\s+([.,;:!?])", r"\1", cleaned).strip()

    citations = [{
        "marker": f"[[{s['kind']}:{s['ref']}]]",
        "kind": s["kind"], "id": s["id"],
        "section_id": s.get("section_id"),
        "title": s["title"],
        "curated": bool(s.get("curated")),
    } for s in used.values()]
    return cleaned, citations