src/flashcard_generator/chains.py

from __future__ import annotations

import json
import re
from collections import Counter

from .models import Flashcard
from .text_processing import clean_text

try:
    from langchain_core.prompts import PromptTemplate
except Exception:  # pragma: no cover - only used when langchain-core is unavailable.
    PromptTemplate = None


CONCEPT_PROMPT = """Extract the top {count} academic concepts from this summary.
Return JSON only as a list of objects with keys "concept" and "definition".

Summary:
{summary}
"""

FLASHCARD_PROMPT = """Create exam revision flashcards from these concepts.
Return JSON only as a list of objects with keys "question", "short_answer", "long_answer", "concept", "difficulty", and "bloom_level".

Concepts:
{concepts_json}
"""

BLOOM_PROMPT = """Rewrite easy recall questions into harder Bloom's Taxonomy questions.
Prefer apply, analyze, evaluate, or create levels. Return JSON only.

Flashcards:
{cards_json}
"""


def build_prompt(template: str, **values: object) -> str:
    if PromptTemplate is None:
        return template.format(**values)
    prompt = PromptTemplate.from_template(template)
    return prompt.format(**values)


def extract_key_concepts(summary: str, count: int = 5) -> list[dict[str, str]]:
    """Extract concepts locally while preserving the JSON contract required by the brief."""
    build_prompt(CONCEPT_PROMPT, count=count, summary=summary)
    cleaned = clean_text(summary)
    if not cleaned:
        return []

    candidate_phrases = _candidate_phrases(cleaned)
    concepts: list[dict[str, str]] = []
    seen: set[str] = set()
    for phrase in candidate_phrases:
        key = phrase.lower()
        if key in seen:
            continue
        seen.add(key)
        concepts.append({"concept": phrase, "definition": _definition_for(cleaned, phrase)})
        if count > 0 and len(concepts) >= count:
            break

    if not concepts:
        concepts.append({"concept": cleaned.split(".")[0][:80], "definition": cleaned[:240]})

    return _json_round_trip(concepts)


def dedupe_concepts(concepts: list[dict[str, str]]) -> list[dict[str, str]]:
    unique: list[dict[str, str]] = []
    seen: set[str] = set()
    for item in concepts:
        concept = clean_text(str(item.get("concept", "")))
        definition = clean_text(str(item.get("definition", "")))
        key = _normalize_key(concept)
        if not concept or not definition or key in seen:
            continue
        seen.add(key)
        unique.append({"concept": concept, "definition": definition})
    return unique


def generate_flashcards(concepts: list[dict[str, str]], cards_per_concept: int = 1) -> list[Flashcard]:
    build_prompt(FLASHCARD_PROMPT, concepts_json=json.dumps(concepts, indent=2))
    cards: list[Flashcard] = []
    question_styles = [
        (
            "explain",
            "medium",
            "understand",
            "What role does {concept} play in this topic?",
        ),
        (
            "compare",
            "hard",
            "analyze",
            "How does {concept} connect to the surrounding ideas in the notes?",
        ),
        (
            "apply",
            "medium",
            "apply",
            "How could {concept} be applied in an exam-style problem?",
        ),
        (
            "cause",
            "hard",
            "analyze",
            "Why would a change in {concept} affect the outcome being studied?",
        ),
        (
            "evidence",
            "medium",
            "evaluate",
            "What evidence or reasoning from the notes supports {concept}?",
        ),
        (
            "misconception",
            "hard",
            "evaluate",
            "What common misconception about {concept} should a student avoid?",
        ),
        (
            "define",
            "easy",
            "remember",
            "What does {concept} mean?",
        ),
        (
            "importance",
            "medium",
            "understand",
            "Why is {concept} important for exam revision?",
        ),
        (
            "difference",
            "hard",
            "analyze",
            "How can {concept} be distinguished from related ideas?",
        ),
        (
            "limitation",
            "hard",
            "evaluate",
            "What limitation or condition should be considered when using {concept}?",
        ),
    ]

    for concept_index, item in enumerate(dedupe_concepts(concepts)):
        concept = clean_text(item.get("concept", ""))
        long_answer = _polish_long_answer(item.get("definition", ""), concept)
        short_answer = _short_answer(long_answer, concept)
        if not concept or not short_answer or not long_answer:
            continue

        for offset in range(cards_per_concept):
            style = question_styles[(concept_index + offset) % len(question_styles)]
            _, difficulty, bloom_level, template = style
            cards.append(
                Flashcard(
                    question=_one_line_question(template.format(concept=concept)),
                    short_answer=short_answer,
                    long_answer=_style_long_answer(long_answer, concept, bloom_level),
                    concept=concept,
                    difficulty=difficulty,
                    bloom_level=bloom_level,
                )
            )
    return dedupe_flashcards(cards)


def harden_flashcards(cards: list[Flashcard]) -> list[Flashcard]:
    build_prompt(BLOOM_PROMPT, cards_json=json.dumps([card.__dict__ for card in cards], indent=2))
    hard_templates = [
        ("How does {concept} influence the larger process described in the notes?", "analyze"),
        ("Why is {concept} important for solving exam problems on this topic?", "evaluate"),
        ("What might happen if {concept} were missing, incorrect, or changed?", "evaluate"),
        ("What short example would demonstrate {concept} in action?", "create"),
    ]
    hardened: list[Flashcard] = []
    for index, card in enumerate(cards):
        if card.difficulty != "easy":
            hardened.append(card)
            continue
        template, bloom_level = hard_templates[index % len(hard_templates)]
        hardened.append(
            Flashcard(
                question=template.format(concept=card.concept),
                short_answer=card.short_answer,
                long_answer=card.long_answer,
                concept=card.concept,
                difficulty="hard",
                bloom_level=bloom_level,
            )
        )
    return dedupe_flashcards(hardened)


def dedupe_flashcards(cards: list[Flashcard]) -> list[Flashcard]:
    unique: list[Flashcard] = []
    seen_questions: set[str] = set()
    seen_pairs: set[tuple[str, str]] = set()
    for card in cards:
        question_key = _normalize_key(card.question)
        pair_key = (question_key, _normalize_key(card.short_answer), _normalize_key(card.long_answer))
        if not card.question.strip() or not card.short_answer.strip() or not card.long_answer.strip():
            continue
        if question_key in seen_questions or pair_key in seen_pairs:
            continue
        seen_questions.add(question_key)
        seen_pairs.add(pair_key)
        unique.append(card)
    return unique


def parse_json_list(raw_text: str) -> list[dict[str, object]]:
    """Parse a JSON list from plain or fenced model output."""
    match = re.search(r"\[[\s\S]*\]", raw_text)
    if not match:
        return []
    try:
        value = json.loads(match.group(0))
    except json.JSONDecodeError:
        return []
    return value if isinstance(value, list) else []


def _candidate_phrases(text: str) -> list[str]:
    words = re.findall(r"[A-Za-z][A-Za-z-]{3,}", text)
    stopwords = {
        "about",
        "after",
        "also",
        "because",
        "between",
        "could",
        "absorbs",
        "convert",
        "converts",
        "drive",
        "drives",
        "during",
        "example",
        "from",
        "have",
        "important",
        "into",
        "lecture",
        "more",
        "notes",
        "other",
        "produce",
        "produces",
        "should",
        "their",
        "there",
        "these",
        "this",
        "through",
        "using",
        "were",
        "when",
        "which",
        "with",
        "would",
    }
    filtered = [word.lower() for word in words if word.lower() not in stopwords]
    counts = Counter(filtered)
    ranked = [word for word, _ in counts.most_common(20)]

    phrases = _technical_phrases(text, stopwords)
    noun_like = re.findall(r"\b(?:[A-Z][a-z]+(?:\s+[A-Z][a-z]+){0,2})\b", text)
    phrases.extend(phrase for phrase in noun_like if phrase.lower() not in stopwords)
    phrases.extend(word.title() for word in ranked if counts[word] > 1)
    return phrases


def _technical_phrases(text: str, stopwords: set[str]) -> list[str]:
    phrases: list[str] = []
    sentences = re.split(r"(?<=[.!?])\s+", text)
    for size in (2, 3):
        counts: Counter[str] = Counter()
        for sentence in sentences:
            words = [word.lower() for word in re.findall(r"[A-Za-z][A-Za-z-]*", sentence)]
            for index in range(0, len(words) - size + 1):
                window = words[index : index + size]
                if any(len(word) < 4 for word in window):
                    continue
                if any(word in stopwords for word in window):
                    continue
                if len(set(window)) < size:
                    continue
                counts[" ".join(window)] += 1
        phrases.extend(phrase.title() for phrase, amount in counts.most_common(12) if amount >= 1)
    return phrases


def _definition_for(text: str, phrase: str) -> str:
    sentences = re.split(r"(?<=[.!?])\s+", text)
    phrase_lower = phrase.lower()
    for sentence in sentences:
        if phrase_lower in sentence.lower():
            return sentence.strip()
    return text[:260].strip()


def _json_round_trip(concepts: list[dict[str, str]]) -> list[dict[str, str]]:
    return json.loads(json.dumps(concepts))


def _normalize_key(text: str) -> str:
    return re.sub(r"[^a-z0-9]+", " ", text.lower()).strip()


def _one_line_question(question: str) -> str:
    cleaned = clean_text(question)
    if not cleaned.endswith("?"):
        cleaned = f"{cleaned.rstrip('.') }?"
    return cleaned


def _short_answer(long_answer: str, concept: str) -> str:
    first_sentence = re.split(r"(?<=[.!?])\s+", clean_text(long_answer))[0]
    if not first_sentence:
        return concept
    words = first_sentence.split()
    if len(words) <= 22:
        return first_sentence
    return " ".join(words[:22]).rstrip(",;:") + "."


def _style_long_answer(answer: str, concept: str, bloom_level: str) -> str:
    cleaned = clean_text(answer)
    if bloom_level in {"analyze", "evaluate", "create"} and concept.lower() not in cleaned.lower():
        return f"{concept} should be understood in context: {cleaned}"
    return cleaned


def _polish_long_answer(answer: str, concept: str) -> str:
    cleaned = clean_text(answer)
    if not cleaned:
        return ""

    if len(cleaned.split()) < 8:
        return f"{concept} refers to {cleaned.lower()}. It is important because it helps explain the main process, relationship, or outcome described in the notes."

    if len(cleaned.split()) < 18:
        return (
            f"{cleaned} In exam terms, focus on what {concept} does, why it matters, "
            "and how it connects to the cause, effect, or process described in the notes."
        )

    if not cleaned.endswith((".", "!", "?")):
        cleaned = f"{cleaned}."
    return cleaned