src/app/main_app.py

###############################################################
# main_app.py — Agentic Language Partner UI (Streamlit)
###############################################################
import pandas as pd
import json
import random
import re
from datetime import datetime
from pathlib import Path
from typing import Dict, List, Any

import streamlit as st
import streamlit.components.v1 as components
from deep_translator import GoogleTranslator
from pydub import AudioSegment
from io import BytesIO

from .auth import (
    authenticate_user,
    register_user,
    get_user_prefs,
    update_user_prefs,
)
from .config import get_user_dir
from .conversation_core import ConversationManager
from .flashcards_tools import (
    list_user_decks,
    load_deck,
    _get_decks_dir,
    save_deck,
    generate_flashcards_from_text,
    generate_flashcards_from_ocr_results,
)
from .ocr_tools import ocr_and_translate_batch
from .viewers import generate_flashcard_viewer_for_user


###############################################################
# PAGE + GLOBAL STYLE
###############################################################

st.set_page_config(
    page_title="Agentic Language Partner",
    layout="wide",
    page_icon="🌐",
)

st.markdown(
    """
<style>
.chat-column {
    display: flex;
    flex-direction: column;
}
/* Input bar at the top */
.chat-input-bar {
    margin-bottom: 0.5rem;
    background-color: #111;
    padding: 0.75rem 0.5rem 0.5rem;
    border: 1px solid #333;
    border-radius: 0.5rem;
}
/* Scrollable chat messages below input */
.chat-window {
    max-height: 65vh;
    overflow-y: auto;
    padding-right: .75rem;
    padding-bottom: 0.5rem;
}
/* Chat bubbles */
.chat-row-user { justify-content:flex-end; display:flex; margin-bottom:.4rem; }
.chat-row-assistant { justify-content:flex-start; display:flex; margin-bottom:.4rem; }
.chat-bubble {
    border-radius:14px;
    padding:.55rem .95rem;
    max-width:80%;
    line-height:1.4;
    box-shadow:0 2px 5px rgba(0,0,0,0.4);
    font-size:1.05rem;  /* larger for readability */
}
.chat-bubble-user { background:#3a3b3c; color:white; }
.chat-bubble-assistant { background:#1a73e8; color:white; }
.chat-aux {
    font-size:1.0rem;   /* larger translation/explanation */
    color:#ccc;
    margin:0.1rem 0.25rem 0.5rem 0.25rem;
}
/* Lock viewport height and avoid infinite page scrolling */
html, body {
    height: 100%;
    overflow: hidden !important;
}
.block-container {
    height: 100vh !important;
    overflow-y: auto !important;
}
.saved-conv-panel { max-width: 360px; }
</style>
""",
    unsafe_allow_html=True,
)


###############################################################
# HELPERS / GLOBALS
###############################################################
# ------------------------------------------------------------
# Model preload / Conversation manager
# ------------------------------------------------------------

def preload_models():
    """
    Loads all heavy models ONCE at startup.
    Safe for HuggingFace Spaces CPU environment.
    """
    from .conversation_core import load_partner_lm, load_whisper_pipe

    # Qwen LM
    try:
        load_partner_lm()
    except Exception as e:
        print("[preload_models] ERROR loading Qwen model:", e)

    # Whisper ASR
    try:
        load_whisper_pipe()
    except Exception as e:
        print("[preload_models] ERROR loading Whisper pipeline:", e)


def get_conv_manager() -> ConversationManager:
    if "conv_manager" not in st.session_state:
        prefs = st.session_state["prefs"]
        st.session_state["conv_manager"] = ConversationManager(
            target_language=prefs.get("target_language", "english"),
            native_language=prefs.get("native_language", "english"),
            cefr_level=prefs.get("cefr_level", "B1"),
            topic=prefs.get("topic", "general conversation"),
        )
    return st.session_state["conv_manager"]


def ensure_default_decks(username: str):
    decks_dir = _get_decks_dir(username)

    alpha = decks_dir / "alphabet.json"
    if not alpha.exists():
        save_deck(alpha, {
            "name": "Alphabet (A–Z)",
            "cards": [{"front": chr(65+i), "back": f"Letter {chr(65+i)}"} for i in range(26)],
            "tags": ["starter"],
        })

    nums = decks_dir / "numbers_1_10.json"
    if not nums.exists():
        save_deck(nums, {
            "name": "Numbers 1–10",
            "cards": [{"front": str(i), "back": f"Number {i}"} for i in range(1, 11)],
            "tags": ["starter"],
        })

    greetings = decks_dir / "greetings_intros.json"
    if not greetings.exists():
        save_deck(greetings, {
            "name": "Greetings & Introductions",
            "cards": [
                {"front": "Hallo!", "back": "Hello!"},
                {"front": "Wie geht's?", "back": "How are you?"},
                {"front": "Ich heiße …", "back": "My name is …"},
                {"front": "Freut mich!", "back": "Nice to meet you!"},
            ],
            "tags": ["starter"],
        })


def ui_clean_assistant_text(text: str) -> str:
    if not text:
        return ""
    text = re.sub(r"(?i)\b(user|assistant|system):\s*", "", text)
    text = re.sub(r"\s{2,}", " ", text)
    return text.strip()


def save_current_conversation(username: str, name: str) -> Path:
    """Save chat_history as JSON, stripping non-serializable fields (audio bytes)."""
    user_dir = get_user_dir(username)
    save_dir = user_dir / "chats" / "saved"
    save_dir.mkdir(parents=True, exist_ok=True)

    cleaned_messages = []
    for m in st.session_state.get("chat_history", []):
        cleaned_messages.append(
            {
                "role": m.get("role"),
                "text": m.get("text"),
                "explanation": m.get("explanation"),
                # store only a flag for audio, not raw bytes
                "audio_present": bool(m.get("audio")),
            }
        )

    payload = {
        "name": name,
        "timestamp": datetime.utcnow().isoformat(),
        "messages": cleaned_messages,
    }
    fname = datetime.utcnow().strftime("%Y%m%d_%H%M%S") + ".json"
    path = save_dir / fname
    path.write_text(json.dumps(payload, ensure_ascii=False, indent=2), encoding="utf-8")
    return path


###############################################################
# CHAT HANDLING
###############################################################

def handle_user_message(username: str, text: str):
    text = text.strip()
    if not text:
        return

    conv = get_conv_manager()

    st.session_state["chat_history"].append(
        {"role": "user", "text": text, "audio": None, "explanation": None}
    )

    with st.spinner("Thinking…"):
        result = conv.reply(text)

    reply_text = ui_clean_assistant_text(result.get("reply_text", ""))
    reply_audio = result.get("audio", None)
    explanation = ui_clean_assistant_text(result.get("explanation", ""))

    st.session_state["chat_history"].append(
        {
            "role": "assistant",
            "text": reply_text,
            "audio": reply_audio,
            "explanation": explanation,
        }
    )


###############################################################
# AUTH
###############################################################

def login_view():
    st.title("🌐 Agentic Language Partner")
    tab1, tab2 = st.tabs(["Login", "Register"])

    with tab1:
        u = st.text_input("Username")
        p = st.text_input("Password", type="password")
        if st.button("Login"):
            if authenticate_user(u, p):
                st.session_state["user"] = u
                st.session_state["prefs"] = get_user_prefs(u)
                st.experimental_rerun()
            else:
                st.error("Invalid login.")

    with tab2:
        u = st.text_input("New username")
        p = st.text_input("New password", type="password")
        if st.button("Register"):
            if register_user(u, p):
                st.success("Registered! Please log in.")
            else:
                st.error("Username already exists.")


###############################################################
# SIDEBAR SETTINGS
###############################################################

def sidebar_settings(username: str):
    st.sidebar.header("⚙ Settings")

    prefs = st.session_state["prefs"]
    langs = ["english", "spanish", "german", "russian", "japanese", "chinese", "korean"]

    tgt = st.sidebar.selectbox(
        "Target language",
        langs,
        index=langs.index(prefs.get("target_language", "english")),
        key="sidebar_target",
    )
    nat = st.sidebar.selectbox(
        "Native language",
        langs,
        index=langs.index(prefs.get("native_language", "english")),
        key="sidebar_native",
    )

    cefr_levels = ["A1", "A2", "B1", "B2", "C1", "C2"]
    level = st.sidebar.selectbox(
        "CEFR Level",
        cefr_levels,
        index=cefr_levels.index(prefs.get("cefr_level", "B1")),
        key="sidebar_cefr",
    )

    topic = st.sidebar.text_input(
        "Conversation Topic",
        prefs.get("topic", "general conversation"),
        key="sidebar_topic",
    )

    show_exp = st.sidebar.checkbox(
        "Show Explanations",
        value=prefs.get("show_explanations", True),
        key="sidebar_show_exp",
    )

    if st.sidebar.button("Save Settings"):
        new = {
            "target_language": tgt,
            "native_language": nat,
            "cefr_level": level,
            "topic": topic,
            "show_explanations": show_exp,
        }
        st.session_state["prefs"] = new
        update_user_prefs(username, new)
        if "conv_manager" in st.session_state:
            del st.session_state["conv_manager"]
        st.sidebar.success("Settings saved!")


###############################################################
# DASHBOARD TAB
###############################################################

def dashboard_tab(username: str):
    st.title("Agentic Language Partner — Dashboard")

    prefs = st.session_state["prefs"]
    langs = ["english", "spanish", "german", "russian", "japanese", "chinese", "korean"]

    st.subheader("Language Settings")

    col1, col2, col3 = st.columns(3)

    with col1:
        native = st.selectbox(
            "Native language",
            langs,
            index=langs.index(prefs.get("native_language", "english")),
            key="dash_native_language",
        )
    with col2:
        target = st.selectbox(
            "Target language",
            langs,
            index=langs.index(prefs.get("target_language", "english")),
            key="dash_target_language",
        )
    with col3:
        cefr_levels = ["A1", "A2", "B1", "B2", "C1", "C2"]
        level = st.selectbox(
            "CEFR Level",
            cefr_levels,
            index=cefr_levels.index(prefs.get("cefr_level", "B1")),
            key="dash_cefr_level",
        )

    topic = st.text_input(
        "Conversation Topic",
        prefs.get("topic", "general conversation"),
        key="dash_topic",
    )

    if st.button("Save Language Settings", key="dash_save_lang"):
        new = {
            "native_language": native,
            "target_language": target,
            "cefr_level": level,
            "topic": topic,
            "show_explanations": prefs.get("show_explanations", True),
        }
        st.session_state["prefs"] = new
        update_user_prefs(username, new)
        if "conv_manager" in st.session_state:
            del st.session_state["conv_manager"]
        st.success("Language settings saved!")

    st.markdown("---")

    ###########################################################
    # MICROPHONE & TRANSCRIPTION CALIBRATION (native phrase)
    ###########################################################
    st.subheader("Microphone & Transcription Calibration")

    st.write(
        "To verify that audio recording and transcription are working, "
        "please repeat this phrase in your native language:\n\n"
        "> \"Hello, my name is [Your Name], and I am here to practice languages.\"\n\n"
        "Upload or record a short clip and then run transcription to check accuracy."
    )

    calib_col1, calib_col2 = st.columns([2, 1])

    with calib_col1:
        calib_file = st.file_uploader(
            "Upload or record a short audio sample (e.g., WAV/MP3)",
            type=["wav", "mp3", "m4a", "ogg"],
            key="calibration_audio",
        )

        if calib_file is not None:
            st.caption("Calibration audio loaded. Click 'Transcribe sample' to test.")

        if st.button("Transcribe sample", key="calibration_transcribe"):
            if calib_file is None:
                st.warning("Please upload or record a short calibration clip first.")
            else:
                conv = get_conv_manager()
                try:
                    raw = calib_file.read()
                    seg = AudioSegment.from_file(BytesIO(raw))
                    seg = seg.set_frame_rate(16000).set_channels(1)
                    with st.spinner("Transcribing calibration audio…"):
                        text_out, det_lang, det_prob = conv.transcribe(
                            seg,
                            spoken_lang=st.session_state["prefs"]["native_language"]
                        )
                    st.session_state["calibration_result"] = {
                        "text": text_out,
                        "det_lang": det_lang,
                        "det_prob": det_prob,
                    }
                    st.success("Calibration transcript updated.")
                except Exception as e:
                    st.error(f"Calibration error: {e}")

    with calib_col2:
        if st.session_state.get("calibration_result"):
            res = st.session_state["calibration_result"]
            st.markdown("**Calibration transcript:**")
            st.info(res.get("text", ""))
            st.caption(
                f"Detected lang: {res.get('det_lang','?')} · Confidence ~ {res.get('det_prob', 0):.2f}"
            )
        else:
            st.caption("No calibration transcript yet.")

    st.markdown("---")

    ###########################################################
    # TOOL OVERVIEW
    ###########################################################
    st.subheader("Tools Overview")

    c1, c2, c3 = st.columns(3)
    with c1:
        st.markdown("### 🎙️ Conversation Partner")
        st.write("Real-time language practice with microphone support (via audio uploads).")

    with c2:
        st.markdown("### 🃏 Flashcards & Quizzes")
        st.write("Starter decks: Alphabet, Numbers, Greetings.")

    with c3:
        st.markdown("### 📷 OCR Helper")
        st.write("Upload images to extract and translate text.")


# ------------------------------------------------------------
# Settings tab  (restore missing function)
# ------------------------------------------------------------
def settings_tab(username: str):
    """Minimal settings tab so main() can call it safely."""
    st.header("Settings")

    st.subheader("User Preferences")
    prefs = st.session_state.get("prefs", {})
    st.json(prefs)

    st.markdown("---")

    st.subheader("System Status")
    st.write("Models preloaded:", st.session_state.get("models_loaded", False))

    st.markdown(
        "This is a placeholder settings panel. "
        "You can customize this later with user-specific configuration."
    )


###############################################################
# CONVERSATION TAB
###############################################################

def conversation_tab(username: str):
    import re
    from datetime import datetime
    from deep_translator import GoogleTranslator

    st.header("Conversation")

    # ------------------------------------------
    # INITIAL STATE
    # ------------------------------------------
    if "chat_history" not in st.session_state:
        st.session_state["chat_history"] = []

    if "pending_transcript" not in st.session_state:
        st.session_state["pending_transcript"] = ""

    if "speech_state" not in st.session_state:
        st.session_state["speech_state"] = "idle"   # idle | pending_speech

    if "recorder_key" not in st.session_state:
        st.session_state["recorder_key"] = 0

    conv = get_conv_manager()
    prefs = st.session_state.get("prefs", {})
    show_exp = prefs.get("show_explanations", True)

    # ------------------------------------------
    # RESET BUTTON (ONLY ONE)
    # ------------------------------------------
    if st.button("🔄 Reset Conversation"):
        st.session_state["chat_history"] = []
        st.session_state["pending_transcript"] = ""
        st.session_state["speech_state"] = "idle"
        st.session_state["recorder_key"] += 1
        st.experimental_rerun()

    # ------------------------------------------
    # FIRST MESSAGE GREETING
    # ------------------------------------------
    if len(st.session_state["chat_history"]) == 0:
        lang = conv.target_language.lower()
        topic = prefs.get("topic", "").strip()

        default_greetings = {
            "english": "Hello! I heard you want to practice with me. How is your day going?",
            "german":  "Hallo! Ich habe gehört, dass du üben möchtest. Wie geht dein Tag bisher?",
            "spanish": "¡Hola! Escuché que querías practicar conmigo. ¿Cómo va tu día?",
            "japanese":"こんにちは！練習したいと聞きました。今日はどんな一日ですか？",
        }

        intro = default_greetings.get(lang, default_greetings["english"])

        if topic and topic.lower() != "general conversation":
            try:
                intro = GoogleTranslator(source="en", target=lang).translate(
                    f"Hello! Let's talk about {topic}. What do you think about it?"
                )
            except Exception:
                pass

        st.session_state["chat_history"].append(
            {"role":"assistant","text":intro,"audio":None,"explanation":None}
        )

    # ------------------------------------------
    # LAYOUT
    # ------------------------------------------
    col_chat, col_saved = st.columns([3,1])

    # ===========================
    # LEFT: CHAT WINDOW
    # ===========================
    with col_chat:

        st.markdown('<div class="chat-window">', unsafe_allow_html=True)

        for msg in st.session_state["chat_history"]:
            role = msg["role"]
            bubble = "chat-bubble-user" if role == "user" else "chat-bubble-assistant"
            row    = "chat-row-user"     if role == "user" else "chat-row-assistant"

            st.markdown(
                f'<div class="{row}"><div class="chat-bubble {bubble}">{msg["text"]}</div></div>',
                unsafe_allow_html=True,
            )

            if role == "assistant" and msg.get("audio"):
                st.audio(msg["audio"], format="audio/mp3")

            if role == "assistant":
                try:
                    tr = GoogleTranslator(source="auto", target=conv.native_language).translate(msg["text"])
                    st.markdown(f'<div class="chat-aux">{tr}</div>', unsafe_allow_html=True)
                except: 
                    pass

                if show_exp and msg.get("explanation"):
                    exp = msg["explanation"]

                    # Force EXACTLY ONE sentence
                    exp = re.split(r"(?<=[.!?])\s+", exp)[0].strip()

                    # Remove any meta nonsense ("version:", "meaning:", "this sentence", etc)
                    exp = re.sub(r"(?i)(english version|the meaning|this sentence|the german sentence).*", "", exp).strip()

                    if exp:
                        st.markdown(f'<div class="chat-aux">{exp}</div>', unsafe_allow_html=True)

        # scroll
        st.markdown("""
            <script>
            setTimeout(() => {
                let w = window.parent.document.getElementsByClassName('chat-window')[0];
                if (w) w.scrollTop = w.scrollHeight;
            }, 200);
            </script>
        """, unsafe_allow_html=True)

        # -------------------------------
        # AUDIO UPLOAD / RECORDING
        # -------------------------------
        st.markdown('<div class="sticky-input-bar">', unsafe_allow_html=True)

        audio_file = st.file_uploader(
            "🎤 Upload or record an audio message",
            type=["wav", "mp3", "m4a", "ogg"],
            key=f"chat_audio_{st.session_state['recorder_key']}",
        )

        # ------------------------------------------
        # STATE: idle → file → transcribe
        # ------------------------------------------
        if st.session_state["speech_state"] == "idle":
            if audio_file is not None:
                raw = audio_file.read()
                try:
                    seg = AudioSegment.from_file(BytesIO(raw))
                    seg = seg.set_frame_rate(16000).set_channels(1)
                    with st.spinner("Transcribing…"):
                        txt, lang, conf = conv.transcribe(seg, spoken_lang=conv.target_language)

                    st.session_state["pending_transcript"] = txt.strip()
                    st.session_state["speech_state"] = "pending_speech"
                    st.session_state["recorder_key"] += 1
                    st.experimental_rerun()
                except Exception as e:
                    st.error(f"Audio decode/transcription error: {e}")

        # ------------------------------------------
        # STATE: pending_speech → confirm
        # ------------------------------------------
        if st.session_state["speech_state"] == "pending_speech":

            st.write("### Confirm your spoken message:")
            st.info(st.session_state["pending_transcript"])

            c1, c2 = st.columns([1,1])
            with c1:
                if st.button("Send message", key="send_pending"):
                    txt = st.session_state["pending_transcript"]

                    with st.spinner("Partner is responding…"):
                        handle_user_message(username, txt)

                    # cleanup
                    st.session_state["speech_state"] = "idle"
                    st.session_state["pending_transcript"] = ""
                    st.session_state["recorder_key"] += 1
                    st.experimental_rerun()

            with c2:
                if st.button("Discard", key="discard_pending"):
                    st.session_state["speech_state"] = "idle"
                    st.session_state["pending_transcript"] = ""
                    st.session_state["recorder_key"] += 1
                    st.experimental_rerun()

        # -------------------------------
        # TYPED TEXT INPUT
        # -------------------------------
        typed = st.text_input("Type your message:", key="typed_input")
        if typed.strip() and st.button("Send typed message"):
            handle_user_message(username, typed.strip())
            st.session_state["typed_input"] = ""
            st.experimental_rerun()

        st.markdown("</div>", unsafe_allow_html=True)

        # ======================================================
        # RIGHT: SAVED CONVERSATIONS
        # ======================================================
        with col_saved:
            from pathlib import Path
            import json

            st.markdown("### Saved Conversations")

            default_name = datetime.utcnow().strftime("Session %Y-%m-%d %H:%M")
            name_box = st.text_input("Name conversation", value=default_name)

            if st.button("Save conversation"):
                if not st.session_state["chat_history"]:
                    st.warning("Nothing to save.")
                else:
                    safe = re.sub(r"[^0-9A-Za-z_-]", "_", name_box)

                    path = save_current_conversation(username, safe)
                    st.success(f"Saved as {path.name}")

            saved_dir = get_user_dir(username) / "chats" / "saved"
            saved_dir.mkdir(parents=True, exist_ok=True)

            files = sorted(saved_dir.glob("*.json"), key=lambda p: p.stat().st_mtime, reverse=True)

            for f in files:
                data = json.loads(f.read_text())
                sess_name = data.get("name", f.stem)
                msgs = data.get("messages", [])

                with st.expander(f"{sess_name} ({len(msgs)} msgs)"):
                    deck_name = st.text_input(f"Deck name for {sess_name}", value=f"deck_{f.stem}")

                    if st.button(f"Export {f.stem}", key=f"export_{f.stem}"):
                        body = "\n".join(m["text"] for m in msgs if m["role"] == "assistant")
                        deck_path = generate_flashcards_from_text(
                            username=username,
                            text=body,
                            deck_name=deck_name,
                            target_lang=prefs["native_language"],
                            tags=["conversation"],
                        )
                        st.success(f"Deck exported: {deck_path.name}")

                    if st.button(f"Delete {f.stem}", key=f"delete_{f.stem}"):
                        f.unlink()
                        st.experimental_rerun()


###############################################################
# OCR TAB
###############################################################

def ocr_tab(username: str):
    st.header("OCR → Flashcards")

    imgs = st.file_uploader("Upload images", ["png", "jpg", "jpeg"], accept_multiple_files=True)
    tgt = st.selectbox("Translate to", ["en", "de", "ja", "zh-cn", "es"])
    deck_name = st.text_input("Deck name", "ocr_vocab")

    if st.button("Create Deck from OCR"):
        if not imgs:
            st.warning("Upload at least one image.")
            return

        with st.spinner("Running OCR…"):
            results = ocr_and_translate_batch([f.read() for f in imgs], target_lang=tgt)

        deck_path = generate_flashcards_from_ocr_results(
            username=username,
            ocr_results=results,
            deck_name=deck_name,
            target_lang=tgt,
            tags=["ocr"],
        )
        st.success(f"Deck saved: {deck_path}")


###############################################################
# FLASHCARDS TAB
###############################################################

def flashcards_tab(username: str):
    import pandas as pd
    import re

    # ---------------------------------------------------------
    # Helpers
    # ---------------------------------------------------------

    def normalize(s: str) -> str:
        """lowercase + strip non-alphanumerics for loose grading."""
        s = s.lower()
        s = re.sub(r"[^a-z0-9]+", "", s)
        return s

    def card_front_html(text: str) -> str:
        return f"""
        <div style="
            background:#1a73e8;
            color:white;
            border-radius:18px;
            padding:50px;
            font-size:2.2rem;
            text-align:center;
            width:70%;
            margin-left:auto;
            margin-right:auto;
            box-shadow:0 4px 12px rgba(0,0,0,0.3);
            ">
            {text}
        </div>
        """

    def card_back_html(front: str, back: str) -> str:
        return f"""
        <div style="margin-bottom:20px;">
            <div style="
                background:#1a73e8;
                color:white;
                border-radius:18px;
                padding:35px;
                font-size:1.8rem;
                text-align:center;
                width:70%;
                margin-left:auto;
                margin-right:auto;
                box-shadow:0 4px 12px rgba(0,0,0,0.25);
            ">{front}</div>
            <div style="
                background:#2b2b2b;
                color:#f5f5f5;
                border-radius:18px;
                padding:40px;
                margin-top:18px;
                font-size:2rem;
                text-align:center;
                width:70%;
                margin-left:auto;
                margin-right:auto;
                box-shadow:0 4px 12px rgba(0,0,0,0.25);
            ">{back}</div>
        </div>
        """

    # ---------------------------------------------------------
    # Load deck
    # ---------------------------------------------------------

    st.header("Flashcards")

    decks = list_user_decks(username)
    if not decks:
        st.info("No decks available yet.")
        return

    deck_name = st.selectbox("Select deck", sorted(decks.keys()))
    deck_path = decks[deck_name]
    deck = load_deck(deck_path)
    cards = deck.get("cards", [])
    tags = deck.get("tags", [])

    if not cards:
        st.warning("Deck is empty.")
        return

    st.write(f"Total cards: **{len(cards)}**")
    if tags:
        st.caption("Tags: " + ", ".join(tags))

    # Delete deck button
    if st.button("Delete deck"):
        deck_path.unlink()
        st.experimental_rerun()

    # ---------------------------------------------------------
    # Session state setup
    # ---------------------------------------------------------

    key = f"fc_{deck_name}_"
    ss = st.session_state

    if key + "init" not in ss:
        ss[key + "mode"] = "Study"
        ss[key + "idx"] = 0
        ss[key + "show_back"] = False

        # test state
        ss[key + "test_active"] = False
        ss[key + "test_order"] = []
        ss[key + "test_pos"] = 0
        ss[key + "test_results"] = []

        ss[key + "init"] = True

    conv = get_conv_manager()

    mode = st.radio("Mode", ["Study", "Test"], horizontal=True, key=key + "mode")

    st.markdown("---")

    # =======================================================
    # CENTER PANEL
    # =======================================================
    with st.container():

        # ---------------------------------------------------
        # STUDY MODE
        # ---------------------------------------------------
        if mode == "Study":
            idx = ss[key + "idx"] % len(cards)
            card = cards[idx]
            show_back = ss[key + "show_back"]

            st.markdown("### Study Mode")
            st.markdown("---")

            # CARD DISPLAY
            if not show_back:
                st.markdown(card_front_html(card["front"]), unsafe_allow_html=True)
                if st.button("🔊 Pronounce", key=key + f"tts_front_{idx}"):
                    audio = conv.text_to_speech(card["front"])
                    if audio:
                        st.audio(audio, format="audio/mp3")
            else:
                st.markdown(card_back_html(card["front"], card["back"]), unsafe_allow_html=True)
                if st.button("🔊 Pronounce", key=key + f"tts_back_{idx}"):
                    audio = conv.text_to_speech(card["back"])
                    if audio:
                        st.audio(audio, format="audio/mp3")

            # FLIPBOOK CONTROLS
            st.markdown("---")
            c1, c2, c3 = st.columns(3)
            with c1:
                if st.button("Flip", key=key + "flip"):
                    ss[key + "show_back"] = not show_back
                    st.experimental_rerun()
            with c2:
                if st.button("Shuffle deck", key=key + "shuf"):
                    random.shuffle(cards)
                    deck["cards"] = cards
                    save_deck(deck_path, deck)
                    ss[key + "idx"] = 0
                    ss[key + "show_back"] = False
                    st.experimental_rerun()
            with c3:
                if st.button("Next →", key=key + "next"):
                    ss[key + "idx"] = (idx + 1) % len(cards)
                    ss[key + "show_back"] = False
                    st.experimental_rerun()

            # DIFFICULTY GRADING (centered)
            st.markdown("### Rate this card")
            cA, cB, cC, cD, cE = st.columns(5)

            def _choose_next_card_index(cards_list):
                # Simple heuristic: prefer lower-score / less-reviewed cards
                scored = []
                for i, c in enumerate(cards_list):
                    score = c.get("score", 0)
                    reviews = c.get("reviews", 0)
                    priority = score - 0.3 * reviews
                    scored.append((priority, i))
                scored.sort(key=lambda x: x[0])
                return scored[0][1] if scored else 0

            def apply_grade(delta):
                card["score"] = max(0, card.get("score", 0) + delta)
                card["reviews"] = card.get("reviews", 0) + 1
                save_deck(deck_path, deck)
                ss[key + "idx"] = _choose_next_card_index(cards)
                ss[key + "show_back"] = False
                st.experimental_rerun()

            with cA:
                if st.button("🔥 Very Difficult", key=key+"g_vd"):
                    apply_grade(-2)
            with cB:
                if st.button("😣 Hard", key=key+"g_h"):
                    apply_grade(-1)
            with cC:
                if st.button("😐 Neutral", key=key+"g_n"):
                    apply_grade(0)
            with cD:
                if st.button("🙂 Easy", key=key+"g_e"):
                    apply_grade(1)
            with cE:
                if st.button("🏆 Mastered", key=key+"g_m"):
                    apply_grade(3)

        # ---------------------------------------------------
        # TEST MODE
        # ---------------------------------------------------
        else:
            # Initial test setup
            if not ss[key + "test_active"]:
                st.markdown("### Test Setup")
                num_q = st.slider("Number of questions", 3, min(20, len(cards)), min(5, len(cards)), key=key+"nq")

                if st.button("Start Test", key=key+"begin"):
                    order = list(range(len(cards)))
                    random.shuffle(order)
                    order = order[:num_q]

                    ss[key + "test_active"] = True
                    ss[key + "test_order"] = order
                    ss[key + "test_pos"] = 0
                    ss[key + "test_results"] = []
                    st.experimental_rerun()

            else:
                order = ss[key + "test_order"]
                pos = ss[key + "test_pos"]
                results = ss[key + "test_results"]

                # Test Complete
                if pos >= len(order):
                    correct = sum(r["correct"] for r in results)
                    st.markdown(f"### Test Complete — Score: {correct}/{len(results)} ({correct/len(results)*100:.1f}%)")
                    st.markdown("---")

                    for i, r in enumerate(results, 1):
                        emoji = "✅" if r["correct"] else "❌"
                        st.write(f"**{i}.** {r['front']} → expected **{r['back']}**, you answered *{r['user_answer']}* {emoji}")

                    if st.button("Restart Test", key=key+"restart"):
                        ss[key + "test_active"] = False
                        ss[key + "test_pos"] = 0
                        ss[key + "test_results"] = []
                        ss[key + "test_order"] = []
                        st.experimental_rerun()
                    return

                # Current question
                cid = order[pos]
                card = cards[cid]

                st.progress(pos / len(order))
                st.caption(f"Question {pos+1} / {len(order)}")

                st.markdown(card_front_html(card["front"]), unsafe_allow_html=True)

                # TTS
                if st.button("🔊 Pronounce", key=key+f"tts_test_{pos}"):
                    audio = conv.text_to_speech(card["front"])
                    if audio:
                        st.audio(audio, format="audio/mp3")

                user_answer = st.text_input("Your answer:", key=key+f"ans_{pos}")

                if st.button("Submit Answer", key=key+f"submit_{pos}"):
                    ua = user_answer.strip()
                    correct = normalize(ua) == normalize(card["back"])

                    # Flash feedback
                    if correct:
                        st.success("Correct!")
                    else:
                        st.error(f"Incorrect — expected: {card['back']}")

                    results.append({
                        "front": card["front"],
                        "back": card["back"],
                        "user_answer": ua,
                        "correct": correct,
                    })
                    ss[key + "test_results"] = results
                    ss[key + "test_pos"] = pos + 1
                    st.experimental_rerun()

    # =======================================================
    # DECK AT A GLANCE (FULL WIDTH)
    # =======================================================
    st.markdown("---")
    st.subheader("Deck at a glance")

    df_rows = []
    for i, c in enumerate(cards, start=1):
        df_rows.append({
            "#": i,
            "Front": c.get("front", ""),
            "Back": c.get("back", ""),
            "Score": c.get("score", 0),
            "Reviews": c.get("reviews", 0),
        })

    st.dataframe(pd.DataFrame(df_rows), height=500, use_container_width=True)


###############################################################
# QUIZ TAB
###############################################################

def quiz_tab(username: str):
    st.header("Quiz")
    ensure_default_decks(username)

    user_dir = get_user_dir(username)
    quiz_dir = user_dir / "quizzes"
    quiz_dir.mkdir(exist_ok=True)

    decks = list_user_decks(username)
    if not decks:
        st.info("No decks.")
        return

    selected = st.multiselect("Use decks", sorted(decks.keys()))
    if not selected:
        return

    num_q = st.slider("Questions", 3, 20, 6)

    if st.button("Generate quiz"):
        pool = []
        for name in selected:
            pool.extend(load_deck(decks[name])["cards"])

        questions = []
        for _ in range(num_q):
            c = random.choice(pool)
            qtype = random.choice(["mc", "fill"])
            if qtype == "mc":
                others = random.sample(pool, min(3, len(pool) - 1))
                opts = [c["back"]] + [x["back"] for x in others]
                random.shuffle(opts)
                questions.append(
                    {"type": "mc", "prompt": c["front"], "options": opts, "answer": c["back"]}
                )
            else:
                questions.append(
                    {"type": "fill", "prompt": c["front"], "answer": c["back"]}
                )

        qid = datetime.utcnow().strftime("quiz_%Y%m%d_%H%M%S")
        quiz = {"id": qid, "questions": questions}
        (quiz_dir / f"{qid}.json").write_text(json.dumps(quiz, indent=2))

        st.session_state["quiz"] = quiz
        st.session_state["quiz_idx"] = 0
        st.session_state["quiz_answers"] = {}
        st.success("Quiz created!")

    if "quiz" not in st.session_state:
        return

    quiz = st.session_state["quiz"]
    qs = quiz["questions"]
    idx = st.session_state["quiz_idx"]

    if idx >= len(qs):
        correct = sum(1 for v in st.session_state["quiz_answers"].values() if v["correct"])
        st.success(f"Score: {correct}/{len(qs)}")
        if st.button("New quiz"):
            del st.session_state["quiz"]
            del st.session_state["quiz_idx"]
            del st.session_state["quiz_answers"]
        return

    q = qs[idx]
    st.subheader(f"Question {idx+1}/{len(qs)}")
    st.markdown(f"**{q['prompt']}**")

    if q["type"] == "mc":
        choice = st.radio("Choose:", q["options"], key=f"mc_{idx}")
        if st.button("Submit", key=f"sub_{idx}"):
            st.session_state["quiz_answers"][idx] = {
                "given": choice,
                "correct": choice == q["answer"],
            }
            st.session_state["quiz_idx"] += 1
            st.experimental_rerun()
    else:
        ans = st.text_input("Your answer", key=f"fill_{idx}")
        if st.button("Submit", key=f"sub_{idx}"):
            st.session_state["quiz_answers"][idx] = {
                "given": ans,
                "correct": ans.strip().lower() == q["answer"].lower(),
            }
            st.session_state["quiz_idx"] += 1
            st.experimental_rerun()


###############################################################
# MAIN
###############################################################

def main():
    # ---------- AUTH ----------
    if "user" not in st.session_state:
        login_view()
        return

    username = st.session_state["user"]
    st.sidebar.write(f"Logged in as **{username}**")

    if st.sidebar.button("Log out"):
        st.session_state.clear()
        st.experimental_rerun()

    # ---------- LOAD MODELS + PREFS ----------
    preload_models()
    sidebar_settings(username)
    ensure_default_decks(username)

    # ---------- TAB PERSISTENCE ----------
    if "active_tab" not in st.session_state:
        st.session_state["active_tab"] = 0

    tab_labels = ["Dashboard", "Conversation", "OCR", "Flashcards", "Quiz", "Settings"]
    tabs = st.tabs(tab_labels)

    tab_dash, tab_conv, tab_ocr, tab_flash, tab_quiz, tab_settings = tabs

    # restore active tab (required so Streamlit shows correct tab on rerun)
    _ = tabs[st.session_state["active_tab"]]

    with tab_dash:
        st.session_state["active_tab"] = 0
        dashboard_tab(username)

    with tab_conv:
        st.session_state["active_tab"] = 1
        conversation_tab(username)

    with tab_ocr:
        st.session_state["active_tab"] = 2
        ocr_tab(username)

    with tab_flash:
        st.session_state["active_tab"] = 3
        flashcards_tab(username)

    with tab_quiz:
        st.session_state["active_tab"] = 4
        quiz_tab(username)

    with tab_settings:
        st.session_state["active_tab"] = 5
        settings_tab(username)


if __name__ == "__main__":
    main()