app/services/context_builder.py

"""Build grounded generation context for DocDoe AI."""
from __future__ import annotations

import json
from collections import Counter
from dataclasses import dataclass
from typing import Any

from sqlalchemy import desc, select
from sqlalchemy.orm import Session

from app.models.document import Document
from app.models.generation import Generation
from app.models.previous_paper import PreviousPaper
from app.models.previous_question import PreviousQuestion
from app.models.study_profile import StudyProfile
from app.models.user import User
from app.services.evidence_contract import build_evidence_context, _syllabus_matches
from app.services.retrieval import (
    chunks_to_context,
    is_context_answerable,
    retrieval_confidence_score,
    retrieve_relevant_chunks,
)


RETRIEVAL_PROFILES: dict[str, str] = {
    "study_map_profile": (
        "definitions formulas diagrams keywords confusing concepts exam questions study order"
    ),
    "notes_profile": "definitions core concepts formulas examples important exam points keywords",
    "simple_explanation_profile": (
        "simple meaning real-life examples steps confusing areas keywords analogy"
    ),
    "quiz_profile": (
        "facts definitions concepts formulas examples tricky areas exam questions numerical"
    ),
    "flashcards_profile": (
        "definitions formulas processes keywords mistakes exam answers active recall"
    ),
    "exam_mode_profile": "exam answers definitions keywords marks diagrams formulas derivation",
    "video_profile": (
        "simple explanation real-life examples process keywords exam answer quick quiz recap"
    ),
    "previous_paper_profile": (
        "previous paper patterns repeated topics chapter weightage expected questions frequency"
    ),
}


@dataclass(frozen=True)
class BuiltGenerationContext:
    context: str
    metadata: dict[str, Any]
    query: str
    retrieved_chunk_count: int
    evidence_label: str = "Based on uploaded material only"


def build_generation_context(
    *,
    db: Session,
    document: Document,
    generation_type: str,
    language: str,
    mode: str | None,
    current_user: User,
    query: str | None = None,
    limit: int = 6,
    include_study_map: bool = True,
) -> BuiltGenerationContext:
    profile_query = query or RETRIEVAL_PROFILES.get(
        f"{generation_type}_profile",
        RETRIEVAL_PROFILES["notes_profile"],
    )
    try:
        retrieved_chunks = retrieve_relevant_chunks(
            db=db,
            document_id=document.id,
            query=profile_query,
            limit=limit,
            user_id=current_user.id,
        )
    except Exception:  # noqa: BLE001 — retrieval failure must not break generation
        try:
            db.rollback()
        except Exception:  # noqa: BLE001
            pass
        retrieved_chunks = []
    source_context = chunks_to_context(
        retrieved_chunks,
        fallback_text=document.extracted_text,
        max_chars=5200,
    )
    retrieval_confidence = retrieval_confidence_score(retrieved_chunks, query=profile_query)
    answerable = is_context_answerable(retrieved_chunks, query=profile_query)
    try:
        study_profile = db.scalar(select(StudyProfile).where(StudyProfile.user_id == current_user.id))
    except Exception:  # noqa: BLE001 — recover from a poisoned transaction
        try:
            db.rollback()
        except Exception:  # noqa: BLE001
            pass
        study_profile = db.scalar(select(StudyProfile).where(StudyProfile.user_id == current_user.id))
    profile_skipped = bool(study_profile and (study_profile.extra or {}).get("onboardingSkipped"))
    board_value = study_profile.board if study_profile is not None and not profile_skipped else None
    grade_value = study_profile.grade if study_profile is not None and not profile_skipped else None
    latest_study_map = _latest_study_map(db, document, current_user) if include_study_map else None
    paper_insights = _previous_paper_insights(
        db,
        current_user,
        document,
        board=board_value,
        class_level=grade_value,
    )
    evidence_ctx = build_evidence_context(
        db, current_user.id,
        subject=document.subject,
        board=board_value,
        class_level=grade_value,
        explicit_source_ids=[document.id],
    )

    metadata: dict[str, Any] = {
        "user_id": current_user.id,
        "document_id": document.id,
        "title": document.title,
        "source_title": document.title,
        "material_type": getattr(document, "material_type", None),
        "subject": document.subject,
        "chapter": document.chapter,
        "syllabus": document.syllabus,
        "board": board_value,
        "class_level": grade_value,
        "file_type": document.file_type,
        "generation_type": generation_type,
        "language": language,
        "mode": mode,
        "retrieval_query": profile_query,
        "retrieved_chunk_count": len(retrieved_chunks),
        "retrieval_confidence": retrieval_confidence,
        "source_answerable": answerable,
        "has_study_map": latest_study_map is not None,
        "previous_paper_insights_count": len(paper_insights),
        "evidence_level": evidence_ctx.evidence_level,
        "evidence_label": evidence_ctx.evidence_label,
    }

    sections = [
        "# Document metadata",
        json.dumps(
            {
                "title": document.title,
                "subject": document.subject,
                "chapter": document.chapter,
                "syllabus": document.syllabus,
                "status": document.status,
            },
            ensure_ascii=True,
        ),
        "# Retrieved source chunks (ranked by relevance)",
        source_context,
    ]

    if latest_study_map:
        sections.extend(["# Existing study map (previously generated)", _safe_json(latest_study_map)])

    if paper_insights:
        sections.extend(
            [
                "# Previous paper insights (topic frequency from past exams)",
                _safe_json(paper_insights),
                "# How to use PYQ insights: Topics with higher frequency and marks are more likely to appear. "
                "Prioritize these in study plans and emphasize their keywords in answers.",
            ],
        )

    # Evidence contract — inject BEFORE strict tutor rules so AI sees it early
    sections.extend(
        [
            "# Evidence-bound answer policy",
            evidence_ctx.prompt_rules,
        ],
    )

    sections.extend(
        [
            "# Strict tutor rules",
            (
                f"Context confidence: {retrieval_confidence:.2f}. Source answerable: {answerable}.\n"
                "1. Use only the retrieved source chunks and PYQ metadata for source-specific facts.\n"
                "2. If the source is low-confidence or missing the requested fact, say what is missing.\n"
                "3. Simple meaning first - explain so a student gets it on first read.\n"
                "4. Use exact exam keywords that board evaluators check for.\n"
                "5. Keep answers concise - every sentence must help score marks.\n"
                "6. Include formulas with SI units in [brackets] for STEM topics.\n"
                "7. Add memory tricks, common mistakes, and quick self-check where schema allows.\n"
                "8. Do not invent facts, formulas, dates, or definitions not in context or standard curriculum.\n"
                "9. For Malayalam: mix Malayalam naturally with English technical terms and keywords."
            ),
        ],
    )

    return BuiltGenerationContext(
        context="\n\n".join(section for section in sections if section),
        metadata=metadata,
        query=profile_query,
        retrieved_chunk_count=len(retrieved_chunks),
        evidence_label=evidence_ctx.evidence_label,
    )


def _latest_study_map(
    db: Session,
    document: Document,
    current_user: User,
) -> dict[str, Any] | None:
    generation = db.scalar(
        select(Generation)
        .where(
            Generation.user_id == current_user.id,
            Generation.document_id == document.id,
            Generation.type == "study_map",
        )
        .order_by(desc(Generation.created_at)),
    )
    if generation is None:
        return None
    return generation.output_json


def _previous_paper_insights(
    db: Session,
    current_user: User,
    document: Document,
    *,
    board: str | None = None,
    class_level: str | None = None,
) -> list[dict[str, Any]]:
    paper_query = (
        select(PreviousPaper)
        .where(PreviousPaper.user_id == current_user.id)
        .where(PreviousPaper.verification_status == "verified")
    )
    if document.subject:
        paper_query = paper_query.where(PreviousPaper.subject == document.subject)
    papers = [
        paper
        for paper in db.scalars(paper_query).all()
        if _syllabus_matches(paper.syllabus, board, class_level)
    ]
    paper_ids = [paper.id for paper in papers]
    if not paper_ids:
        return []
    questions = list(
        db.scalars(
            select(PreviousQuestion)
            .where(PreviousQuestion.previous_paper_id.in_(paper_ids))
            .limit(250),
        ).all(),
    )
    subject = (document.subject or "").strip().lower()
    if subject:
        questions = [
            question
            for question in questions
            if not question.subject or question.subject.strip().lower() == subject
        ]
    topic_counts = Counter(
        (question.topic or question.chapter or "unknown").strip()
        for question in questions
    )
    insights: list[dict[str, Any]] = []
    for topic, count in topic_counts.most_common(8):
        related = [
            question
            for question in questions
            if (question.topic or question.chapter or "unknown").strip() == topic
        ]
        marks_list = [question.marks or 1 for question in related]
        insights.append(
            {
                "topic": topic,
                "times_asked": count,
                "total_marks": sum(marks_list),
                "avg_marks": round(sum(marks_list) / len(marks_list), 1) if marks_list else 0,
                "question_types": sorted({question.question_type for question in related}),
                "last_asked_year": max((question.year or 0 for question in related), default=0)
                or None,
                "marks_range": f"{min(marks_list)}-{max(marks_list)}" if marks_list else "1",
            },
        )
    return insights


def _safe_json(value: Any, max_chars: int = 4000) -> str:
    text = json.dumps(value, ensure_ascii=True)
    if len(text) <= max_chars:
        return text
    return f"{text[:max_chars].rstrip()} [trimmed]"