app.py GOOOD

import os
import re
import json
import pickle
import base64
import mimetypes
from datetime import datetime

import numpy as np
import gradio as gr
from openai import OpenAI
from rank_bm25 import BM25Okapi
from sentence_transformers import SentenceTransformer


# =====================================================
# CONFIG
# =====================================================
BUILD_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")

LOGO_FILE = "logo.png"
OPENAI_MODEL = os.getenv("OPENAI_MODEL", "gpt-4o-mini")

BM25 = None
CHUNKS = None
EMBEDDINGS = None
EMBED_MODEL = None
CLIENT = None

ANALYTICS_LOG = []


# =====================================================
# LOADERS
# =====================================================
def tokenize(text: str):
    return re.findall(r"\w+", text.lower(), flags=re.UNICODE)


def ensure_loaded():
    global BM25, CHUNKS, EMBEDDINGS, EMBED_MODEL, CLIENT

    if CHUNKS is None:
        missing = []
        for p in [CHUNKS_PATH, TOKENS_PATH, EMBED_PATH, CONFIG_PATH]:
            if not os.path.exists(p):
                missing.append(p)

        if missing:
            raise FileNotFoundError("Missing build files:\n" + "\n".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"])

    if CLIENT is None:
        api_key = os.getenv("OPENAI_API_KEY")
        if not api_key:
            raise ValueError("OPENAI_API_KEY is missing in Hugging Face Space Secrets.")
        CLIENT = OpenAI(api_key=api_key)


# =====================================================
# SOURCE CLEANING AND PRIORITY
# =====================================================
def clean_source_name(book_name: str) -> str:
    name = (book_name or "").strip()

    if "ilovepdf" in name.lower() or "merged" in name.lower():
        return "Professor Handouts"

    if name.lower().endswith(".pdf"):
        name = name[:-4]

    return name or "Professor Handouts"


def source_priority_label(book_name: str) -> str:
    return "Primary source" if clean_source_name(book_name) == "Professor Handouts" else "Supporting textbook"


def prioritize_professor_handouts(records):
    return sorted(
        records,
        key=lambda r: (
            0 if clean_source_name(r.get("book", "")) == "Professor Handouts" else 1,
            -float(r.get("final_score", r.get("similarity_score", 0)))
        )
    )


# =====================================================
# GENERAL CHAT
# =====================================================
def is_general_chat(text: str) -> bool:
    t = text.lower().strip()

    general_phrases = [
        "hi", "hello", "hola", "hey",
        "good morning", "good afternoon", "good evening",
        "thanks", "thank you", "gracias",
        "ok", "okay", "who are you", "what can you do", "help"
    ]

    return t in general_phrases


def general_chat_reply(text: str, language_mode: str) -> str:
    t = text.lower().strip()

    if language_mode == "English":
        return (
            "Hello! I am BrainChat, your AI tutor for Neurology and PMQSN. "
            "You can ask me to explain topics, create short notes, generate flashcards, "
            "or test you with quiz questions. I first use Professor Handouts, "
            "and then supporting textbooks if needed."
        )

    if language_mode == "Spanish":
        return (
            "¡Hola! Soy BrainChat, tu tutor de IA para Neurología y PMQSN. "
            "Puedes pedirme explicaciones, apuntes breves, flashcards o preguntas tipo quiz. "
            "Primero usaré los apuntes del profesor y, si es necesario, otros libros de apoyo."
        )

    if t in ["hola", "gracias"]:
        return (
            "¡Hola! Soy BrainChat, tu tutor de IA para Neurología y PMQSN. "
            "Primero uso los apuntes del profesor y después otros libros de apoyo si es necesario."
        )

    return (
        "Hello! I am BrainChat, your AI tutor for Neurology and PMQSN. "
        "I first use Professor Handouts and then supporting textbooks if needed."
    )


# =====================================================
# RETRIEVAL WITH PROFESSOR HANDOUT BOOST
# =====================================================
def search_hybrid(query: str, shortlist_k: int = 30, final_k: int = 5):
    ensure_loaded()

    q_tokens = tokenize(query)
    bm25_scores = BM25.get_scores(q_tokens)

    shortlist_idx = np.argsort(bm25_scores)[::-1][:shortlist_k]
    shortlist_emb = EMBEDDINGS[shortlist_idx]

    qvec = EMBED_MODEL.encode([query], normalize_embeddings=True).astype("float32")[0]
    dense_scores = shortlist_emb @ qvec

    results = []

    for idx, score in zip(shortlist_idx, dense_scores):
        record = CHUNKS[int(idx)].copy()
        clean_book = clean_source_name(record.get("book", ""))

        priority_boost = 0.15 if clean_book == "Professor Handouts" else 0.0
        final_score = float(score) + priority_boost

        record["similarity_score"] = float(score)
        record["final_score"] = final_score
        record["source_priority"] = (
            "Professor Handouts"
            if clean_book == "Professor Handouts"
            else "Supporting textbook"
        )

        results.append(record)

    results = sorted(results, key=lambda r: r["final_score"], reverse=True)

    return prioritize_professor_handouts(results[:final_k])


def build_context(records):
    blocks = []

    records = prioritize_professor_handouts(records)

    for i, r in enumerate(records, start=1):
        clean_book = clean_source_name(r.get("book", ""))

        blocks.append(
            f"""[Source {i}]
Book: {clean_book}
Source priority: {source_priority_label(clean_book)}
Section: {r.get('section_title','')}
Pages: {r.get('page_start','')}-{r.get('page_end','')}
Similarity Score: {r.get('similarity_score', 0):.3f}
Final Score: {r.get('final_score', r.get('similarity_score', 0)):.3f}
Text:
{r.get('text','')}"""
        )

    return "\n\n".join(blocks)


def make_sources(records):
    seen = set()
    lines = []

    records = prioritize_professor_handouts(records)

    for r in records:
        clean_book = clean_source_name(r.get("book", ""))

        key = (
            clean_book,
            r.get("section_title"),
            r.get("page_start"),
            r.get("page_end"),
        )

        if key in seen:
            continue

        seen.add(key)

        section = r.get("section_title", "Course Material")
        page_start = r.get("page_start", "")
        page_end = r.get("page_end", "")
        score = r.get("final_score", r.get("similarity_score", 0))

        if page_start and page_end and page_start != page_end:
            page_text = f"pages {page_start}-{page_end}"
        elif page_start:
            page_text = f"page {page_start}"
        else:
            page_text = "page not specified"

        source_type = source_priority_label(clean_book)

        lines.append(
            f"• {clean_book} ({source_type}) | {section} | {page_text} | relevance: {score:.2f}"
        )

    return "\n".join(lines)


# =====================================================
# CONFIDENCE LOGIC
# =====================================================
def is_not_found_answer(answer: str) -> bool:
    a = (answer or "").lower().strip()

    return (
        "not found in the course material" in a
        or "no encontrado en el material del curso" in a
        or "no se encontró información" in a
        or a == "no encontrado"
    )


def compute_confidence(records, answer: str):
    if is_not_found_answer(answer):
        return {
            "level": "red",
            "label": "Not found",
            "score": 0.0,
        }

    if not records:
        return {
            "level": "red",
            "label": "Not found",
            "score": 0.0,
        }

    scores = [float(r.get("final_score", r.get("similarity_score", 0))) for r in records]
    raw_scores = [float(r.get("similarity_score", 0)) for r in records]

    top_score = max(scores)
    top_raw = max(raw_scores)

    professor_found = any(
        clean_source_name(r.get("book", "")) == "Professor Handouts"
        for r in records[:3]
    )

    if professor_found and top_score >= 0.48:
        return {
            "level": "green",
            "label": "High confidence",
            "score": top_raw,
        }

    if top_score >= 0.52:
        return {
            "level": "green",
            "label": "High confidence",
            "score": top_raw,
        }

    if top_score >= 0.35:
        return {
            "level": "orange",
            "label": "Medium confidence",
            "score": top_raw,
        }

    return {
        "level": "red",
        "label": "Low confidence",
        "score": top_raw,
    }


def confidence_html(conf):
    color_map = {
        "green": "#16a34a",
        "orange": "#f97316",
        "red": "#dc2626",
    }

    color = color_map.get(conf["level"], "#999999")

    return f"""
    <div class="bc-confidence">
        <span class="bc-dot" style="background:{color};"></span>
        <span><strong>{conf['label']}</strong> — similarity score: {conf['score']:.2f}</span>
    </div>
    """


# =====================================================
# ANALYTICS DASHBOARD
# =====================================================
def log_event(event_type, mode, language, confidence_level, similarity, query):
    ANALYTICS_LOG.append({
        "time": datetime.now().strftime("%Y-%m-%d %H:%M:%S"),
        "event": event_type,
        "mode": mode,
        "language": language,
        "confidence": confidence_level,
        "similarity": round(float(similarity), 3),
        "query": query[:120],
    })


def render_dashboard():
    total = len(ANALYTICS_LOG)

    if total == 0:
        return """
        <div class="bc-dashboard">
            <div class="bc-dashboard-grid">
                <div>
                    <h3>Progress Analytics Dashboard</h3>
                    <p>No interactions recorded yet.</p>
                </div>
                <div class="bc-dashboard-help">
                    <h4>What this dashboard shows</h4>
                    <p>This dashboard summarizes how students are using BrainChat.</p>
                    <p><strong>Total interactions:</strong> number of questions or quiz actions.</p>
                    <p><strong>High confidence:</strong> answers strongly supported by course material.</p>
                    <p><strong>Medium confidence:</strong> answers with partial support.</p>
                    <p><strong>Low / Not found:</strong> questions not clearly supported by the material.</p>
                    <p><strong>Average similarity:</strong> how closely the retrieved material matches the question.</p>
                </div>
            </div>
        </div>
        """

    green = sum(1 for x in ANALYTICS_LOG if x["confidence"] == "green")
    orange = sum(1 for x in ANALYTICS_LOG if x["confidence"] == "orange")
    red = sum(1 for x in ANALYTICS_LOG if x["confidence"] == "red")
    quizzes = sum(1 for x in ANALYTICS_LOG if x["event"] in ["quiz_generated", "quiz_evaluated"])
    avg_sim = sum(x["similarity"] for x in ANALYTICS_LOG) / total

    recent_rows = ""

    for item in ANALYTICS_LOG[-8:][::-1]:
        recent_rows += f"""
        <tr>
            <td>{item['time']}</td>
            <td>{item['event']}</td>
            <td>{item['mode']}</td>
            <td><span class="bc-pill bc-{item['confidence']}">{item['confidence']}</span></td>
            <td>{item['similarity']}</td>
            <td>{item['query']}</td>
        </tr>
        """

    return f"""
    <div class="bc-dashboard">
        <div class="bc-dashboard-grid">
            <div>
                <h3>Progress Analytics Dashboard</h3>

                <div class="bc-metrics">
                    <div class="bc-card total"><strong>{total}</strong><br>Total interactions</div>
                    <div class="bc-card green"><strong>{green}</strong><br>High confidence</div>
                    <div class="bc-card orange"><strong>{orange}</strong><br>Medium confidence</div>
                    <div class="bc-card red"><strong>{red}</strong><br>Low / Not found</div>
                    <div class="bc-card quiz"><strong>{quizzes}</strong><br>Quiz actions</div>
                    <div class="bc-card avg"><strong>{avg_sim:.2f}</strong><br>Avg similarity</div>
                </div>
            </div>

            <div class="bc-dashboard-help">
                <h4>What this dashboard shows</h4>
                <p>This dashboard helps teachers monitor BrainChat usage and answer quality.</p>
                <p><strong>🟢 High confidence:</strong> retrieved material strongly supports the answer.</p>
                <p><strong>🟠 Medium confidence:</strong> answer may need checking with handouts.</p>
                <p><strong>🔴 Low / Not found:</strong> material is weak or not available.</p>
                <p><strong>Avg similarity:</strong> higher value means a better match between the question and course material.</p>
            </div>
        </div>

        <h4>Recent activity</h4>
        <table class="bc-table">
            <tr>
                <th>Time</th>
                <th>Event</th>
                <th>Mode</th>
                <th>Confidence</th>
                <th>Similarity</th>
                <th>Query</th>
            </tr>
            {recent_rows}
        </table>
    </div>
    """


def refresh_dashboard():
    return render_dashboard()


def clear_analytics():
    ANALYTICS_LOG.clear()
    return render_dashboard()


# =====================================================
# PROMPTS
# =====================================================
def language_instruction(language_mode: str) -> str:
    if language_mode == "English":
        return "Answer only in English."

    if language_mode == "Spanish":
        return "Answer only in Spanish."

    if language_mode == "Bilingual":
        return "Answer first in English, then provide a Spanish version under the heading 'Español:'."

    return "If the user's message is in Spanish, answer in Spanish; otherwise answer in English."


def choose_quiz_count(user_text: str, selector: str) -> int:
    if selector in {"3", "5", "7"}:
        return int(selector)

    t = user_text.lower()

    if any(k in t for k in ["mock test", "final exam", "exam practice", "full test"]):
        return 7

    if any(k in t for k in ["detailed", "revision", "comprehensive", "study"]):
        return 5

    return 3


def build_tutor_prompt(mode: str, language_mode: str, question: str, context: str) -> str:
    styles = {
        "Explain": """
Explain clearly like a friendly clinical tutor.
Use simple language.
Give the concept first, then key clinical points.
If useful, include one common mistake to avoid.
""",
        "Detailed": """
Give a detailed explanation with clinical relevance.
Structure the answer using clear headings.
Only include details supported by the context.
""",
        "Short Notes": """
Write concise revision notes using short bullet points.
Focus on exam-useful points from the professor handouts.
""",
        "Flashcards": """
Create 6 flashcards in Q/A format using only the context.
Keep them useful for exam revision.
""",
        "Case-Based": """
Create a short clinical case scenario.
Then guide the student using clinical reasoning.
Use the Socratic method where possible.
Do not simply give the answer immediately if reasoning is expected.
""",
    }

    return f"""
You are BrainChat, an interactive neurology tutor for PMQSN.

Core rules:
- Use ONLY the provided context.
- Always prioritize Professor Handouts first.
- Use supporting textbooks only when Professor Handouts are insufficient.
- Clearly keep Professor Handouts as the primary course source.
- If the answer is not supported by the context, say exactly:
  Not found in the course material.
- Do not invent facts outside the context.
- Do not invent references.
- {language_instruction(language_mode)}

Teaching behavior:
- Act as a Socratic clinical tutor.
- Prefer guiding the student with reasoning rather than only giving direct answers.
- Keep the answer clear, structured, and useful for medical students.
- If the question asks for treatment, diagnosis, definition, or comparison, focus directly on that requested point.

Teaching style:
{styles.get(mode, "Explain clearly like a friendly clinical tutor.")}

Context:
{context}

Student question:
{question}
""".strip()


def build_quiz_generation_prompt(language_mode: str, topic: str, context: str, n_questions: int) -> str:
    return f"""
You are BrainChat, an interactive neurology tutor.

Rules:
- Use ONLY the provided context.
- Always prioritize Professor Handouts first.
- Use supporting textbooks only when needed.
- Create exactly {n_questions} quiz questions.
- Questions should support autonomous study.
- Keep questions short and clear.
- Include a short answer key for each.
- Return VALID JSON only.
- {language_instruction(language_mode)}

Return JSON in this format:
{{
  "title": "short quiz title",
  "questions": [
    {{"q": "question 1", "answer_key": "expected short answer"}},
    {{"q": "question 2", "answer_key": "expected short answer"}}
  ]
}}

Context:
{context}

Topic:
{topic}
""".strip()


def build_quiz_eval_prompt(language_mode: str, quiz_data: dict, user_answers: str) -> str:
    quiz_json = json.dumps(quiz_data, ensure_ascii=False)

    return f"""
You are BrainChat, an interactive neurology tutor.

Evaluate the student's answers fairly using the answer keys.
Accept semantically correct answers even if wording differs.
Give constructive feedback.
Return VALID JSON only.

Return JSON in this format:
{{
  "score_obtained": 0,
  "score_total": 0,
  "summary": "short overall feedback",
  "results": [
    {{
      "question": "question text",
      "answer_key": "expected answer",
      "student_answer": "student answer",
      "result": "Correct / Partially Correct / Incorrect",
      "feedback": "short explanation"
    }}
  ],
  "improvement_tip": "one short study suggestion"
}}

Quiz:
{quiz_json}

Student answers:
{user_answers}

Language:
{language_instruction(language_mode)}
""".strip()


# =====================================================
# OPENAI
# =====================================================
def oai_text(prompt: str) -> str:
    ensure_loaded()

    resp = CLIENT.chat.completions.create(
        model=OPENAI_MODEL,
        temperature=0.2,
        messages=[
            {
                "role": "system",
                "content": "You are BrainChat, a careful educational assistant for neurology students."
            },
            {"role": "user", "content": prompt},
        ],
    )

    return resp.choices[0].message.content.strip()


def oai_json(prompt: str) -> dict:
    ensure_loaded()

    resp = CLIENT.chat.completions.create(
        model=OPENAI_MODEL,
        temperature=0.2,
        response_format={"type": "json_object"},
        messages=[
            {"role": "system", "content": "Return only valid JSON."},
            {"role": "user", "content": prompt},
        ],
    )

    return json.loads(resp.choices[0].message.content)


# =====================================================
# LOGO
# =====================================================
def get_logo_data_uri():
    if not os.path.exists(LOGO_FILE):
        return None

    mime_type, _ = mimetypes.guess_type(LOGO_FILE)

    if not mime_type:
        mime_type = "image/png"

    with open(LOGO_FILE, "rb") as f:
        encoded = base64.b64encode(f.read()).decode("utf-8")

    return f"data:{mime_type};base64,{encoded}"


def render_logo():
    data_uri = get_logo_data_uri()

    if data_uri:
        return f'<img src="{data_uri}" alt="BrainChat logo" class="bc-logo-img">'

    return '<div class="bc-logo-fallback">BRAIN<br>CHAT</div>'


# =====================================================
# CHAT HTML
# =====================================================
def format_text(text: str) -> str:
    safe = (
        text.replace("&", "&amp;")
            .replace("<", "&lt;")
            .replace(">", "&gt;")
    )

    safe = re.sub(r"\*\*(.+?)\*\*", r"<strong>\1</strong>", safe)
    safe = safe.replace("\n", "<br>")

    return safe


def render_chat(history):
    if not history:
        return """
        <div class="bc-empty">
            <div class="bc-empty-text">
                <strong>Welcome to BrainChat.</strong><br><br>
                I am your AI tutor for Neurology and PMQSN.<br>
                You can ask questions, request explanations, practise clinical cases,<br>
                or generate quizzes. I first use Professor Handouts,<br>
                then supporting textbooks if needed.
            </div>
        </div>
        """

    rows = []

    for item in history:
        role = item["role"]
        content = format_text(item["content"])
        confidence_block = item.get("confidence_html", "")

        if role == "user":
            rows.append(
                f'<div class="bc-row bc-user-row"><div class="bc-bubble bc-user-bubble">{content}</div></div>'
            )
        else:
            rows.append(
                f'<div class="bc-row bc-bot-row"><div class="bc-bubble bc-bot-bubble">{confidence_block}{content}</div></div>'
            )

    return f"""
    <div class="bc-chat-wrap" id="bc-chat-wrap">
        {''.join(rows)}
    </div>
    <script>
        const chatWrap = document.getElementById("bc-chat-wrap");
        if (chatWrap) {{
            chatWrap.scrollTop = chatWrap.scrollHeight;
        }}
    </script>
    """


# =====================================================
# MAIN LOGIC
# =====================================================
def respond(user_msg, history, mode, language_mode, quiz_count_mode, show_sources, quiz_state):
    history = history or []
    quiz_state = quiz_state or {
        "active": False,
        "quiz_data": None,
        "language_mode": "Auto"
    }

    text = (user_msg or "").strip()

    if not text:
        return "", history, render_chat(history), quiz_state, render_dashboard()

    try:
        history = history + [{"role": "user", "content": text}]

        if is_general_chat(text):
            reply = general_chat_reply(text, language_mode)

            conf = {
                "level": "green",
                "label": "Ready",
                "score": 1.0
            }

            log_event(
                event_type="general_chat",
                mode=mode,
                language=language_mode,
                confidence_level="green",
                similarity=1.0,
                query=text
            )

            history = history + [
                {
                    "role": "assistant",
                    "content": reply,
                    "confidence_html": confidence_html(conf)
                }
            ]

            return "", history, render_chat(history), quiz_state, render_dashboard()

        if quiz_state.get("active", False):
            evaluation = oai_json(
                build_quiz_eval_prompt(
                    quiz_state.get("language_mode", language_mode),
                    quiz_state.get("quiz_data", {}),
                    text
                )
            )

            lines = []
            lines.append(
                f"**Score:** {evaluation.get('score_obtained', 0)}/{evaluation.get('score_total', 0)}"
            )

            if evaluation.get("summary"):
                lines.append(f"\n**Overall feedback:** {evaluation['summary']}")

            if evaluation.get("improvement_tip"):
                lines.append(f"\n**Study tip:** {evaluation['improvement_tip']}\n")

            results = evaluation.get("results", [])

            if results:
                lines.append("**Question-wise feedback:**")

                for item in results:
                    lines.append("")
                    lines.append(f"**Q:** {item.get('question','')}")
                    lines.append(f"**Your answer:** {item.get('student_answer','')}")
                    lines.append(f"**Expected answer:** {item.get('answer_key','')}")
                    lines.append(f"**Result:** {item.get('result','')}")
                    lines.append(f"**Feedback:** {item.get('feedback','')}")

            log_event(
                event_type="quiz_evaluated",
                mode=mode,
                language=language_mode,
                confidence_level="green",
                similarity=1.0,
                query=text
            )

            conf = {
                "level": "green",
                "label": "Quiz evaluated",
                "score": 1.0
            }

            history = history + [
                {
                    "role": "assistant",
                    "content": "\n".join(lines).strip(),
                    "confidence_html": confidence_html(conf)
                }
            ]

            quiz_state = {
                "active": False,
                "quiz_data": None,
                "language_mode": language_mode
            }

            return "", history, render_chat(history), quiz_state, render_dashboard()

        records = search_hybrid(text, shortlist_k=30, final_k=5)
        context = build_context(records)

        if mode == "Quiz Me":
            n_questions = choose_quiz_count(text, quiz_count_mode)

            quiz_data = oai_json(
                build_quiz_generation_prompt(
                    language_mode,
                    text,
                    context,
                    n_questions
                )
            )

            conf = compute_confidence(records, "quiz generated")

            lines = []
            lines.append(f"**{quiz_data.get('title', 'Quiz')}**")
            lines.append(f"\n**Total questions:** {len(quiz_data.get('questions', []))}\n")
            lines.append("Reply in one message using numbered answers.")
            lines.append("Example: 1. ... 2. ...\n")

            for i, q in enumerate(quiz_data.get("questions", []), start=1):
                lines.append(f"**Q{i}.** {q.get('q','')}")

            if show_sources and conf["level"] != "red":
                lines.append("\n\n**References used to create this quiz:**")
                lines.append(make_sources(records))

            log_event(
                event_type="quiz_generated",
                mode=mode,
                language=language_mode,
                confidence_level=conf["level"],
                similarity=conf["score"],
                query=text
            )

            history = history + [
                {
                    "role": "assistant",
                    "content": "\n".join(lines).strip(),
                    "confidence_html": confidence_html(conf)
                }
            ]

            quiz_state = {
                "active": True,
                "quiz_data": quiz_data,
                "language_mode": language_mode
            }

            return "", history, render_chat(history), quiz_state, render_dashboard()

        answer = oai_text(
            build_tutor_prompt(
                mode,
                language_mode,
                text,
                context
            )
        )

        conf = compute_confidence(records, answer)

        if conf["level"] == "red":
            if language_mode == "English":
                final_answer = "Not found in the course material."
            else:
                final_answer = "No encontrado en el material del curso."
        else:
            final_answer = answer.strip()

            if show_sources:
                final_answer += "\n\n**References used:**\n" + make_sources(records)

        log_event(
            event_type="answer",
            mode=mode,
            language=language_mode,
            confidence_level=conf["level"],
            similarity=conf["score"],
            query=text
        )

        history = history + [
            {
                "role": "assistant",
                "content": final_answer.strip(),
                "confidence_html": confidence_html(conf)
            }
        ]

        return "", history, render_chat(history), quiz_state, render_dashboard()

    except Exception as e:
        history = history + [{"role": "assistant", "content": f"Error: {str(e)}"}]

        quiz_state = {
            "active": False,
            "quiz_data": None,
            "language_mode": language_mode
        }

        return "", history, render_chat(history), quiz_state, render_dashboard()


def clear_all():
    empty_history = []
    empty_quiz = {
        "active": False,
        "quiz_data": None,
        "language_mode": "Auto"
    }

    return "", empty_history, render_chat(empty_history), empty_quiz, render_dashboard()


# =====================================================
# CSS
# =====================================================
CSS = """
:root{
  --page-bg: #d9d9dd;
  --uva-purple: #5a2d77;
  --uva-purple-light: #7b3f98;
  --uva-gold: #c7a008;
  --uva-gold-light: #fff8cc;
  --uva-soft-purple: #efe7f6;
  --text-dark: #241336;
  --shadow: rgba(30,20,50,0.18);
}

html, body, .gradio-container{
  background: var(--page-bg) !important;
  font-family: Arial, Helvetica, sans-serif !important;
}

footer{
  display:none !important;
}

#bc_app{
  max-width: 1100px;
  margin: 18px auto;
}

/* SETTINGS BOX - UVa style light theme */
.bc-settings{
  background:#ffffff;
  border-radius:22px;
  padding:18px;
  box-shadow:0 12px 28px rgba(0,0,0,0.22);
  margin-bottom:16px;
  border-top:8px solid #5a2d77;
  color:#241336 !important;
}

.bc-settings label{
  color:#5a2d77 !important;
  font-weight:800 !important;
}

.bc-settings input,
.bc-settings textarea,
.bc-settings select{
  color:#241336 !important;
  background:#ffffff !important;
}

.bc-howto{
  margin-top:12px;
  padding:16px;
  border-radius:16px;
  background:#f4edf7;
  color:#241336 !important;
  font-size:14px;
  line-height:1.55;
  border-left:6px solid #c7a008;
}

.bc-howto strong{
  color:#5a2d77 !important;
}

/* CHAT WINDOW */
.bc-phone{
  position: relative;
  background: #ffffff;
  border-radius: 30px;
  padding: 92px 14px 14px 14px;
  box-shadow: 0 16px 34px rgba(0,0,0,0.22);
  min-height: 620px;
  border-top: 8px solid #5a2d77;
}

.bc-logo-holder{
  position: absolute;
  top: 16px;
  left: 50%;
  transform: translateX(-50%);
  width: 104px;
  height: 104px;
  border-radius: 999px;
  background: #c7a008;
  display: flex;
  align-items: center;
  justify-content: center;
  box-shadow: 0 10px 22px rgba(0,0,0,0.18);
}

.bc-logo-img{
  width: 88px;
  height: 88px;
  object-fit: contain;
  display:block;
}

.bc-logo-fallback{
  width: 88px;
  height: 88px;
  border-radius: 999px;
  display:flex;
  align-items:center;
  justify-content:center;
  text-align:center;
  font-size: 13px;
  font-weight: 900;
  color: #241336;
  background: rgba(255,255,255,0.55);
  line-height: 1.05;
}

.bc-chat-shell{
  background:#ffffff;
  border-radius:20px;
  padding:16px;
  min-height:460px;
  box-shadow: inset 0 0 0 2px #d8c6e8;
}

.bc-chat-wrap{
  display: flex;
  flex-direction: column;
  gap: 14px;
  max-height: 460px;
  overflow-y: auto;
  padding-right: 4px;
}

.bc-chat-wrap::-webkit-scrollbar{
  width: 8px;
}

.bc-chat-wrap::-webkit-scrollbar-thumb{
  background: #c7a008;
  border-radius: 999px;
}

.bc-row{
  display:flex;
  width:100%;
}

.bc-user-row{
  justify-content: flex-start;
}

.bc-bot-row{
  justify-content: flex-end;
}

.bc-bubble{
  max-width: 82%;
  padding: 15px 18px;
  border-radius: 22px;
  line-height: 1.6;
  font-size: 15px;
  box-shadow: 0 10px 18px rgba(0,0,0,0.10);
  word-wrap: break-word;
  font-weight: 500;
}

.bc-user-bubble{
  background: #efe7f6;
  color: #241336 !important;
  border: 2px solid #d8c6e8;
  border-bottom-left-radius: 8px;
}

.bc-bot-bubble{
  background: #fff8cc;
  color: #241336 !important;
  border: 2px solid #c7a008;
  border-bottom-right-radius: 8px;
}

.bc-bubble strong{
  color: #241336 !important;
}

.bc-confidence{
  display:flex;
  align-items:center;
  gap:8px;
  margin-bottom:10px;
  padding:7px 10px;
  background:rgba(255,255,255,0.75);
  border-radius:999px;
  font-size:13px;
  color:#111827;
  border:1px solid #e5d8ef;
}

.bc-dot{
  width:15px;
  height:15px;
  border-radius:999px;
  display:inline-block;
  box-shadow:0 0 0 3px rgba(255,255,255,0.75);
}

.bc-empty{
  display:flex;
  justify-content:center;
  align-items:center;
  min-height: 400px;
}

.bc-empty-text{
  color:#5a2d77 !important;
  text-align:center;
  opacity:1 !important;
  font-size:16px;
  line-height:1.7;
  font-weight:700;
}

.bc-input-bar{
  margin-top: 12px;
  background: #5a2d77;
  border-radius: 999px;
  padding: 8px 10px;
  display:flex;
  align-items:center;
  gap: 10px;
  box-shadow: 0 10px 22px rgba(0,0,0,0.14);
}

.bc-plus{
  width: 38px;
  height: 38px;
  border-radius: 999px;
  background: #c7a008;
  display:flex;
  align-items:center;
  justify-content:center;
  font-size: 30px;
  font-weight: 900;
  color: #ffffff;
  user-select:none;
}

#bc_msg textarea{
  background: #ffffff !important;
  border: 2px solid #c7a008 !important;
  box-shadow: none !important;
  border-radius: 999px !important;
  color: #241336 !important;
  padding: 11px 14px !important;
  min-height: 42px !important;
}

#bc_msg textarea::placeholder{
  color: rgba(34,23,53,0.72) !important;
}

#bc_send button{
  min-width: 48px !important;
  height: 42px !important;
  border-radius: 999px !important;
  border: none !important;
  background: #c7a008 !important;
  color: #ffffff !important;
  font-size: 20px !important;
  font-weight: 900 !important;
  box-shadow: none !important;
}

#bc_send button:hover{
  background: #9f8006 !important;
}

#bc_clear button, #bc_refresh button, #bc_clear_analytics button{
  border-radius: 14px !important;
}

/* DASHBOARD */
.bc-dashboard{
  background:#ffffff;
  border-radius:22px;
  padding:22px;
  box-shadow:0 12px 28px rgba(0,0,0,0.22);
  margin-top:18px;
  color:#241336 !important;
  border-top:8px solid #5a2d77;
}

.bc-dashboard h3{
  color:#5a2d77 !important;
  font-size:22px;
  font-weight:800;
  margin-bottom:10px;
}

.bc-dashboard h4{
  color:#5a2d77 !important;
  font-size:17px;
  font-weight:800;
}

.bc-dashboard p{
  color:#241336 !important;
  font-size:14px;
  line-height:1.5;
}

.bc-dashboard-grid{
  display:grid;
  grid-template-columns: 2fr 1fr;
  gap:20px;
  align-items:start;
}

.bc-dashboard-help{
  background:#f4edf7;
  border-left:6px solid #c7a008;
  border-radius:16px;
  padding:16px;
  color:#241336 !important;
}

.bc-dashboard-help strong{
  color:#5a2d77 !important;
}

.bc-metrics{
  display:grid;
  grid-template-columns: repeat(3, 1fr);
  gap:14px;
  margin:16px 0;
}

.bc-card{
  border-radius:16px;
  padding:16px;
  text-align:center;
  font-size:14px;
  color:#241336 !important;
  border:2px solid #e5d8ef;
  font-weight:600;
}

.bc-card strong{
  display:block;
  font-size:28px;
  color:#5a2d77 !important;
  margin-bottom:4px;
}

.bc-card.total{ background:#efe7f6; }
.bc-card.green{ background:#dff7e7; border-color:#22c55e; }
.bc-card.orange{ background:#fff1d6; border-color:#f59e0b; }
.bc-card.red{ background:#ffe1e1; border-color:#dc2626; }
.bc-card.quiz{ background:#f7edff; border-color:#8b5cf6; }
.bc-card.avg{ background:#fff8cc; border-color:#c7a008; }

.bc-table{
  width:100%;
  border-collapse:collapse;
  font-size:13px;
  background:#ffffff;
  color:#241336 !important;
  margin-top:12px;
}

.bc-table th{
  background:#5a2d77;
  color:#ffffff !important;
  padding:10px;
  border:1px solid #ddd;
  font-weight:700;
}

.bc-table td{
  border:1px solid #ddd;
  padding:9px;
  vertical-align:top;
  color:#241336 !important;
  background:#ffffff;
}

.bc-table tr:nth-child(even) td{
  background:#faf7fc;
}

.bc-pill{
  padding:5px 10px;
  border-radius:999px;
  font-weight:800;
  color:#241336 !important;
}

.bc-green{ background:#86efac; }
.bc-orange{ background:#fdba74; }
.bc-red{ background:#fca5a5; }

@media (max-width: 768px){
  #bc_app{
    max-width: 96vw;
  }

  .bc-bubble{
    max-width: 90%;
  }

  .bc-dashboard-grid{
    grid-template-columns: 1fr;
  }

  .bc-metrics{
    grid-template-columns: 1fr;
  }
}
"""


# =====================================================
# UI
# =====================================================
with gr.Blocks() as demo:
    history_state = gr.State([])
    quiz_state = gr.State({
        "active": False,
        "quiz_data": None,
        "language_mode": "Auto"
    })

    with gr.Column(elem_id="bc_app"):

        with gr.Group(elem_classes="bc-settings"):
            with gr.Row():
                mode = gr.Dropdown(
                    choices=[
                        "Explain",
                        "Detailed",
                        "Short Notes",
                        "Flashcards",
                        "Case-Based",
                        "Quiz Me"
                    ],
                    value="Explain",
                    label="Tutor Mode"
                )

                language_mode = gr.Dropdown(
                    choices=[
                        "Auto",
                        "Spanish",
                        "English",
                        "Bilingual"
                    ],
                    value="Spanish",
                    label="Answer Language"
                )

            with gr.Row():
                quiz_count_mode = gr.Dropdown(
                    choices=[
                        "Auto",
                        "3",
                        "5",
                        "7"
                    ],
                    value="Auto",
                    label="Quiz Questions"
                )

                show_sources = gr.Checkbox(
                    value=True,
                    label="Show References"
                )

            gr.HTML("""
            <div class="bc-howto">
              <strong>Welcome to BrainChat</strong><br>
              BrainChat is an AI-based neurology tutor designed to support PMQSN learning.<br>
              It first searches <strong>Professor Handouts</strong>, and then uses other textbooks only when needed.<br><br>

              <strong>Confidence indicator</strong><br>
              🟢 Strong support from course material &nbsp; | &nbsp;
              🟠 Partial support &nbsp; | &nbsp;
              🔴 Not found / weak evidence<br><br>

              <strong>How to use</strong><br>
              1. Choose a tutor mode: Explain, Detailed, Short Notes, Flashcards, Case-Based, or Quiz Me.<br>
              2. Select the answer language: Spanish, English, Bilingual, or Auto.<br>
              3. Type your question in the message box below.<br>
              4. Use Quiz Me to practise questions and receive automatic feedback.<br><br>

              <strong>Example prompts</strong><br>
              • Explícame la afasia de Broca según los apuntes.<br>
              • Ponme 3 preguntas tipo test sobre ictus.<br>
              • Explícame la diferencia diagnóstica entre EM y NMOSD.<br>
              • Dame un caso clínico sobre epilepsia.
            </div>
            """)

        with gr.Group(elem_classes="bc-phone"):
            gr.HTML(f'<div class="bc-logo-holder">{render_logo()}</div>')

            chat_html = gr.HTML(
                f'<div class="bc-chat-shell">{render_chat([])}</div>'
            )

            with gr.Row(elem_classes="bc-input-bar"):
                gr.HTML('<div class="bc-plus">+</div>')

                msg = gr.Textbox(
                    placeholder="Type a message...",
                    show_label=False,
                    container=False,
                    scale=8,
                    elem_id="bc_msg"
                )

                send_btn = gr.Button(
                    "➤",
                    elem_id="bc_send",
                    scale=1
                )

        with gr.Row():
            clear_btn = gr.Button("Clear Chat", elem_id="bc_clear")
            refresh_btn = gr.Button("Refresh Dashboard", elem_id="bc_refresh")
            clear_analytics_btn = gr.Button("Clear Analytics", elem_id="bc_clear_analytics")

        dashboard_html = gr.HTML(render_dashboard())

    msg.submit(
        respond,
        inputs=[
            msg,
            history_state,
            mode,
            language_mode,
            quiz_count_mode,
            show_sources,
            quiz_state
        ],
        outputs=[
            msg,
            history_state,
            chat_html,
            quiz_state,
            dashboard_html
        ]
    )

    send_btn.click(
        respond,
        inputs=[
            msg,
            history_state,
            mode,
            language_mode,
            quiz_count_mode,
            show_sources,
            quiz_state
        ],
        outputs=[
            msg,
            history_state,
            chat_html,
            quiz_state,
            dashboard_html
        ]
    )

    clear_btn.click(
        clear_all,
        inputs=[],
        outputs=[
            msg,
            history_state,
            chat_html,
            quiz_state,
            dashboard_html
        ],
        queue=False
    )

    refresh_btn.click(
        refresh_dashboard,
        inputs=[],
        outputs=[dashboard_html],
        queue=False
    )

    clear_analytics_btn.click(
        clear_analytics,
        inputs=[],
        outputs=[dashboard_html],
        queue=False
    )


if __name__ == "__main__":
    demo.queue()
    demo.launch(css=CSS)