src/streamlit_app.py

import os
import re
import json
import html
import pickle
import sqlite3
import hashlib
import uuid
import base64
from difflib import SequenceMatcher
from datetime import datetime
from typing import Dict, List, Any, Optional, Tuple

import numpy as np
import pandas as pd
import streamlit as st
import plotly.express as px
from openai import OpenAI
from rank_bm25 import BM25Okapi
from sentence_transformers import SentenceTransformer

# =====================================================
# CONFIGURATION
# =====================================================
APP_TITLE = "BrainChat PMQSN"
BASE_DIR = "src"
BUILD_DIR = os.path.join(BASE_DIR, "brainchat_build")
CHUNKS_PATH = os.path.join(BUILD_DIR, "chunks.pkl")
TOKENS_PATH = os.path.join(BUILD_DIR, "tokenized_chunks.pkl")
EMBED_PATH = os.path.join(BUILD_DIR, "embeddings.npy")
CONFIG_PATH = os.path.join(BUILD_DIR, "config.json")
QUESTION_BANK_FILE = os.path.join(BASE_DIR, "exam_questions_pmqs.json")
LOGO_FILE = os.path.join(BASE_DIR, "logo.png")
SOURCE_ALIASES_FILE = os.path.join(BASE_DIR, "source_aliases.json")
DB_PATH = os.getenv("BRAINCHAT_DB", "brainchat.db")
OPENAI_MODEL = os.getenv("OPENAI_MODEL", "gpt-4o-mini")
OPENAI_IMAGE_MODEL = os.getenv("OPENAI_IMAGE_MODEL", "gpt-image-1")
ENABLE_AI_IMAGES = os.getenv("ENABLE_AI_IMAGES", "true").lower() == "true"
TEACHER_PASSWORD = os.getenv("TEACHER_PASSWORD", "teacher123")

TOPICS = [
    "Stroke / Cerebrovascular",
    "Epilepsy",
    "Headache",
    "Multiple Sclerosis / Demyelination",
    "Parkinson / Movement Disorders",
    "Dementia",
    "Neuropathy / Neuromuscular",
    "Neuroanatomy / Topography",
    "General Neurology",
]

DEPTH_LEVELS = ["Basic", "Intermediate", "Advanced"]
QUIZ_DIFFICULTIES = ["Easy", "Medium", "Exam level"]
QUESTION_COUNTS = [3, 5, 10, 15, 20]

SOURCE_PRIORITY = {
    "official": 0.50,
    "course": 0.20,
    "external": 0.00,
}

SOURCE_LABELS = {
    "official": "Official Neurology Guiones",
    "course": "Other Course Material",
    "external": "Supplementary External Source",
}

DEPTH_INSTRUCTIONS = {
    "Basic": """
Use beginner-friendly language.
Structure the response as:
1. Simple definition
2. Essential concepts
3. Main symptoms or clinical features
4. One clear example
5. Five short revision points
Avoid unnecessary research terminology.
""",
    "Intermediate": """
Use standard medical curriculum depth.
Structure the response as:
1. Definition
2. Relevant anatomy and pathophysiology
3. Clinical presentation
4. Diagnosis and investigations
5. Treatment or management
6. Differential diagnosis
7. Important examination points
""",
    "Advanced": """
Use detailed clinical and research-oriented depth.
Structure the response as:
1. Detailed mechanisms
2. Advanced diagnostic reasoning
3. Differential diagnosis
4. Current management principles
5. Complications and difficult cases
6. Areas of uncertainty or controversy
7. Important research or guideline considerations
Clearly distinguish established course knowledge from supplementary evidence.
""",
}

TRANSLATIONS = {
    "English": {
        "app_subtitle": "AI tutor and quiz platform for Neurology / PMQSN learning",
        "language": "Interface language",
        "mode": "Choose mode",
        "student_mode": "Student Mode",
        "teacher_mode": "Teacher Mode",
        "student_id": "Student ID",
        "student_name": "Student name",
        "topic": "Choose topic",
        "difficulty": "Quiz difficulty",
        "depth": "Explanation depth",
        "num_questions": "Number of MCQ questions",
        "start_quiz": "Generate quiz",
        "submit_quiz": "Submit quiz",
        "chat": "Tutor Chat",
        "quiz": "Topic Quiz",
        "report": "Learning Report",
        "teacher_password": "Teacher password",
        "login": "Open teacher dashboard",
        "download_html": "Download HTML report",
        "no_data": "No student data available yet.",
        "score": "Score",
        "weak_areas": "Weak areas",
        "badges": "Badges earned",
        "ask_question": "Write your neurology question here",
        "send": "Ask BrainChat",
        "saved": "Attempt saved successfully.",
        "sources": "Sources used",
        "evidence_mix": "Retrieved evidence composition",
        "activity": "Tutor activity",
    },
    "Spanish": {
        "app_subtitle": "Tutor de IA y plataforma de cuestionarios para Neurología / PMQSN",
        "language": "Idioma de la interfaz",
        "mode": "Elegir modo",
        "student_mode": "Modo estudiante",
        "teacher_mode": "Modo profesor",
        "student_id": "ID del estudiante",
        "student_name": "Nombre del estudiante",
        "topic": "Elegir tema",
        "difficulty": "Dificultad del cuestionario",
        "depth": "Nivel de explicación",
        "num_questions": "Número de preguntas tipo test",
        "start_quiz": "Generar cuestionario",
        "submit_quiz": "Enviar cuestionario",
        "chat": "Tutor Chat",
        "quiz": "Cuestionario por tema",
        "report": "Informe de aprendizaje",
        "teacher_password": "Contraseña del profesor",
        "login": "Abrir panel del profesor",
        "download_html": "Descargar informe HTML",
        "no_data": "Todavía no hay datos de estudiantes.",
        "score": "Puntuación",
        "weak_areas": "Áreas débiles",
        "badges": "Insignias obtenidas",
        "ask_question": "Escribe aquí tu pregunta de neurología",
        "send": "Preguntar a BrainChat",
        "saved": "Intento guardado correctamente.",
        "sources": "Fuentes utilizadas",
        "evidence_mix": "Composición de la evidencia recuperada",
        "activity": "Actividad del tutor",
    },
}

st.set_page_config(page_title=APP_TITLE, page_icon="🧠", layout="wide")

# =====================================================
# DATABASE AND MIGRATIONS
# =====================================================
def get_conn() -> sqlite3.Connection:
    conn = sqlite3.connect(DB_PATH, check_same_thread=False)
    conn.row_factory = sqlite3.Row
    return conn


def ensure_column(conn: sqlite3.Connection, table: str, column: str, definition: str) -> None:
    existing = {row[1] for row in conn.execute(f"PRAGMA table_info({table})").fetchall()}
    if column not in existing:
        conn.execute(f"ALTER TABLE {table} ADD COLUMN {column} {definition}")


def init_db() -> None:
    conn = get_conn()
    cur = conn.cursor()

    cur.execute("""
        CREATE TABLE IF NOT EXISTS students (
            student_id TEXT PRIMARY KEY,
            name TEXT,
            language TEXT,
            created_at TEXT
        )
    """)

    cur.execute("""
        CREATE TABLE IF NOT EXISTS quiz_attempts (
            id INTEGER PRIMARY KEY AUTOINCREMENT,
            student_id TEXT,
            student_name TEXT,
            language TEXT,
            topic TEXT,
            difficulty TEXT,
            score INTEGER,
            total INTEGER,
            percent REAL,
            confidence_color TEXT,
            weak_areas TEXT,
            badges TEXT,
            quiz_json TEXT,
            answers_json TEXT,
            created_at TEXT
        )
    """)

    cur.execute("""
        CREATE TABLE IF NOT EXISTS chat_logs (
            id INTEGER PRIMARY KEY AUTOINCREMENT,
            student_id TEXT,
            language TEXT,
            topic TEXT,
            question TEXT,
            answer TEXT,
            confidence_color TEXT,
            similarity REAL,
            created_at TEXT
        )
    """)

    cur.execute("""
        CREATE TABLE IF NOT EXISTS generated_questions (
            question_id TEXT PRIMARY KEY,
            student_id TEXT,
            language TEXT,
            topic TEXT,
            difficulty TEXT,
            question TEXT,
            options_json TEXT,
            correct_option TEXT,
            explanation TEXT,
            subtopic TEXT,
            source_refs_json TEXT,
            question_hash TEXT,
            status TEXT DEFAULT 'unreviewed',
            created_at TEXT
        )
    """)

    cur.execute("""
        CREATE TABLE IF NOT EXISTS question_reviews (
            id INTEGER PRIMARY KEY AUTOINCREMENT,
            question_id TEXT,
            reporter_type TEXT,
            reporter_id TEXT,
            issue_type TEXT,
            reporter_comment TEXT,
            professor_comment TEXT,
            corrected_question TEXT,
            corrected_options_json TEXT,
            corrected_answer TEXT,
            corrected_explanation TEXT,
            review_status TEXT DEFAULT 'pending',
            reviewer TEXT,
            created_at TEXT,
            reviewed_at TEXT
        )
    """)

    cur.execute("""
        CREATE TABLE IF NOT EXISTS approved_questions (
            id INTEGER PRIMARY KEY AUTOINCREMENT,
            question_id TEXT UNIQUE,
            topic TEXT,
            difficulty TEXT,
            question TEXT,
            options_json TEXT,
            correct_option TEXT,
            explanation TEXT,
            source_refs_json TEXT,
            approved_by TEXT,
            approved_at TEXT,
            active INTEGER DEFAULT 1
        )
    """)

    cur.execute("""
        CREATE TABLE IF NOT EXISTS rejected_question_patterns (
            id INTEGER PRIMARY KEY AUTOINCREMENT,
            question_id TEXT,
            topic TEXT,
            question TEXT,
            question_hash TEXT,
            reason TEXT,
            rejected_by TEXT,
            created_at TEXT,
            active INTEGER DEFAULT 1
        )
    """)

    cur.execute("""
        CREATE TABLE IF NOT EXISTS feedback_rules (
            id INTEGER PRIMARY KEY AUTOINCREMENT,
            topic TEXT,
            rule_text TEXT,
            source_question_id TEXT,
            decision_type TEXT,
            created_by TEXT,
            created_at TEXT,
            active INTEGER DEFAULT 1
        )
    """)

    # Migrate older installations without deleting data.
    ensure_column(conn, "quiz_attempts", "source_refs_json", "TEXT")
    ensure_column(conn, "quiz_attempts", "source_mix_json", "TEXT")
    ensure_column(conn, "chat_logs", "depth_level", "TEXT")
    ensure_column(conn, "chat_logs", "source_refs_json", "TEXT")
    ensure_column(conn, "chat_logs", "source_mix_json", "TEXT")

    conn.commit()
    conn.close()


def now_iso() -> str:
    return datetime.now().isoformat(timespec="seconds")


def upsert_student(student_id: str, name: str, language: str) -> None:
    sid = (student_id or "Guest").strip() or "Guest"
    display_name = (name or sid).strip() or sid
    conn = get_conn()
    conn.execute("""
        INSERT INTO students(student_id, name, language, created_at)
        VALUES (?, ?, ?, ?)
        ON CONFLICT(student_id) DO UPDATE SET
            name=excluded.name,
            language=excluded.language
    """, (sid, display_name, language, now_iso()))
    conn.commit()
    conn.close()


def save_chat_log(
    student_id: str,
    language: str,
    topic: str,
    question: str,
    answer: str,
    confidence_color: str,
    similarity: float,
    depth_level: str,
    source_refs: List[Dict[str, Any]],
    source_mix: Dict[str, float],
) -> None:
    conn = get_conn()
    conn.execute("""
        INSERT INTO chat_logs(
            student_id, language, topic, question, answer,
            confidence_color, similarity, created_at,
            depth_level, source_refs_json, source_mix_json
        ) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)
    """, (
        student_id, language, topic, question, answer,
        confidence_color, similarity, now_iso(), depth_level,
        json.dumps(source_refs, ensure_ascii=False),
        json.dumps(source_mix, ensure_ascii=False),
    ))
    conn.commit()
    conn.close()


def save_quiz_attempt(
    student_id: str,
    name: str,
    language: str,
    topic: str,
    difficulty: str,
    score: int,
    total: int,
    confidence_color: str,
    weak_areas: List[str],
    badges: List[str],
    quiz: List[Dict[str, Any]],
    answers: Dict[str, str],
    source_refs: Optional[List[Dict[str, Any]]] = None,
    source_mix: Optional[Dict[str, float]] = None,
) -> None:
    upsert_student(student_id, name, language)
    percent = round((score / max(total, 1)) * 100, 2)
    conn = get_conn()
    conn.execute("""
        INSERT INTO quiz_attempts(
            student_id, student_name, language, topic, difficulty,
            score, total, percent, confidence_color, weak_areas,
            badges, quiz_json, answers_json, created_at,
            source_refs_json, source_mix_json
        ) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)
    """, (
        student_id, name, language, topic, difficulty,
        score, total, percent, confidence_color,
        json.dumps(weak_areas, ensure_ascii=False),
        json.dumps(badges, ensure_ascii=False),
        json.dumps(quiz, ensure_ascii=False),
        json.dumps(answers, ensure_ascii=False),
        now_iso(),
        json.dumps(source_refs or [], ensure_ascii=False),
        json.dumps(source_mix or {}, ensure_ascii=False),
    ))
    conn.commit()
    conn.close()


def load_attempts_df() -> pd.DataFrame:
    conn = get_conn()
    try:
        return pd.read_sql_query("SELECT * FROM quiz_attempts ORDER BY created_at DESC", conn)
    finally:
        conn.close()


def load_chat_df() -> pd.DataFrame:
    conn = get_conn()
    try:
        return pd.read_sql_query("SELECT * FROM chat_logs ORDER BY created_at DESC", conn)
    finally:
        conn.close()

# =====================================================
# SOURCE NORMALISATION AND RAG
# =====================================================
def tokenize(text: str) -> List[str]:
    return re.findall(r"\w+", (text or "").lower(), flags=re.UNICODE)


@st.cache_data(show_spinner=False)
def load_source_aliases() -> Dict[str, Any]:
    """Load optional filename-to-display-name mappings.

    Keys are matched as case-insensitive substrings against all source metadata.
    This lets a generic build filename such as ``ilovepdf_merged.pdf`` be shown
    with its real course title without rebuilding the application code.
    """
    if not os.path.exists(SOURCE_ALIASES_FILE):
        return {}
    try:
        with open(SOURCE_ALIASES_FILE, "r", encoding="utf-8") as f:
            data = json.load(f)
        return data if isinstance(data, dict) else {}
    except Exception:
        return {}


def _source_candidates(record: Dict[str, Any]) -> List[str]:
    fields = [
        "document_title", "source_title", "title", "filename", "file_name",
        "pdf_name", "document_name", "source", "path", "book", "collection",
    ]
    values: List[str] = []
    for field in fields:
        value = str(record.get(field, "")).strip()
        if value and value not in values:
            values.append(value)
    return values


def clean_source_name(book_name: str) -> str:
    name = os.path.basename((book_name or "").strip())
    if name.lower().endswith(".pdf"):
        name = name[:-4]
    readable = re.sub(r"[_-]+", " ", name).strip()
    low = readable.lower()

    if any(x in low for x in ["guiones", "guion neurolog", "neurology guideline", "neurología"]):
        return "Guiones de Neurología"
    if any(x in low for x in ["professor", "teacher", "lecture", "handout", "course notes"]):
        return "Professor Handouts"
    if any(x in low for x in ["ilovepdf", "i love pdf", "merged", "combinepdf", "combined pdf"]):
        return "Course Material (merged document)"
    return readable or "Course Material"


def resolve_source_metadata(record: Dict[str, Any]) -> Tuple[str, str, str]:
    candidates = _source_candidates(record)
    raw_name = candidates[0] if candidates else "Course Material"
    searchable = " ".join(candidates).lower().replace("_", " ").replace("-", " ")

    for pattern, metadata in load_source_aliases().items():
        if str(pattern).lower() not in searchable:
            continue

        rules = metadata if isinstance(metadata, list) else [metadata]
        fallback_rule: Optional[Dict[str, Any]] = None
        for rule in rules:
            if not isinstance(rule, dict):
                continue
            has_page_rule = "page_start" in rule or "page_end" in rule
            if not has_page_rule:
                fallback_rule = rule
                continue

            try:
                record_start = int(record.get("page_start", 0) or 0)
                record_end = int(record.get("page_end", record_start) or record_start)
                rule_start = int(rule.get("page_start", 1) or 1)
                rule_end = int(rule.get("page_end", 10**9) or 10**9)
                page_match = record_end >= rule_start and record_start <= rule_end
            except (TypeError, ValueError):
                page_match = False

            if page_match:
                display_name = str(rule.get("display_name", "")).strip() or clean_source_name(raw_name)
                source_type = str(rule.get("source_type", "course")).strip().lower()
                if source_type not in SOURCE_PRIORITY:
                    source_type = "course"
                return display_name, source_type, raw_name

        if fallback_rule:
            display_name = str(fallback_rule.get("display_name", "")).strip() or clean_source_name(raw_name)
            source_type = str(fallback_rule.get("source_type", "course")).strip().lower()
            if source_type not in SOURCE_PRIORITY:
                source_type = "course"
            return display_name, source_type, raw_name

    explicit = str(record.get("source_type", "")).strip().lower()
    if explicit in SOURCE_PRIORITY:
        source_type = explicit
    elif any(x in searchable for x in ["guiones", "guion neurolog", "official script", "official course"]):
        source_type = "official"
    elif any(x in searchable for x in ["professor", "teacher", "lecture", "handout", "course", "notes", "pmqsn", "merged"]):
        source_type = "course"
    else:
        source_type = "external"

    return clean_source_name(raw_name), source_type, raw_name


def clean_section_title(section_title: Any) -> str:
    section = str(section_title or "").strip()
    if not section:
        return ""
    if re.fullmatch(r"section[_\s-]*\d+", section, flags=re.I):
        return ""
    return re.sub(r"[_]+", " ", section).strip()


def expand_short_query(query: str) -> str:
    q = (query or "").strip()
    q_lower = q.lower()
    expansions = {
        "mri": "MRI magnetic resonance imaging resonancia magnética RM neuroimaging brain scan",
        "rm": "RM MRI resonancia magnética magnetic resonance imaging neuroimaging brain scan",
        "ct": "CT computed tomography tomografía computarizada TC brain scan",
        "tc": "TC CT tomografía computarizada computed tomography brain scan",
        "csf": "CSF cerebrospinal fluid LCR líquido cefalorraquídeo",
        "lcr": "LCR líquido cefalorraquídeo CSF cerebrospinal fluid",
        "eeg": "EEG electroencephalography electroencefalograma epilepsy seizure crisis",
    }
    return expansions.get(q_lower, q)


@st.cache_resource(show_spinner=False)
def load_rag_resources():
    required = [CHUNKS_PATH, TOKENS_PATH, EMBED_PATH, CONFIG_PATH]
    missing = [p for p in required if not os.path.exists(p)]
    if missing:
        return None, None, None, None, f"Missing course build files: {', '.join(missing)}"

    with open(CHUNKS_PATH, "rb") as f:
        chunks = pickle.load(f)
    with open(TOKENS_PATH, "rb") as f:
        tokenized_chunks = pickle.load(f)
    embeddings = np.load(EMBED_PATH)
    with open(CONFIG_PATH, "r", encoding="utf-8") as f:
        cfg = json.load(f)

    bm25 = BM25Okapi(tokenized_chunks)
    embed_model = SentenceTransformer(cfg["embedding_model"])
    return chunks, embeddings, bm25, embed_model, None


@st.cache_resource(show_spinner=False)
def get_client() -> Optional[OpenAI]:
    api_key = os.getenv("OPENAI_API_KEY")
    return OpenAI(api_key=api_key) if api_key else None


def search_hybrid(query: str, final_k: int = 8) -> Tuple[List[Dict[str, Any]], Optional[str]]:
    chunks, embeddings, bm25, embed_model, err = load_rag_resources()
    if err:
        return [], err

    expanded_query = expand_short_query(query)
    q_tokens = tokenize(expanded_query)
    bm25_scores = bm25.get_scores(q_tokens)
    shortlist_idx = np.argsort(bm25_scores)[::-1][:60]
    shortlist_emb = embeddings[shortlist_idx]
    qvec = embed_model.encode([expanded_query], normalize_embeddings=True).astype("float32")[0]
    dense_scores = shortlist_emb @ qvec

    results: List[Dict[str, Any]] = []
    for idx, dense_score in zip(shortlist_idx, dense_scores):
        r = chunks[int(idx)].copy()
        display_name, source_type, raw_name = resolve_source_metadata(r)
        r["source_original"] = raw_name
        r["book"] = display_name
        r["section_title"] = clean_section_title(r.get("section_title", ""))
        bm25_score = float(bm25_scores[idx])
        bm25_norm = min(max(bm25_score, 0.0) / 10.0, 0.20)

        # Source priority applies only when the passage has minimum semantic relevance.
        priority_boost = SOURCE_PRIORITY[source_type] if float(dense_score) >= 0.25 else 0.0
        final_score = float(dense_score) + bm25_norm + priority_boost

        r["source_type"] = source_type
        r["similarity_score"] = float(dense_score)
        r["bm25_score"] = bm25_score
        r["final_score"] = final_score
        results.append(r)

    results.sort(key=lambda x: x.get("final_score", 0.0), reverse=True)

    # Keep only the best chunk for each displayed document/page range.
    # Chunking often creates overlapping text windows with identical pages;
    # showing all of them makes one document look like several sources.
    selected: List[Dict[str, Any]] = []
    seen = set()
    for r in results:
        key = (
            str(r.get("book", "")).strip().lower(),
            str(r.get("source_type", "")).strip().lower(),
            str(r.get("page_start", "")).strip(),
            str(r.get("page_end", "")).strip(),
            clean_section_title(r.get("section_title", "")).lower(),
        )
        if key in seen:
            continue
        seen.add(key)
        selected.append(r)
        if len(selected) >= final_k:
            break

    return selected, None


def build_context(records: List[Dict[str, Any]]) -> str:
    blocks = []
    for i, r in enumerate(records, 1):
        blocks.append(
            f"""[Source {i}]
Book: {r.get('book', 'Course Material')}
Source category: {SOURCE_LABELS.get(r.get('source_type', 'external'), 'Supplementary External Source')}
Section: {r.get('section_title', '')}
Pages: {r.get('page_start', '')}-{r.get('page_end', '')}
Similarity: {r.get('similarity_score', 0):.3f}
Text:
{str(r.get('text', ''))[:3000]}"""
        )
    return "\n\n".join(blocks)


def compact_source_refs(records: List[Dict[str, Any]]) -> List[Dict[str, Any]]:
    refs = []
    for i, r in enumerate(records, 1):
        refs.append({
            "source_number": i,
            "book": r.get("book", "Course Material"),
            "source_type": r.get("source_type", "external"),
            "source_label": SOURCE_LABELS.get(r.get("source_type", "external"), "Supplementary External Source"),
            "section_title": r.get("section_title", ""),
            "page_start": r.get("page_start", ""),
            "page_end": r.get("page_end", ""),
            "similarity_score": round(float(r.get("similarity_score", 0)), 3),
            "final_score": round(float(r.get("final_score", 0)), 3),
        })
    return refs


def calculate_source_mix(records: List[Dict[str, Any]]) -> Dict[str, float]:
    totals = {"official": 0.0, "course": 0.0, "external": 0.0}
    for r in records:
        source_type = r.get("source_type", "external")
        relevance = max(float(r.get("similarity_score", 0)), 0.01)
        text_length = max(min(len(str(r.get("text", ""))), 3000), 1)
        totals[source_type] += relevance * text_length

    denominator = sum(totals.values())
    if denominator <= 0:
        return {"official": 0.0, "course": 0.0, "external": 0.0}

    rounded = {k: round(v / denominator * 100, 1) for k, v in totals.items()}
    # Keep displayed total at 100 after rounding without creating a negative category.
    delta = round(100.0 - sum(rounded.values()), 1)
    largest_key = max(rounded, key=rounded.get)
    rounded[largest_key] = round(rounded[largest_key] + delta, 1)
    return rounded


def confidence_from_similarity(similarity: float) -> str:
    if similarity >= 0.58:
        return "green"
    if similarity >= 0.42:
        return "orange"
    return "red"

# =====================================================
# QUESTION BANK, APPROVAL MEMORY AND REJECTION MEMORY
# =====================================================
@st.cache_data(show_spinner=False)
def load_question_bank() -> List[Dict[str, Any]]:
    if not os.path.exists(QUESTION_BANK_FILE):
        return []
    try:
        with open(QUESTION_BANK_FILE, "r", encoding="utf-8") as f:
            data = json.load(f)
        return data if isinstance(data, list) else []
    except Exception:
        return []


def detect_topic(text: str) -> str:
    t = (text or "").lower()
    topics = {
        "Stroke / Cerebrovascular": ["stroke", "ictus", "acm", "mca", "reperfusion", "trombol", "carótida", "hemipares", "afasia", "aspects", "vascular"],
        "Epilepsy": ["epile", "seizure", "crisis", "convuls", "eeg", "antiepil", "valpro", "levetiracetam"],
        "Headache": ["headache", "cefalea", "migraine", "migraña", "racimos", "trigémino", "cluster"],
        "Multiple Sclerosis / Demyelination": ["multiple sclerosis", "esclerosis", "desmiel", "nmosd", "neuromielitis", "lcr", "oligoclon"],
        "Parkinson / Movement Disorders": ["parkinson", "temblor", "bradicinesia", "levodopa", "diston", "movimiento", "supranuclear", "multisist"],
        "Dementia": ["dementia", "demencia", "alzheimer", "cognit", "memoria", "alucinaciones", "lewy"],
        "Neuropathy / Neuromuscular": ["neurop", "miasten", "myasthen", "guillain", "ela", "motoneur", "fascicul", "miopat"],
        "Neuroanatomy / Topography": ["topograf", "localiza", "lesion", "lesión", "médula", "tronco", "arteria", "quiasma", "reflejo", "sensibilidad"],
    }
    for topic, keys in topics.items():
        if any(k in t for k in keys):
            return topic
    return "General Neurology"


def filter_question_examples(topic: str, limit: int = 6) -> List[Dict[str, Any]]:
    bank = load_question_bank()
    matches = []
    for q in bank:
        option_text = " ".join(
            o.get("text", "") if isinstance(o, dict) else str(o)
            for o in q.get("options", [])
        )
        if detect_topic(q.get("question", "") + " " + option_text) == topic:
            matches.append(q)
    return (matches or bank)[:limit]


def normalise_question_text(question: str) -> str:
    text = re.sub(r"[^\w\s]", " ", (question or "").lower(), flags=re.UNICODE)
    return re.sub(r"\s+", " ", text).strip()


def question_hash(question: str) -> str:
    return hashlib.sha256(normalise_question_text(question).encode("utf-8")).hexdigest()


def token_jaccard(a: str, b: str) -> float:
    a_set = set(tokenize(normalise_question_text(a)))
    b_set = set(tokenize(normalise_question_text(b)))
    if not a_set or not b_set:
        return 0.0
    return len(a_set & b_set) / len(a_set | b_set)


def load_rejected_patterns(topic: str) -> List[Dict[str, Any]]:
    conn = get_conn()
    rows = conn.execute("""
        SELECT question, question_hash, reason
        FROM rejected_question_patterns
        WHERE active=1 AND (topic=? OR topic='General Neurology')
        ORDER BY created_at DESC
    """, (topic,)).fetchall()
    conn.close()
    return [dict(r) for r in rows]


def is_rejected_or_too_similar(question: str, rejected: List[Dict[str, Any]]) -> bool:
    q_hash = question_hash(question)
    normalised = normalise_question_text(question)
    for item in rejected:
        rejected_question = item.get("question", "")
        if item.get("question_hash") == q_hash:
            return True
        if token_jaccard(question, rejected_question) >= 0.68:
            return True
        if SequenceMatcher(None, normalised, normalise_question_text(rejected_question)).ratio() >= 0.86:
            return True
    return False


def load_feedback_rules(topic: str, limit: int = 30) -> List[Dict[str, Any]]:
    conn = get_conn()
    rows = conn.execute("""
        SELECT id, topic, rule_text, source_question_id, decision_type, created_by, created_at
        FROM feedback_rules
        WHERE active=1 AND (topic=? OR topic='General Neurology')
        ORDER BY created_at DESC
        LIMIT ?
    """, (topic, limit)).fetchall()
    conn.close()
    return [dict(r) for r in rows]


def save_feedback_rule(
    topic: str,
    rule_text: str,
    question_id: str,
    decision_type: str,
    reviewer: str,
) -> None:
    cleaned = re.sub(r"\s+", " ", (rule_text or "").strip())
    if not cleaned:
        return
    conn = get_conn()
    duplicate = conn.execute("""
        SELECT id FROM feedback_rules
        WHERE active=1 AND topic=? AND lower(rule_text)=lower(?)
    """, (topic, cleaned)).fetchone()
    if not duplicate:
        conn.execute("""
            INSERT INTO feedback_rules(
                topic, rule_text, source_question_id, decision_type,
                created_by, created_at, active
            ) VALUES (?, ?, ?, ?, ?, ?, 1)
        """, (topic, cleaned, question_id, decision_type, reviewer, now_iso()))
        conn.commit()
    conn.close()


def load_generated_questions_df(limit: int = 500) -> pd.DataFrame:
    conn = get_conn()
    try:
        return pd.read_sql_query("""
            SELECT question_id, topic, difficulty, question, options_json,
                   correct_option, explanation, subtopic, source_refs_json,
                   status, created_at, student_id, language
            FROM generated_questions
            ORDER BY created_at DESC
            LIMIT ?
        """, conn, params=(limit,))
    finally:
        conn.close()


def load_approved_questions(topic: str, difficulty: str, limit: int = 12) -> List[Dict[str, Any]]:
    conn = get_conn()
    rows = conn.execute("""
        SELECT * FROM approved_questions
        WHERE active=1 AND topic=? AND (difficulty=? OR difficulty='Any')
        ORDER BY approved_at DESC
        LIMIT ?
    """, (topic, difficulty, limit)).fetchall()
    conn.close()

    output = []
    for r in rows:
        item = dict(r)
        item["options"] = json.loads(item.get("options_json") or "[]")
        item["source_refs"] = json.loads(item.get("source_refs_json") or "[]")
        output.append(item)
    return output


def save_generated_questions(
    questions: List[Dict[str, Any]],
    student_id: str,
    language: str,
    topic: str,
    difficulty: str,
) -> None:
    conn = get_conn()
    for item in questions:
        conn.execute("""
            INSERT OR IGNORE INTO generated_questions(
                question_id, student_id, language, topic, difficulty,
                question, options_json, correct_option, explanation,
                subtopic, source_refs_json, question_hash, status, created_at
            ) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)
        """, (
            item["question_id"], student_id, language, topic, difficulty,
            item["question"], json.dumps(item["options"], ensure_ascii=False),
            item["correct_option"], item["explanation"], item.get("subtopic", topic),
            json.dumps(item.get("source_refs", []), ensure_ascii=False),
            question_hash(item["question"]), item.get("status", "unreviewed"), now_iso(),
        ))
    conn.commit()
    conn.close()


def create_question_review(
    question: Dict[str, Any],
    reporter_type: str,
    reporter_id: str,
    issue_type: str,
    comment: str,
) -> bool:
    qid = question.get("question_id") or str(uuid.uuid4())
    conn = get_conn()

    existing = conn.execute("SELECT question_id FROM generated_questions WHERE question_id=?", (qid,)).fetchone()
    if not existing:
        conn.execute("""
            INSERT INTO generated_questions(
                question_id, student_id, language, topic, difficulty,
                question, options_json, correct_option, explanation,
                subtopic, source_refs_json, question_hash, status, created_at
            ) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)
        """, (
            qid, reporter_id, question.get("language", ""),
            question.get("topic", "General Neurology"), question.get("difficulty", "Any"),
            question.get("question", ""), json.dumps(question.get("options", []), ensure_ascii=False),
            question.get("correct_option", "A"), question.get("explanation", ""),
            question.get("subtopic", ""), json.dumps(question.get("source_refs", []), ensure_ascii=False),
            question_hash(question.get("question", "")), "flagged", now_iso(),
        ))
    else:
        conn.execute("UPDATE generated_questions SET status='flagged' WHERE question_id=?", (qid,))

    pending = conn.execute("""
        SELECT id FROM question_reviews
        WHERE question_id=? AND review_status='pending'
    """, (qid,)).fetchone()
    if pending:
        conn.commit()
        conn.close()
        return False

    conn.execute("""
        INSERT INTO question_reviews(
            question_id, reporter_type, reporter_id, issue_type,
            reporter_comment, review_status, created_at
        ) VALUES (?, ?, ?, ?, ?, 'pending', ?)
    """, (qid, reporter_type, reporter_id, issue_type, comment, now_iso()))
    conn.commit()
    conn.close()
    return True


def load_pending_reviews() -> pd.DataFrame:
    conn = get_conn()
    query = """
        SELECT
            r.id AS review_id,
            r.question_id,
            r.reporter_type,
            r.reporter_id,
            r.issue_type,
            r.reporter_comment,
            r.created_at AS reported_at,
            g.topic,
            g.difficulty,
            g.question,
            g.options_json,
            g.correct_option,
            g.explanation,
            g.subtopic,
            g.source_refs_json,
            g.status AS question_status
        FROM question_reviews r
        JOIN generated_questions g ON g.question_id = r.question_id
        WHERE r.review_status='pending'
        ORDER BY r.created_at ASC
    """
    try:
        return pd.read_sql_query(query, conn)
    finally:
        conn.close()


def approve_review(
    review_id: int,
    question_id: str,
    reviewer: str,
    professor_comment: str,
    corrected_question: str,
    corrected_options: List[str],
    corrected_answer: str,
    corrected_explanation: str,
) -> None:
    conn = get_conn()
    row = conn.execute("SELECT * FROM generated_questions WHERE question_id=?", (question_id,)).fetchone()
    if not row:
        conn.close()
        raise ValueError("Generated question not found.")

    source_refs_json = row["source_refs_json"] or "[]"
    conn.execute("""
        INSERT INTO approved_questions(
            question_id, topic, difficulty, question, options_json,
            correct_option, explanation, source_refs_json,
            approved_by, approved_at, active
        ) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, 1)
        ON CONFLICT(question_id) DO UPDATE SET
            topic=excluded.topic,
            difficulty=excluded.difficulty,
            question=excluded.question,
            options_json=excluded.options_json,
            correct_option=excluded.correct_option,
            explanation=excluded.explanation,
            source_refs_json=excluded.source_refs_json,
            approved_by=excluded.approved_by,
            approved_at=excluded.approved_at,
            active=1
    """, (
        question_id, row["topic"], row["difficulty"], corrected_question,
        json.dumps(corrected_options, ensure_ascii=False), corrected_answer,
        corrected_explanation, source_refs_json, reviewer, now_iso(),
    ))

    conn.execute("""
        UPDATE question_reviews SET
            professor_comment=?, corrected_question=?, corrected_options_json=?,
            corrected_answer=?, corrected_explanation=?, review_status='approved',
            reviewer=?, reviewed_at=?
        WHERE id=?
    """, (
        professor_comment, corrected_question,
        json.dumps(corrected_options, ensure_ascii=False), corrected_answer,
        corrected_explanation, reviewer, now_iso(), review_id,
    ))
    conn.execute("UPDATE generated_questions SET status='approved' WHERE question_id=?", (question_id,))

    original_changed = (
        normalise_question_text(row["question"]) != normalise_question_text(corrected_question)
        or row["correct_option"] != corrected_answer
        or json.loads(row["options_json"] or "[]") != corrected_options
        or (row["explanation"] or "").strip() != (corrected_explanation or "").strip()
    )
    if original_changed:
        reason = professor_comment.strip() or "Original formulation replaced by a professor-corrected version."
        duplicate = conn.execute("""
            SELECT id FROM rejected_question_patterns
            WHERE active=1 AND question_hash=?
        """, (row["question_hash"],)).fetchone()
        if not duplicate:
            conn.execute("""
                INSERT INTO rejected_question_patterns(
                    question_id, topic, question, question_hash, reason,
                    rejected_by, created_at, active
                ) VALUES (?, ?, ?, ?, ?, ?, ?, 1)
            """, (
                question_id, row["topic"], row["question"], row["question_hash"],
                reason, reviewer, now_iso(),
            ))

    conn.commit()
    conn.close()
    save_feedback_rule(row["topic"], professor_comment, question_id, "approved_correction", reviewer)


def reject_review(
    review_id: int,
    question_id: str,
    reviewer: str,
    reason: str,
) -> None:
    conn = get_conn()
    row = conn.execute("SELECT * FROM generated_questions WHERE question_id=?", (question_id,)).fetchone()
    if not row:
        conn.close()
        raise ValueError("Generated question not found.")

    duplicate = conn.execute("""
        SELECT id FROM rejected_question_patterns
        WHERE active=1 AND question_hash=?
    """, (row["question_hash"],)).fetchone()
    if not duplicate:
        conn.execute("""
            INSERT INTO rejected_question_patterns(
                question_id, topic, question, question_hash, reason,
                rejected_by, created_at, active
            ) VALUES (?, ?, ?, ?, ?, ?, ?, 1)
        """, (
            question_id, row["topic"], row["question"], row["question_hash"],
            reason, reviewer, now_iso(),
        ))
    conn.execute("""
        UPDATE question_reviews SET
            professor_comment=?, review_status='rejected', reviewer=?, reviewed_at=?
        WHERE id=?
    """, (reason, reviewer, now_iso(), review_id))
    conn.execute("UPDATE generated_questions SET status='rejected' WHERE question_id=?", (question_id,))
    conn.execute("UPDATE approved_questions SET active=0 WHERE question_id=?", (question_id,))
    conn.commit()
    conn.close()
    save_feedback_rule(row["topic"], reason, question_id, "rejected", reviewer)


# =====================================================
# AI HELPERS
# =====================================================
def safe_json_from_text(text: str) -> Any:
    text = (text or "").strip()
    text = re.sub(r"^```json", "", text, flags=re.I).strip()
    text = re.sub(r"^```", "", text).strip()
    text = re.sub(r"```$", "", text).strip()
    start = text.find("[")
    end = text.rfind("]")
    if start != -1 and end != -1 and end > start:
        text = text[start:end + 1]
    return json.loads(text)


def strip_code_fences(text: str) -> str:
    text = (text or "").strip()
    text = re.sub(r"^```(?:dot|graphviz)?", "", text, flags=re.I).strip()
    text = re.sub(r"```$", "", text).strip()
    return text


def normalize_mcq_option(opt: Any, index: int) -> str:
    letter = chr(65 + index)
    if isinstance(opt, dict):
        opt_letter = str(opt.get("letter", letter)).strip().upper()[:1] or letter
        text = str(opt.get("text", opt.get("option", opt.get("value", "")))).strip()
        return f"{opt_letter}. {text or str(opt)}"

    text = str(opt).strip()
    match = re.match(r"^([A-Ea-e])\s*[\.|\)]\s*(.+)$", text)
    if match:
        return f"{match.group(1).upper()}. {match.group(2).strip()}"
    return f"{letter}. {text}"


def approved_to_quiz_item(item: Dict[str, Any]) -> Dict[str, Any]:
    return {
        "question_id": item.get("question_id") or str(uuid.uuid4()),
        "question": item.get("question", ""),
        "options": item.get("options", []),
        "correct_option": item.get("correct_option", "A"),
        "explanation": item.get("explanation", ""),
        "subtopic": item.get("topic", "General Neurology"),
        "source_refs": item.get("source_refs", []),
        "status": "approved",
    }


def fallback_mcqs(topic: str, n: int, language: str, source_refs: List[Dict[str, Any]]) -> List[Dict[str, Any]]:
    if language == "Spanish":
        q = f"Pregunta de práctica sobre {topic}: ¿cuál opción es más correcta?"
        exp = "Pregunta de demostración. Configure OPENAI_API_KEY y los materiales del curso para generar preguntas reales."
    else:
        q = f"Practice question on {topic}: which option is most correct?"
        exp = "Demo question. Configure OPENAI_API_KEY and course materials to generate real questions."

    return [{
        "question_id": str(uuid.uuid4()),
        "question": q,
        "options": ["A. Option A", "B. Option B", "C. Option C", "D. Option D", "E. Option E"],
        "correct_option": "A",
        "explanation": exp,
        "subtopic": topic,
        "source_refs": source_refs,
        "status": "demo",
    } for _ in range(n)]


def generate_mcqs(
    topic: str,
    difficulty: str,
    n_questions: int,
    language: str,
    student_id: str,
) -> Tuple[List[Dict[str, Any]], str, List[Dict[str, Any]], Dict[str, float]]:
    records, err = search_hybrid(topic + " neurology PMQSN exam questions", final_k=10)
    source_refs = compact_source_refs(records)
    source_mix = calculate_source_mix(records)
    context = build_context(records)

    approved = load_approved_questions(topic, difficulty, limit=n_questions)
    approved_items = [approved_to_quiz_item(x) for x in approved]
    rejected = load_rejected_patterns(topic)
    feedback_rules = load_feedback_rules(topic)

    # Use some approved questions directly so professor corrections affect the next quiz immediately.
    final_questions = approved_items[: min(len(approved_items), max(1, n_questions // 2))]
    needed = n_questions - len(final_questions)

    if err and not final_questions:
        quiz = fallback_mcqs(topic, n_questions, language, source_refs)
        save_generated_questions(quiz, student_id, language, topic, difficulty)
        return quiz, err, source_refs, source_mix

    client = get_client()
    if needed <= 0:
        save_generated_questions(final_questions[:n_questions], student_id, language, topic, difficulty)
        return final_questions[:n_questions], "", source_refs, source_mix

    if client is None:
        additions = fallback_mcqs(topic, needed, language, source_refs)
        final_questions.extend(additions)
        save_generated_questions(final_questions, student_id, language, topic, difficulty)
        return final_questions, "OPENAI_API_KEY missing. Approved and demo questions are shown.", source_refs, source_mix

    examples = filter_question_examples(topic, limit=6)
    approved_examples = [
        {
            "question": x.get("question"),
            "options": x.get("options"),
            "correct_option": x.get("correct_option"),
            "explanation": x.get("explanation"),
        }
        for x in approved[:8]
    ]
    rejected_guidance = [
        {"question": x.get("question", ""), "reason": x.get("reason", "")}
        for x in rejected[:10]
    ]
    professor_rules = [x.get("rule_text", "") for x in feedback_rules if x.get("rule_text")]

    lang_instruction = "Write everything in English." if language == "English" else "Escribe todo en español."
    requested = needed + 4

    prompt = f"""
You are BrainChat, an exam-focused neurology tutor.
Generate at least {requested} MCQs for: {topic}.
Difficulty: {difficulty}.
{lang_instruction}

Source hierarchy:
1. Official Neurology Guiones
2. Other course material
3. Supplementary external material only when necessary

Rules:
- Output ONLY a valid JSON array.
- Each item must contain: question, options, correct_option, explanation, subtopic.
- options must contain exactly five strings labelled A, B, C, D and E.
- There must be exactly one correct answer.
- Every question and explanation must be supported by the supplied course context.
- Avoid vague wording, trick wording and multiple defensible answers.
- Do not repeat any rejected question or its pattern.
- Do not mention files, retrieval, prompts or JSON.

Professor-approved examples:
{json.dumps(approved_examples, ensure_ascii=False)[:7000]}

Past exam-style examples:
{json.dumps(examples, ensure_ascii=False)[:6000]}

Rejected patterns and reasons:
{json.dumps(rejected_guidance, ensure_ascii=False)[:5000]}

Professor feedback rules learned from earlier reviews:
{json.dumps(professor_rules, ensure_ascii=False)[:5000]}

Course context:
{context}
"""

    warning = ""
    try:
        response = client.chat.completions.create(
            model=OPENAI_MODEL,
            messages=[{"role": "user", "content": prompt}],
            temperature=0.25,
        )
        raw_items = safe_json_from_text(response.choices[0].message.content or "[]")
    except Exception as exc:
        raw_items = []
        warning = f"AI generation failed: {exc}"

    generated: List[Dict[str, Any]] = []
    seen_hashes = {question_hash(q["question"]) for q in final_questions}

    for item in raw_items if isinstance(raw_items, list) else []:
        if len(generated) >= needed:
            break
        question_text = str(item.get("question", "")).strip()
        options = item.get("options", [])
        if isinstance(options, dict):
            options = [{"letter": k, "text": v} for k, v in options.items()]
        if not question_text or not isinstance(options, list):
            continue

        formatted_options = [normalize_mcq_option(opt, i) for i, opt in enumerate(options[:5])]
        if len(formatted_options) != 5:
            continue

        q_hash = question_hash(question_text)
        if q_hash in seen_hashes or is_rejected_or_too_similar(question_text, rejected):
            continue

        correct = str(item.get("correct_option", item.get("answer", "A"))).strip().upper()[:1]
        if correct not in "ABCDE":
            continue

        explanation = str(item.get("explanation", "")).strip()
        if not explanation:
            continue

        generated.append({
            "question_id": str(uuid.uuid4()),
            "question": question_text,
            "options": formatted_options,
            "correct_option": correct,
            "explanation": explanation,
            "subtopic": str(item.get("subtopic", topic)).strip() or topic,
            "source_refs": source_refs,
            "status": "unreviewed",
        })
        seen_hashes.add(q_hash)

    final_questions.extend(generated)

    if len(final_questions) < n_questions:
        missing = n_questions - len(final_questions)
        final_questions.extend(fallback_mcqs(topic, missing, language, source_refs))
        warning = warning or "Some demo questions were added because too few valid questions were generated."

    final_questions = final_questions[:n_questions]
    save_generated_questions(final_questions, student_id, language, topic, difficulty)
    return final_questions, warning, source_refs, source_mix


def answer_tutor_question(
    question: str,
    topic: str,
    language: str,
    depth_level: str,
) -> Tuple[str, str, float, List[Dict[str, Any]], Dict[str, float], Optional[str]]:
    records, err = search_hybrid(question + " " + topic, final_k=8)
    source_refs = compact_source_refs(records)
    source_mix = calculate_source_mix(records)
    similarity = max([r.get("similarity_score", 0) for r in records], default=0.0)
    color = confidence_from_similarity(similarity)

    if err:
        return err, "red", 0.0, source_refs, source_mix, err

    client = get_client()
    if client is None:
        msg = "OPENAI_API_KEY is missing. Add it in Hugging Face Space Secrets."
        return msg, "red", similarity, source_refs, source_mix, msg

    lang_instruction = "Answer fully in English." if language == "English" else "Responde completamente en español."
    depth_instruction = DEPTH_INSTRUCTIONS.get(depth_level, DEPTH_INSTRUCTIONS["Intermediate"])
    context = build_context(records)

    prompt = f"""
You are BrainChat, a neurology tutor. {lang_instruction}

Required explanation level: {depth_level}
{depth_instruction}

Evidence rules:
- Use the Official Neurology Guiones first.
- Use other course material second.
- Use supplementary external material only where the course material is insufficient.
- Add [Source 1], [Source 2], etc. after important factual statements.
- Do not invent page numbers or sources.
- If the supplied evidence does not support part of the question, state that clearly.
- End with a short revision summary and one revision tip.

Topic: {topic}
Question: {question}

Course context:
{context}
"""

    try:
        response = client.chat.completions.create(
            model=OPENAI_MODEL,
            messages=[{"role": "user", "content": prompt}],
            temperature=0.20,
        )
        answer = response.choices[0].message.content or ""
        return answer, color, similarity, source_refs, source_mix, None
    except Exception as exc:
        return f"AI response failed: {exc}", "red", similarity, source_refs, source_mix, str(exc)


def generate_dot_visual(
    topic: str,
    visual_type: str,
    depth_level: str,
    language: str,
) -> Tuple[str, List[Dict[str, Any]], Dict[str, float], Optional[str]]:
    records, err = search_hybrid(f"{topic} {visual_type} diagnosis management", final_k=8)
    source_refs = compact_source_refs(records)
    source_mix = calculate_source_mix(records)
    if err:
        return "", source_refs, source_mix, err

    client = get_client()
    if client is None:
        return "", source_refs, source_mix, "OPENAI_API_KEY is missing."

    lang_instruction = "Use English labels." if language == "English" else "Usa etiquetas en español."
    context = build_context(records)
    prompt = f"""
Create a clear {visual_type} for the neurology topic: {topic}.
Depth: {depth_level}.
{lang_instruction}

Return ONLY valid Graphviz DOT code beginning with digraph.
Rules:
- Use short node labels.
- Use simple top-to-bottom flow.
- Do not use colours.
- Do not use HTML labels.
- Keep the diagram educational and readable.
- Include only claims supported by the supplied context.
- For a clinical pathway, include decision diamonds only where a true decision exists.
- Do not include exact medication doses unless explicitly supported in the context.

Context:
{context}
"""
    try:
        response = client.chat.completions.create(
            model=OPENAI_MODEL,
            messages=[{"role": "user", "content": prompt}],
            temperature=0.15,
        )
        dot = strip_code_fences(response.choices[0].message.content or "")
        if not dot.lower().startswith("digraph"):
            return "", source_refs, source_mix, "The model did not return valid Graphviz DOT code."
        return dot, source_refs, source_mix, None
    except Exception as exc:
        return "", source_refs, source_mix, str(exc)


def generate_ai_medical_image(
    topic: str,
    depth_level: str,
    language: str,
    visual_focus: str,
    visual_style: str,
) -> Tuple[Optional[bytes], List[Dict[str, Any]], Dict[str, float], str, Optional[str]]:
    records, rag_error = search_hybrid(
        f"{topic} {visual_focus} anatomy mechanism diagnosis educational illustration",
        final_k=4,
    )
    source_refs = compact_source_refs(records)
    source_mix = calculate_source_mix(records)
    context = build_context(records)

    if rag_error:
        return None, source_refs, source_mix, "", rag_error
    if not ENABLE_AI_IMAGES:
        return None, source_refs, source_mix, "", (
            "AI image generation is disabled. Set ENABLE_AI_IMAGES=true to enable it."
        )

    client = get_client()
    if client is None:
        return None, source_refs, source_mix, "", "OPENAI_API_KEY is missing."

    label_language = "English" if language == "English" else "Spanish"
    focus = visual_focus.strip() or topic
    brief_prompt = f"""
You are preparing a source-grounded prompt for a medical education image generator.

Topic: {topic}
Requested focus: {focus}
Learner level: {depth_level}
Visual style: {visual_style}
Label language: {label_language}

Create one concise image-generation brief. Use only facts supported by the supplied course context.
The image must be educational, uncluttered and medically cautious. Prefer a simplified labelled mechanism, anatomy overview or process illustration. Avoid exact medication doses, diagnostic certainty, photorealistic patients, identifiable people and decorative imagery. Use minimal text because image models may misspell labels. Do not include citations inside the image.

Course context:
{context}
"""

    try:
        brief_response = client.chat.completions.create(
            model=OPENAI_MODEL,
            messages=[{"role": "user", "content": brief_prompt}],
            temperature=0.10,
        )
        visual_brief = (brief_response.choices[0].message.content or "").strip()
    except Exception:
        visual_brief = (
            f"Create a clean {visual_style.lower()} for {focus} within {topic}, suitable for "
            f"{depth_level.lower()} medical learners, using {label_language} labels."
        )

    final_prompt = (
        f"{visual_brief} White background, clear hierarchy, high-resolution educational medical illustration. "
        "No patient-identifying features. No diagnosis claim. No exact dosage. No decorative border. "
        "Use only a few large, legible labels. The output is an educational illustration, not a diagnostic image."
    )

    try:
        response = client.images.generate(
            model=OPENAI_IMAGE_MODEL,
            prompt=final_prompt,
            size="1024x1024",
        )
        item = response.data[0]
        if getattr(item, "b64_json", None):
            return base64.b64decode(item.b64_json), source_refs, source_mix, visual_brief, None
        return None, source_refs, source_mix, visual_brief, "The image API returned no image data."
    except Exception as exc:
        return None, source_refs, source_mix, visual_brief, str(exc)


# =====================================================
# BADGES AND REPORTS
# =====================================================
def badges_for_student(student_id: str) -> List[str]:
    df = load_attempts_df()
    if df.empty:
        return []
    sdf = df[df["student_id"] == student_id].copy()
    if sdf.empty:
        return []

    badges = set()
    for topic in TOPICS:
        topic_df = sdf[sdf["topic"] == topic]
        short = topic.split(" /")[0]
        if len(topic_df[topic_df["percent"] >= 70]) >= 2:
            badges.add(f"🥉 {short} Bronze")
        if len(topic_df[topic_df["percent"] >= 80]) >= 3:
            badges.add(f"🥈 {short} Silver")
        if len(topic_df[topic_df["percent"] >= 90]) >= 5:
            badges.add(f"🥇 {short} Gold")
    if len(sdf) >= 10:
        badges.add("📘 Consistent Learner")
    if len(sdf) >= 5 and float(sdf["percent"].mean()) >= 85:
        badges.add("🏆 Neurology Master")
    return sorted(badges)


def html_report_student(student_id: str, name: str, language: str) -> str:
    df = load_attempts_df()
    sdf = df[df["student_id"] == student_id] if not df.empty else pd.DataFrame()
    title = "Learning Report" if language == "English" else "Informe de aprendizaje"
    rows = ""
    if not sdf.empty:
        for _, r in sdf.iterrows():
            rows += (
                f"<tr><td>{html.escape(str(r['created_at']))}</td>"
                f"<td>{html.escape(str(r['topic']))}</td>"
                f"<td>{html.escape(str(r['difficulty']))}</td>"
                f"<td>{r['score']}/{r['total']}</td>"
                f"<td>{r['percent']:.1f}%</td>"
                f"<td>{html.escape(str(r['confidence_color']))}</td></tr>"
            )
    else:
        rows = "<tr><td colspan='6'>No attempts yet.</td></tr>"

    badges = ", ".join(badges_for_student(student_id)) or "None"
    avg = sdf["percent"].mean() if not sdf.empty else 0
    return f"""
<!doctype html>
<html><head><meta charset='utf-8'><title>{title}</title>
<style>
body{{font-family:Arial;margin:30px;line-height:1.5}}
.card{{border:1px solid #ddd;border-radius:12px;padding:18px;margin:12px 0}}
table{{border-collapse:collapse;width:100%}}
th,td{{border:1px solid #ddd;padding:8px;text-align:left}}
th{{background:#f3f3f3}}
</style></head><body>
<h1>{title}</h1>
<div class='card'><b>Student:</b> {html.escape(name)}<br><b>ID:</b> {html.escape(student_id)}<br><b>Generated:</b> {datetime.now().strftime('%Y-%m-%d %H:%M')}</div>
<div class='card'><h2>Summary</h2><p><b>Average score:</b> {avg:.1f}%</p><p><b>Badges:</b> {html.escape(badges)}</p></div>
<div class='card'><h2>Quiz attempts</h2><table><tr><th>Date</th><th>Topic</th><th>Difficulty</th><th>Score</th><th>Percent</th><th>Confidence</th></tr>{rows}</table></div>
</body></html>
"""


def html_report_teacher() -> str:
    df = load_attempts_df()
    if df.empty:
        body = "<p>No student data available yet.</p>"
    else:
        summary = df.groupby("topic").agg(attempts=("id", "count"), avg_score=("percent", "mean")).reset_index()
        body = "<h2>Topic summary</h2><table><tr><th>Topic</th><th>Attempts</th><th>Average score</th></tr>"
        for _, r in summary.iterrows():
            body += f"<tr><td>{html.escape(str(r['topic']))}</td><td>{int(r['attempts'])}</td><td>{r['avg_score']:.1f}%</td></tr>"
        body += "</table>"
    return f"""
<!doctype html><html><head><meta charset='utf-8'><title>Teacher Class Report</title>
<style>body{{font-family:Arial;margin:30px;line-height:1.5}}table{{border-collapse:collapse;width:100%}}th,td{{border:1px solid #ddd;padding:8px;text-align:left}}th{{background:#f3f3f3}}</style></head>
<body><h1>Teacher Class Report</h1><p>Generated: {datetime.now().strftime('%Y-%m-%d %H:%M')}</p>{body}</body></html>
"""

# =====================================================
# UI HELPERS
# =====================================================
def t(key: str) -> str:
    lang = st.session_state.get("language", "English")
    return TRANSLATIONS[lang].get(key, key)


def render_header() -> None:
    c1, c2 = st.columns([1, 5])
    with c1:
        if os.path.exists(LOGO_FILE):
            st.image(LOGO_FILE, width=90)
        else:
            st.markdown("# 🧠")
    with c2:
        st.title(APP_TITLE)
        st.caption(t("app_subtitle"))


def confidence_badge(color: str) -> None:
    label = {"green": "🟢 Green", "orange": "🟠 Orange", "red": "🔴 Red"}.get(color, color)
    st.markdown(f"**Confidence:** {label}")


def format_option(opt: Any) -> str:
    if isinstance(opt, dict):
        letter = str(opt.get("letter", "")).strip()
        text = str(opt.get("text", "")).strip()
        return f"{letter}. {text}" if letter and text else text or str(opt)
    return str(opt)


def display_source_mix(source_mix: Dict[str, float]) -> None:
    st.markdown(f"#### {t('evidence_mix')}")
    c1, c2, c3 = st.columns(3)
    c1.metric("Official Guiones", f"{source_mix.get('official', 0):.1f}%")
    c2.metric("Other course material", f"{source_mix.get('course', 0):.1f}%")
    c3.metric("Supplementary", f"{source_mix.get('external', 0):.1f}%")
    st.caption(
        "These percentages describe the retrieved supporting passages, not an exact count of generated words."
        if st.session_state.get("language") == "English"
        else "Estos porcentajes describen los pasajes de apoyo recuperados, no un recuento exacto de las palabras generadas."
    )


def display_sources(source_refs: List[Dict[str, Any]]) -> None:
    with st.expander(t("sources"), expanded=False):
        if not source_refs:
            st.info("No source metadata was available.")
            return
        for src in source_refs:
            pages = f"{src.get('page_start', '?')}–{src.get('page_end', '?')}"
            section = clean_section_title(src.get("section_title", ""))
            details = [
                f"Category: {src.get('source_label')}",
                f"Pages: {pages}",
                f"Relevance: {src.get('similarity_score', 0):.2f}",
            ]
            if section:
                details.insert(1, f"Section: {section}")
            st.markdown(
                f"**Source {src.get('source_number')}: {src.get('book', 'Course Material')}**  \n"
                + "  \n".join(details)
            )
            st.divider()


def _merge_page_ranges(page_ranges: List[Tuple[Any, Any]]) -> str:
    numeric: List[Tuple[int, int]] = []
    text_ranges: List[str] = []
    for start, end in page_ranges:
        try:
            a, b = int(start), int(end)
            if b < a:
                a, b = b, a
            numeric.append((a, b))
        except (TypeError, ValueError):
            label = f"{start}–{end}" if start or end else "Not available"
            if label not in text_ranges:
                text_ranges.append(label)

    numeric.sort()
    merged: List[List[int]] = []
    for start, end in numeric:
        if not merged or start > merged[-1][1] + 1:
            merged.append([start, end])
        else:
            merged[-1][1] = max(merged[-1][1], end)

    labels = [str(a) if a == b else f"{a}–{b}" for a, b in merged]
    labels.extend(text_ranges)
    return ", ".join(labels) or "Not available"


def display_image_content_references(source_refs: List[Dict[str, Any]]) -> None:
    """Display concise grouped references for an AI-generated image.

    These are not image copyrights or original image sources. They are the
    retrieved course passages used to prepare the image-generation prompt.
    """
    title = (
        "Course material used to prepare the image"
        if st.session_state.get("language") == "English"
        else "Material del curso utilizado para preparar la imagen"
    )
    with st.expander(title, expanded=False):
        if not source_refs:
            st.info("No course-reference metadata was available.")
            return

        grouped: Dict[Tuple[str, str], Dict[str, Any]] = {}
        for src in source_refs:
            key = (str(src.get("book", "Course Material")), str(src.get("source_label", "")))
            item = grouped.setdefault(key, {"ranges": [], "best": 0.0, "sections": []})
            item["ranges"].append((src.get("page_start", ""), src.get("page_end", "")))
            item["best"] = max(item["best"], float(src.get("similarity_score", 0) or 0))
            section = clean_section_title(src.get("section_title", ""))
            if section and section not in item["sections"]:
                item["sections"].append(section)

        for index, ((book, label), item) in enumerate(grouped.items(), 1):
            lines = [
                f"Category: {label}",
                f"Relevant pages: {_merge_page_ranges(item['ranges'])}",
                f"Best passage relevance: {item['best']:.2f}",
            ]
            if item["sections"]:
                lines.insert(1, f"Sections: {', '.join(item['sections'][:3])}")
            st.markdown(f"**{index}. {book}**  \n" + "  \n".join(lines))
            st.divider()

        st.caption(
            "The picture is newly generated by AI. These references support the medical content of the prompt; they are not the source of the image itself."
            if st.session_state.get("language") == "English"
            else "La imagen ha sido generada por IA. Estas referencias respaldan el contenido médico de las instrucciones; no son la fuente de la imagen en sí."
        )


def issue_options(language: str) -> List[str]:
    if language == "Spanish":
        return [
            "Respuesta incorrecta",
            "Redacción ambigua",
            "Más de una respuesta posible",
            "No está respaldada por el material del curso",
            "Explicación incorrecta",
            "Dificultad inadecuada",
            "Otro",
        ]
    return [
        "Incorrect answer",
        "Ambiguous wording",
        "More than one possible answer",
        "Not supported by course material",
        "Incorrect explanation",
        "Too easy or too difficult",
        "Other",
    ]


def professor_issue_options() -> List[str]:
    return [
        "Incorrect correct answer",
        "Ambiguous wording",
        "More than one defensible answer",
        "Weak or implausible distractors",
        "Incorrect or incomplete explanation",
        "Unsupported by Neurology Guiones",
        "Wrong difficulty level",
        "Duplicate or near-duplicate question",
        "Other",
    ]

# =====================================================
# STABLE AI IMAGE PANEL
# =====================================================
# Streamlit fragments rerun only this panel instead of rebuilding the whole app.
# The fallback keeps the app compatible if an older Streamlit version is used.
fragment = getattr(st, "fragment", lambda func: func)


@fragment
def render_ai_image_generator(topic: str, depth_level: str, language: str) -> None:
    """Render a stable AI-image form and persist the result across reruns."""
    state_defaults = {
        "ai_image_bytes": None,
        "ai_image_refs": [],
        "ai_image_mix": {},
        "ai_image_brief": "",
        "ai_image_error": "",
        "ai_image_caption": "",
    }
    for key, default in state_defaults.items():
        if key not in st.session_state:
            st.session_state[key] = default

    st.warning(
        "AI-generated educational illustration. It may contain inaccuracies and must not be used for diagnosis or exact anatomical measurement."
        if language == "English"
        else "Ilustración educativa generada por IA. Puede contener inexactitudes y no debe utilizarse para diagnóstico ni mediciones anatómicas exactas."
    )

    # A form batches input changes, so typing and changing the style do not
    # repeatedly rerun and redraw the interface.
    with st.form("ai_image_generation_form", clear_on_submit=False):
        visual_focus = st.text_input(
            "What should the image show?" if language == "English" else "¿Qué debe mostrar la imagen?",
            value=st.session_state.get("ai_visual_focus", topic),
            key="ai_visual_focus_form",
        )
        visual_style = st.selectbox(
            "Image format" if language == "English" else "Formato de imagen",
            ["Labelled medical illustration", "Mechanism diagram", "Clinical infographic"],
            key="ai_visual_style_form",
        )
        submitted = st.form_submit_button(
            "Generate educational image" if language == "English" else "Generar imagen educativa",
            type="primary",
            use_container_width=True,
        )

    status_slot = st.empty()

    if submitted:
        # Keep the previous image visible in Session State until the new image is ready.
        st.session_state["ai_image_error"] = ""
        status_slot.info(
            "Generating the source-grounded educational illustration…"
            if language == "English"
            else "Generando la ilustración educativa basada en las fuentes…"
        )

        image_bytes, refs, mix, visual_brief, image_error = generate_ai_medical_image(
            topic,
            depth_level,
            language,
            visual_focus,
            visual_style,
        )

        if image_error:
            st.session_state["ai_image_error"] = image_error
        elif image_bytes:
            st.session_state["ai_image_bytes"] = image_bytes
            st.session_state["ai_image_refs"] = refs
            st.session_state["ai_image_mix"] = mix
            st.session_state["ai_image_brief"] = visual_brief
            st.session_state["ai_image_caption"] = visual_focus

        status_slot.empty()

    if st.session_state.get("ai_image_error"):
        st.error(st.session_state["ai_image_error"])

    image_bytes = st.session_state.get("ai_image_bytes")
    if image_bytes:
        st.markdown("#### Generated educational image" if language == "English" else "#### Imagen educativa generada")

        # Constrain the image to a centered, fixed display width. This prevents
        # browser-width changes from continuously resizing the whole page.
        left, centre, right = st.columns([1, 6, 1])
        with centre:
            st.image(
                image_bytes,
                caption=(
                    f"AI-generated educational illustration: {st.session_state.get('ai_image_caption', topic)}"
                    if language == "English"
                    else f"Ilustración educativa generada por IA: {st.session_state.get('ai_image_caption', topic)}"
                ),
                width=700,
            )

        with st.expander(
            "Image-generation brief" if language == "English" else "Instrucciones usadas para generar la imagen",
            expanded=False,
        ):
            st.write(st.session_state.get("ai_image_brief", ""))

        display_source_mix(st.session_state.get("ai_image_mix", {}))
        display_image_content_references(st.session_state.get("ai_image_refs", []))

        if st.button(
            "Clear generated image" if language == "English" else "Borrar imagen generada",
            key="clear_ai_generated_image",
        ):
            for key, default in state_defaults.items():
                st.session_state[key] = default
            st.rerun(scope="fragment") if hasattr(st, "fragment") else st.rerun()


# =====================================================
# STUDENT MODE
# =====================================================
def student_mode() -> None:
    with st.sidebar:
        student_id = st.text_input(t("student_id"), value=st.session_state.get("student_id", ""))
        student_name = st.text_input(t("student_name"), value=st.session_state.get("student_name", ""))
        topic = st.selectbox(t("topic"), TOPICS)
        st.session_state["student_id"] = student_id
        st.session_state["student_name"] = student_name

    sid = (student_id or "Guest").strip() or "Guest"
    display_name = (student_name or sid).strip() or sid
    upsert_student(sid, display_name, st.session_state["language"])

    tab_chat, tab_quiz, tab_report = st.tabs([t("chat"), t("quiz"), t("report")])

    with tab_chat:
        c1, c2 = st.columns(2)
        with c1:
            depth_level = st.selectbox(t("depth"), DEPTH_LEVELS, index=1)
        with c2:
            if st.session_state["language"] == "English":
                activities = [
                    "Free question",
                    "Explanation for selected topic",
                    "Flashcards for selected topic",
                    "Case study for selected topic",
                    "Structured outline",
                    "Concept map",
                    "Clinical decision pathway",
                    "AI-generated educational image",
                ]
            else:
                activities = [
                    "Pregunta libre",
                    "Explicación del tema seleccionado",
                    "Tarjetas de estudio",
                    "Caso clínico",
                    "Esquema estructurado",
                    "Mapa conceptual",
                    "Ruta de decisión clínica",
                    "Imagen educativa generada por IA",
                ]
            tutor_activity = st.selectbox(t("activity"), activities)

        free_question = tutor_activity in ["Free question", "Pregunta libre"]
        q = st.text_area(t("ask_question"), height=120) if free_question else ""

        if tutor_activity in ["AI-generated educational image", "Imagen educativa generada por IA"]:
            render_ai_image_generator(
                topic,
                depth_level,
                st.session_state["language"],
            )

        elif tutor_activity in ["Concept map", "Mapa conceptual", "Clinical decision pathway", "Ruta de decisión clínica"]:
            visual_type = "concept map" if tutor_activity in ["Concept map", "Mapa conceptual"] else "clinical decision pathway"
            if st.button(t("send"), key="generate_visual"):
                with st.spinner("Generating source-grounded diagram..."):
                    dot, refs, mix, visual_error = generate_dot_visual(
                        topic, visual_type, depth_level, st.session_state["language"]
                    )
                if visual_error:
                    st.error(visual_error)
                else:
                    st.graphviz_chart(dot, use_container_width=True)
                    display_source_mix(mix)
                    display_sources(refs)

        else:
            if st.button(t("send"), key="ask_btn"):
                if tutor_activity in ["Explanation for selected topic", "Explicación del tema seleccionado"]:
                    q_to_send = f"Explain the selected topic for a medical student: {topic}"
                elif tutor_activity in ["Flashcards for selected topic", "Tarjetas de estudio"]:
                    q_to_send = f"Create 8 flashcards with question and answer for: {topic}"
                elif tutor_activity in ["Case study for selected topic", "Caso clínico"]:
                    q_to_send = f"Create one clinical case study with questions, answers and explanations for: {topic}"
                elif tutor_activity in ["Structured outline", "Esquema estructurado"]:
                    q_to_send = f"Create a structured study outline for: {topic}"
                else:
                    q_to_send = q.strip()

                if not q_to_send:
                    st.warning("Please write a question." if st.session_state["language"] == "English" else "Por favor escribe una pregunta.")
                else:
                    with st.spinner("BrainChat is preparing the answer..."):
                        answer, color, similarity, refs, mix, answer_error = answer_tutor_question(
                            q_to_send, topic, st.session_state["language"], depth_level
                        )
                    save_chat_log(
                        sid, st.session_state["language"], topic, q_to_send,
                        answer, color, similarity, depth_level, refs, mix
                    )
                    confidence_badge(color)
                    st.caption(f"Similarity: {similarity:.2f} | Level: {depth_level}")
                    if answer_error:
                        st.error(answer)
                    else:
                        st.markdown(answer)
                    display_source_mix(mix)
                    display_sources(refs)

    with tab_quiz:
        c1, c2 = st.columns(2)
        with c1:
            difficulty = st.selectbox(t("difficulty"), QUIZ_DIFFICULTIES)
        with c2:
            n_questions = st.selectbox(t("num_questions"), QUESTION_COUNTS, index=1)

        if st.button(t("start_quiz"), key="gen_quiz"):
            with st.spinner("Generating course-grounded MCQ quiz..."):
                quiz, warning, refs, mix = generate_mcqs(
                    topic, difficulty, n_questions,
                    st.session_state["language"], sid
                )
            st.session_state["current_quiz"] = quiz
            st.session_state["quiz_topic"] = topic
            st.session_state["quiz_difficulty"] = difficulty
            st.session_state["quiz_source_refs"] = refs
            st.session_state["quiz_source_mix"] = mix
            st.session_state["quiz_submitted"] = False
            if warning:
                st.warning(warning)

        quiz = st.session_state.get("current_quiz", [])
        if quiz:
            answers: Dict[str, str] = {}
            for i, item in enumerate(quiz, 1):
                status_label = "Professor-approved" if item.get("status") == "approved" else "Generated"
                st.markdown(f"### Q{i}. {item['question']}")
                st.caption(status_label)
                options_display = [format_option(opt) for opt in item.get("options", [])]
                choice = st.radio(
                    "Select answer" if st.session_state["language"] == "English" else "Selecciona la respuesta",
                    options_display,
                    key=f"quiz_{item.get('question_id', i)}",
                )
                answers[str(i)] = choice.strip()[0].upper() if choice else ""

            if st.button(t("submit_quiz"), key="submit_quiz"):
                score = 0
                weak: List[str] = []
                st.session_state["quiz_submitted"] = True
                st.session_state["submitted_answers"] = answers
                for i, item in enumerate(quiz, 1):
                    correct = item.get("correct_option", "A").upper()
                    selected = answers.get(str(i), "")
                    if selected == correct:
                        score += 1
                    else:
                        weak.append(item.get("subtopic", topic))

                percent = score / max(len(quiz), 1) * 100
                color = "green" if percent >= 70 else "orange" if percent >= 45 else "red"
                badges = badges_for_student(sid)
                save_quiz_attempt(
                    sid, display_name, st.session_state["language"],
                    st.session_state.get("quiz_topic", topic),
                    st.session_state.get("quiz_difficulty", difficulty),
                    score, len(quiz), color, sorted(set(weak)), badges,
                    quiz, answers,
                    st.session_state.get("quiz_source_refs", []),
                    st.session_state.get("quiz_source_mix", {}),
                )
                st.success(f"{t('score')}: {score}/{len(quiz)} ({percent:.1f}%). {t('saved')}")

            if st.session_state.get("quiz_submitted"):
                submitted_answers = st.session_state.get("submitted_answers", {})
                st.markdown("## Results" if st.session_state["language"] == "English" else "## Resultados")
                for i, item in enumerate(quiz, 1):
                    correct = item.get("correct_option", "A").upper()
                    selected = submitted_answers.get(str(i), "")
                    ok = selected == correct
                    st.markdown(f"**Q{i}: {'✅' if ok else '❌'} Selected: {selected} | Correct: {correct}**")
                    st.write(item.get("explanation", ""))

                    with st.expander(f"Report a problem with Question {i}"):
                        issue_type = st.selectbox(
                            "Problem type",
                            issue_options(st.session_state["language"]),
                            key=f"issue_{item.get('question_id', i)}",
                        )
                        comment = st.text_area(
                            "Explain the problem",
                            key=f"comment_{item.get('question_id', i)}",
                        )
                        if st.button("Send to professor", key=f"flag_{item.get('question_id', i)}"):
                            created = create_question_review(item, "student", sid, issue_type, comment)
                            if created:
                                st.success("Question sent for professor review.")
                            else:
                                st.info("This question is already waiting for professor review.")

                display_source_mix(st.session_state.get("quiz_source_mix", {}))
                display_sources(st.session_state.get("quiz_source_refs", []))

    with tab_report:
        df = load_attempts_df()
        sdf = df[df["student_id"] == sid] if not df.empty else pd.DataFrame()
        if sdf.empty:
            st.info(t("no_data"))
        else:
            c1, c2, c3 = st.columns(3)
            c1.metric("Average score", f"{sdf['percent'].mean():.1f}%")
            c2.metric("Attempts", len(sdf))
            c3.metric("Badges", len(badges_for_student(sid)))
            st.dataframe(
                sdf[["created_at", "topic", "difficulty", "score", "total", "percent", "confidence_color"]],
                use_container_width=True,
            )
            fig = px.line(
                sdf.sort_values("created_at"), x="created_at", y="percent",
                color="topic", markers=True, title="Progress over time",
            )
            st.plotly_chart(fig, use_container_width=True)

        report_html = html_report_student(sid, display_name, st.session_state["language"])
        st.download_button(
            t("download_html"), data=report_html,
            file_name=f"brainchat_report_{sid}.html", mime="text/html",
        )

# =====================================================
# TEACHER MODE
# =====================================================
def render_question_review_tab() -> None:
    st.subheader("Human-in-the-loop Question Improvement")
    st.caption(
        "A professor can flag any generated question, correct and approve it, or reject it. "
        "Approved corrections become trusted examples. Rejected questions and professor rules are used to block similar future errors."
    )

    reviewer = st.text_input("Reviewer name", value=st.session_state.get("reviewer_name", "Professor"))
    st.session_state["reviewer_name"] = reviewer

    pending = load_pending_reviews()
    conn = get_conn()
    approved_count = conn.execute("SELECT COUNT(*) FROM approved_questions WHERE active=1").fetchone()[0]
    rejected_count = conn.execute("SELECT COUNT(*) FROM rejected_question_patterns WHERE active=1").fetchone()[0]
    rule_count = conn.execute("SELECT COUNT(*) FROM feedback_rules WHERE active=1").fetchone()[0]
    conn.close()

    c1, c2, c3, c4 = st.columns(4)
    c1.metric("Pending reviews", len(pending))
    c2.metric("Approved questions", approved_count)
    c3.metric("Rejected patterns", rejected_count)
    c4.metric("Professor rules", rule_count)

    pending_tab, history_tab, memory_tab = st.tabs([
        "Pending Corrections", "Generated Question History", "Learning Memory"
    ])

    with pending_tab:
        if pending.empty:
            st.success("No pending question reviews.")
        else:
            review_labels = [
                f"#{int(row.review_id)} | {row.topic} | {str(row.question)[:75]}"
                for row in pending.itertuples()
            ]
            selected_label = st.selectbox("Select review", review_labels, key="pending_review_select")
            selected_index = review_labels.index(selected_label)
            row = pending.iloc[selected_index]

            options = json.loads(row["options_json"] or "[]")
            source_refs = json.loads(row["source_refs_json"] or "[]")

            st.markdown(f"### Original question\n{row['question']}")
            st.write("Original options:")
            for option in options:
                st.write(option)
            st.markdown(f"**Current correct answer:** {row['correct_option']}")
            st.markdown(f"**Current explanation:** {row['explanation']}")
            st.warning(
                f"Reported by {row['reporter_type']}: {row['issue_type']} — "
                f"{row['reporter_comment'] or 'No comment provided'}"
            )
            display_sources(source_refs)

            st.markdown("### Professor correction")
            corrected_question = st.text_area(
                "Corrected question", value=row["question"], key=f"corrected_q_{row['review_id']}"
            )
            corrected_options = []
            for i in range(5):
                default = options[i] if i < len(options) else f"{chr(65+i)}. "
                corrected_options.append(
                    st.text_input(
                        f"Option {chr(65+i)}", value=default,
                        key=f"corrected_opt_{row['review_id']}_{i}",
                    )
                )

            answer_index = "ABCDE".find(str(row["correct_option"]).upper())
            corrected_answer = st.selectbox(
                "Correct answer", list("ABCDE"), index=max(answer_index, 0),
                key=f"corrected_answer_{row['review_id']}",
            )
            corrected_explanation = st.text_area(
                "Corrected explanation", value=row["explanation"],
                key=f"corrected_exp_{row['review_id']}",
            )
            professor_comment = st.text_area(
                "Professor rule or reason",
                placeholder="Example: Avoid absolute wording such as 'always'; treatment depends on seizure type and contraindications.",
                key=f"prof_comment_{row['review_id']}",
            )
            st.caption(
                "This comment is saved as a reusable rule for future question generation. "
                "Use a general instruction, not only a description of this single question."
            )

            b1, b2 = st.columns(2)
            with b1:
                if st.button("Correct and approve", type="primary", key=f"approve_{row['review_id']}"):
                    if not corrected_question.strip() or any(not x.strip() for x in corrected_options):
                        st.error("Question and all five options are required.")
                    else:
                        approve_review(
                            int(row["review_id"]), row["question_id"], reviewer or "Professor",
                            professor_comment, corrected_question, corrected_options,
                            corrected_answer, corrected_explanation,
                        )
                        st.success("Correction approved. Future quizzes will use it as a trusted example.")
                        st.rerun()
            with b2:
                if st.button("Reject and block pattern", key=f"reject_{row['review_id']}"):
                    reason = professor_comment or row["issue_type"] or "Rejected by professor"
                    reject_review(
                        int(row["review_id"]), row["question_id"],
                        reviewer or "Professor", reason,
                    )
                    st.success("Question rejected. Its question pattern and professor rule are now blocked in future generation.")
                    st.rerun()

    with history_tab:
        st.markdown("### Professor direct flagging")
        st.caption(
            "Use this screen to flag a poorly formulated question even when no student has reported it."
        )
        history = load_generated_questions_df()
        if history.empty:
            st.info("No generated questions have been stored yet.")
        else:
            f1, f2 = st.columns(2)
            with f1:
                topic_filter = st.selectbox(
                    "Filter topic", ["All"] + TOPICS, key="history_topic_filter"
                )
            with f2:
                status_values = sorted(history["status"].fillna("unreviewed").unique().tolist())
                status_filter = st.selectbox(
                    "Filter status", ["All"] + status_values, key="history_status_filter"
                )

            filtered = history.copy()
            if topic_filter != "All":
                filtered = filtered[filtered["topic"] == topic_filter]
            if status_filter != "All":
                filtered = filtered[filtered["status"] == status_filter]

            if filtered.empty:
                st.info("No questions match the selected filters.")
            else:
                labels = [
                    f"{r.created_at} | {r.topic} | {r.status} | {str(r.question)[:80]}"
                    for r in filtered.itertuples()
                ]
                chosen = st.selectbox("Select generated question", labels, key="history_question_select")
                row = filtered.iloc[labels.index(chosen)]
                options = json.loads(row["options_json"] or "[]")
                refs = json.loads(row["source_refs_json"] or "[]")

                st.markdown(f"### {row['question']}")
                for option in options:
                    st.write(option)
                st.markdown(f"**Correct answer:** {row['correct_option']}")
                st.markdown(f"**Explanation:** {row['explanation']}")
                st.caption(
                    f"Topic: {row['topic']} | Difficulty: {row['difficulty']} | Status: {row['status']}"
                )
                display_sources(refs)

                issue = st.selectbox(
                    "Why is this question poor?", professor_issue_options(),
                    key="professor_direct_issue",
                )
                comment = st.text_area(
                    "Initial correction note",
                    placeholder="Describe the error and the quality rule that future questions should follow.",
                    key="professor_direct_comment",
                )
                if st.button("Flag for correction", type="primary", key="professor_direct_flag"):
                    item = {
                        "question_id": row["question_id"],
                        "topic": row["topic"],
                        "difficulty": row["difficulty"],
                        "language": row["language"],
                        "question": row["question"],
                        "options": options,
                        "correct_option": row["correct_option"],
                        "explanation": row["explanation"],
                        "subtopic": row["subtopic"],
                        "source_refs": refs,
                    }
                    created = create_question_review(
                        item, "professor", reviewer or "Professor", issue, comment
                    )
                    if created:
                        st.success("Question added to Pending Corrections.")
                    else:
                        st.info("This question is already waiting for review.")

    with memory_tab:
        st.markdown("### Supervised learning memory")
        st.info(
            "This is immediate human-in-the-loop learning through retrieval and filtering, not automatic model fine-tuning. "
            "Approved questions are reused directly and as examples; rejected questions are blocked; professor comments become generation rules."
        )
        conn = get_conn()
        approved_df = pd.read_sql_query(
            "SELECT id, question_id, topic, difficulty, question, correct_option, approved_by, approved_at, active FROM approved_questions ORDER BY approved_at DESC",
            conn,
        )
        rejected_df = pd.read_sql_query(
            "SELECT id, topic, question, reason, rejected_by, created_at, active FROM rejected_question_patterns ORDER BY created_at DESC",
            conn,
        )
        rules_df = pd.read_sql_query(
            "SELECT id, topic, rule_text, decision_type, created_by, created_at, active FROM feedback_rules ORDER BY created_at DESC",
            conn,
        )
        conn.close()

        st.markdown("#### Approved question bank")
        st.dataframe(approved_df, use_container_width=True)
        st.markdown("#### Rejected question memory")
        st.dataframe(rejected_df, use_container_width=True)
        st.markdown("#### Professor feedback rules")
        st.dataframe(rules_df, use_container_width=True)


def teacher_mode() -> None:
    pwd = st.text_input(t("teacher_password"), type="password")
    if not st.button(t("login")) and not st.session_state.get("teacher_ok"):
        return
    if pwd == TEACHER_PASSWORD or st.session_state.get("teacher_ok"):
        st.session_state["teacher_ok"] = True
    else:
        st.error("Incorrect password")
        return

    analytics_tab, review_tab, content_tab, reports_tab = st.tabs([
        "Analytics", "Question Review", "Content & Sources", "Reports"
    ])

    with analytics_tab:
        df = load_attempts_df()
        chat_df = load_chat_df()
        if df.empty:
            st.warning(t("no_data"))
        else:
            c1, c2, c3, c4 = st.columns(4)
            c1.metric("Students", df["student_id"].nunique())
            c2.metric("Quiz attempts", len(df))
            c3.metric("Average score", f"{df['percent'].mean():.1f}%")
            c4.metric("Low confidence", int((df["confidence_color"] == "red").sum()))

            topic_summary = df.groupby("topic").agg(
                attempts=("id", "count"), avg_score=("percent", "mean")
            ).reset_index()
            fig1 = px.bar(
                topic_summary, x="topic", y="avg_score", hover_data=["attempts"],
                title="Average score by topic",
            )
            st.plotly_chart(fig1, use_container_width=True)

            student_summary = df.groupby(["student_id", "student_name"]).agg(
                attempts=("id", "count"), avg_score=("percent", "mean")
            ).reset_index()
            st.dataframe(student_summary, use_container_width=True)

            selected_student = st.selectbox("Select student", sorted(df["student_id"].unique()))
            st.dataframe(
                df[df["student_id"] == selected_student][[
                    "created_at", "student_name", "topic", "difficulty",
                    "score", "total", "percent", "weak_areas", "badges",
                ]],
                use_container_width=True,
            )

        with st.expander("Tutor chat logs"):
            if chat_df.empty:
                st.info("No chat logs yet.")
            else:
                st.dataframe(
                    chat_df[[
                        "created_at", "student_id", "topic", "depth_level",
                        "question", "confidence_color", "similarity",
                    ]],
                    use_container_width=True,
                )

    with review_tab:
        render_question_review_tab()

    with content_tab:
        st.subheader("Source Priority and Visual Content")
        st.markdown(
            """
**Retrieval priority**
1. Official Neurology Guiones
2. Other course material
3. Supplementary external sources

The source boost is applied only when a passage meets a minimum semantic-relevance threshold. This prevents an irrelevant official passage from replacing a relevant passage.
"""
        )
        st.code(json.dumps(SOURCE_PRIORITY, indent=2), language="json")

        st.markdown("### AI-generated visual content")
        c1, c2 = st.columns(2)
        c1.metric("AI image generation", "Enabled" if ENABLE_AI_IMAGES else "Disabled")
        c2.metric("Image model", OPENAI_IMAGE_MODEL)
        st.caption(
            "No medical image folder or manifest is required. Each AI image request is prepared from retrieved course passages and displayed with its supporting-source composition."
        )

        st.markdown("### Question-learning protocol")
        st.markdown(
            """
1. A student or professor flags a question.
2. The professor corrects and approves it, or rejects it.
3. Corrected questions enter the approved bank and are reused directly and as examples.
4. Rejected questions are blocked through exact-hash, word-overlap and near-text similarity checks.
5. Professor comments become reusable generation rules.
6. The original formulation is blocked whenever the professor replaces or materially corrects it.
"""
        )

        st.markdown("### RAG build status")
        missing = [p for p in [CHUNKS_PATH, TOKENS_PATH, EMBED_PATH, CONFIG_PATH] if not os.path.exists(p)]
        if missing:
            st.error("Missing: " + ", ".join(missing))
        else:
            st.success("All RAG build files are available.")

            chunks, _, _, _, inventory_error = load_rag_resources()
            if not inventory_error and chunks:
                inventory_rows = []
                seen_inventory = set()
                for record in chunks:
                    display_name, source_type, raw_name = resolve_source_metadata(record)
                    key = (raw_name, display_name, source_type)
                    if key in seen_inventory:
                        continue
                    seen_inventory.add(key)
                    inventory_rows.append({
                        "Raw source metadata": raw_name,
                        "Displayed name": display_name,
                        "Category": SOURCE_LABELS.get(source_type, source_type),
                    })
                with st.expander("Source-name and category preview", expanded=False):
                    st.dataframe(pd.DataFrame(inventory_rows), use_container_width=True)
                    st.caption(
                        "Use src/source_aliases.json when a raw filename is generic or when different page ranges in a merged PDF belong to different source categories."
                    )

        st.markdown("### Optional source aliases")
        if os.path.exists(SOURCE_ALIASES_FILE):
            st.success("src/source_aliases.json is available.")
        else:
            st.info(
                "No source_aliases.json file is present. The app will infer source names and categories from the metadata stored in chunks.pkl."
            )

    with reports_tab:
        st.download_button(
            "Download teacher HTML report",
            data=html_report_teacher(),
            file_name="brainchat_teacher_report.html",
            mime="text/html",
        )

# =====================================================
# MAIN
# =====================================================
def main() -> None:
    init_db()
    if "language" not in st.session_state:
        st.session_state["language"] = "English"

    with st.sidebar:
        st.session_state["language"] = st.radio(
            "Interface language / Idioma", ["English", "Spanish"], horizontal=True
        )
        mode = st.radio(t("mode"), [t("student_mode"), t("teacher_mode")])

    render_header()

    with st.expander("How evidence and confidence are shown", expanded=False):
        st.markdown(
            """
- **Green:** the retrieved course material strongly supports the question.
- **Orange:** support is partial and the answer should be revised carefully.
- **Red:** support is weak or insufficient.
- **Retrieved evidence composition:** estimated share of retrieved supporting passages from official Guiones, other course material and supplementary sources.
- The percentage is not presented as an exact measure of generated words.
"""
        )

    if mode == t("student_mode"):
        student_mode()
    else:
        teacher_mode()


if __name__ == "__main__":
    main()