core/structured_output.py

import json
import re
from datetime import datetime
from pathlib import Path


def parse_quiz(raw_text):
    questions = []

    # split on question markers like Q1:, Q2:, 1., 2. etc
    question_blocks = re.split(
        r"(?:^|\n)\s*(?:Q\d*[\.:]\s*|(?:\d+)[\.)\s]+)",
        raw_text.strip(),
    )

    # fallback: split on blank lines
    if len(question_blocks) <= 1:
        question_blocks = re.split(r"\n\s*\n", raw_text.strip())

    for block in question_blocks:
        block = block.strip()
        if not block:
            continue

        question = _parse_single_question(block)
        if question:
            questions.append(question)

    return questions


def _parse_single_question(block):
    lines = block.strip().split("\n")
    if not lines:
        return None

    question_text = ""
    option_lines = []
    answer_line = ""

    for line in lines:
        stripped = line.strip()
        if not stripped:
            continue

        # option line like A), B), etc
        if re.match(r"^[A-Da-d][).:\s]", stripped):
            option_lines.append(stripped)
        # answer line
        elif re.match(r"^(?:Answer|Correct|Ans)[:\s]", stripped, re.IGNORECASE):
            answer_line = stripped
        # question text (before any options or answer)
        elif not option_lines and not answer_line:
            if question_text:
                question_text += " " + stripped
            else:
                question_text = stripped

    if not question_text:
        return None

    options = []
    for opt in option_lines:
        clean_opt = re.sub(r"^[A-Da-d][).:\s]+", "", opt).strip()
        options.append(clean_opt)

    correct_answer = ""
    if answer_line:
        match = re.search(r"[A-Da-d]", answer_line)
        if match:
            correct_answer = match.group().upper()

    return {
        "question": question_text,
        "options": options,
        "correct_answer": correct_answer,
    }


def format_quiz(questions):
    if not questions:
        return ""

    lines = ["", "  ╔══════════════════════════════════════════╗"]
    lines.append("  ║              📝  QUIZ                    ║")
    lines.append("  ╚══════════════════════════════════════════╝")
    lines.append("")

    option_letters = ["A", "B", "C", "D"]

    for index, question in enumerate(questions, start=1):
        lines.append(f"  Q{index}: {question['question']}")

        if question["options"]:
            for opt_index, option in enumerate(question["options"]):
                letter = option_letters[opt_index] if opt_index < 4 else f"{opt_index + 1}"
                lines.append(f"      {letter}) {option}")

        if question["correct_answer"]:
            lines.append(f"      ✅ Answer: {question['correct_answer']}")

        lines.append("")

    lines.append("  " + "─" * 44)
    return "\n".join(lines)


def parse_flashcards(raw_text):
    # Try to parse complete JSON array first
    try:
        start = raw_text.find("[")
        end = raw_text.rfind("]")
        if start != -1 and end != -1 and end > start:
            json_str = raw_text[start:end+1]
            data = json.loads(json_str)
            if isinstance(data, list):
                cards = []
                for item in data:
                    if isinstance(item, dict) and "front" in item and "back" in item:
                        card = _clean_card(str(item["front"]), str(item["back"]))
                        if card:
                            cards.append(card)
                if cards:
                    return cards
    except Exception:
        pass

    # Recover truncated JSON — LLM may have run out of tokens mid-card
    try:
        start = raw_text.find("[")
        if start != -1:
            fragment = raw_text[start:]
            # find the last complete card object (last "}")
            last_brace = fragment.rfind("}")
            if last_brace != -1:
                recovered = fragment[:last_brace + 1] + "]"
                data = json.loads(recovered)
                if isinstance(data, list):
                    cards = []
                    for item in data:
                        if isinstance(item, dict) and "front" in item and "back" in item:
                            card = _clean_card(str(item["front"]), str(item["back"]))
                            if card:
                                cards.append(card)
                    if cards:
                        return cards
    except Exception:
        pass

    cards = []

    # try Front: / Back: pattern first
    front_back_pairs = re.findall(
        r"(?:Front|Question|Q)\s*[:\-]\s*(.+?)(?:\n|\s{2,})(?:Back|Answer|A)\s*[:\-]\s*(.+?)(?=\n\s*(?:Front|Question|Q)\s*[:\-]|\Z)",
        raw_text,
        re.IGNORECASE | re.DOTALL,
    )
    if front_back_pairs:
        for front, back in front_back_pairs:
            card = _clean_card(front.strip(), back.strip())
            if card:
                cards.append(card)
        return cards

    # try numbered pairs like "1. term - definition"
    numbered_pairs = re.findall(
        r"\d+[.)]\s*(.+?)\s*[-–→:]\s*(.+?)(?=\n\s*\d+[.)]|\Z)",
        raw_text,
        re.DOTALL,
    )
    if numbered_pairs:
        for front, back in numbered_pairs:
            card = _clean_card(front.strip(), back.strip())
            if card:
                cards.append(card)
        return cards

    # try bullet points like "- Term: Definition"
    term_defs = re.findall(
        r"[-•*]\s*\**(.+?)\**\s*[:\-–]\s*(.+)",
        raw_text,
    )
    if term_defs:
        for front, back in term_defs:
            card = _clean_card(front.strip(), back.strip())
            if card:
                cards.append(card)
        return cards

    return cards


def _clean_card(front, back):
    front = front.strip().rstrip(":")
    back = back.strip()
    if not front or not back:
        return None
    if len(front) < 2 or len(back) < 2:
        return None
    return {"front": front, "back": back}


def format_flashcards(cards):
    if not cards:
        return ""

    lines = ["", "  ╔══════════════════════════════════════════╗"]
    lines.append("  ║           🃏  FLASHCARDS                 ║")
    lines.append("  ╚══════════════════════════════════════════╝")
    lines.append("")

    for index, card in enumerate(cards, start=1):
        lines.append(f"  Card {index}")
        lines.append(f"    Front: {card['front']}")
        lines.append(f"    Back:  {card['back']}")
        lines.append("")

    lines.append("  " + "─" * 44)
    return "\n".join(lines)


def save_progress(data_dir, mode, structured_data, user_prompt):
    if not structured_data:
        return

    progress_dir = Path(data_dir) / "progress"
    progress_dir.mkdir(parents=True, exist_ok=True)

    entry = {
        "timestamp": datetime.now().isoformat(timespec="seconds"),
        "mode": mode,
        "prompt": user_prompt[:200],
        "item_count": len(structured_data),
        "items": structured_data,
    }

    file_name = f"{mode}_history.jsonl"
    file_path = progress_dir / file_name

    try:
        with file_path.open("a", encoding="utf-8") as progress_file:
            progress_file.write(json.dumps(entry, ensure_ascii=False) + "\n")
    except OSError:
        pass


def process_structured_output(raw_text, mode, data_dir, user_prompt):
    if mode == "quiz":
        questions = parse_quiz(raw_text)
        if questions:
            formatted = format_quiz(questions)
            save_progress(data_dir, mode, questions, user_prompt)
            return formatted, questions
        return raw_text, []

    if mode == "flashcard":
        cards = parse_flashcards(raw_text)
        if cards:
            formatted = format_flashcards(cards)
            save_progress(data_dir, mode, cards, user_prompt)
            return formatted, cards
        return raw_text, []

    return raw_text, []