src/app/main_app.py

###############################################################
# main_app.py — Agentic Language Partner UI (Streamlit)
###############################################################
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.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.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.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.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.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.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.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
###############################################################

# Lightweight TTS function with caching (no heavy model loading)
@st.cache_data(ttl=3600, show_spinner=False)
def _cached_tts(text: str, lang: str = "en") -> bytes:
    """Generate TTS audio using gTTS with caching."""
    import io
    from gtts import gTTS

    # Language mapping
    lang_map = {
        "english": "en", "spanish": "es", "german": "de",
        "russian": "ru", "japanese": "ja", "chinese": "zh-cn",
        "korean": "ko", "french": "fr",
    }
    tts_lang = lang_map.get(lang.lower(), lang)

    try:
        tts = gTTS(text=text, lang=tts_lang)
        buf = io.BytesIO()
        tts.write_to_fp(buf)
        return buf.getvalue()
    except Exception:
        return None


def flashcards_tab(username: str):
    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.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

    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("---")

            # Get target language for TTS
            prefs = st.session_state.get("prefs", {})
            tts_lang = prefs.get("target_language", "english")

            # CARD DISPLAY
            if not show_back:
                st.markdown(card_front_html(card["front"]), unsafe_allow_html=True)
                # Pronounce button for front
                if st.button("🔊 Pronounce", key=key + f"tts_front_{idx}"):
                    audio = _cached_tts(card["front"], tts_lang)
                    if audio:
                        st.audio(audio, format="audio/mp3")
                    else:
                        st.warning("TTS not available for this text.")
            else:
                st.markdown(card_back_html(card["front"], card["back"]), unsafe_allow_html=True)
                # Pronounce button for back (use native language)
                native_lang = prefs.get("native_language", "english")
                if st.button("🔊 Pronounce", key=key + f"tts_back_{idx}"):
                    audio = _cached_tts(card["back"], native_lang)
                    if audio:
                        st.audio(audio, format="audio/mp3")
                    else:
                        st.warning("TTS not available for this text.")

            # FLIPBOOK CONTROLS
            st.markdown("---")

            # Use number input for card navigation (no rerun needed)
            new_idx = st.number_input(
                f"Card ({idx + 1} of {len(cards)})",
                min_value=1,
                max_value=len(cards),
                value=idx + 1,
                key=key + "card_num"
            ) - 1

            if new_idx != idx:
                ss[key + "idx"] = new_idx
                ss[key + "show_back"] = False

            c1, c2 = st.columns(2)
            with c1:
                if st.button("Flip Card", key=key + "flip"):
                    ss[key + "show_back"] = not show_back
            with c2:
                if st.button("Next →", key=key + "next"):
                    ss[key + "idx"] = (idx + 1) % len(cards)
                    ss[key + "show_back"] = False

        # ---------------------------------------------------
        # 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"] = []

            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"] = []

                else:
                    # 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)

                    # Pronounce button for test
                    prefs = st.session_state.get("prefs", {})
                    tts_lang = prefs.get("target_language", "english")
                    if st.button("🔊 Pronounce", key=key+f"tts_test_{pos}"):
                        audio = _cached_tts(card["front"], tts_lang)
                        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

    # =======================================================
    # DECK AT A GLANCE (FULL WIDTH)
    # =======================================================
    st.markdown("---")
    with st.expander("Deck at a glance", expanded=False):
        if cards:
            for i, c in enumerate(cards, start=1):
                front = c.get("front", "")
                back = c.get("back", "")
                score = c.get("score", 0)
                reviews = c.get("reviews", 0)
                st.markdown(f"**{i}.** {front} → {back} *(Score: {score}, Reviews: {reviews})*")
        else:
            st.info("No cards in this deck.")


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

def quiz_tab(username: str):
    # Custom CSS for Quiz UI
    st.markdown("""
    <style>
    .quiz-card {
        background: linear-gradient(135deg, #667eea 0%, #764ba2 100%);
        border-radius: 20px;
        padding: 30px;
        margin: 20px 0;
        box-shadow: 0 10px 40px rgba(0,0,0,0.3);
    }
    .quiz-question {
        background: rgba(255,255,255,0.15);
        border-radius: 15px;
        padding: 25px;
        margin: 15px 0;
        backdrop-filter: blur(10px);
    }
    .quiz-prompt {
        font-size: 1.8rem;
        font-weight: bold;
        color: white;
        text-align: center;
        margin-bottom: 20px;
    }
    .quiz-progress {
        background: rgba(255,255,255,0.2);
        border-radius: 10px;
        padding: 10px 20px;
        margin-bottom: 20px;
    }
    .result-correct {
        background: linear-gradient(135deg, #11998e 0%, #38ef7d 100%);
        border-radius: 15px;
        padding: 20px;
        margin: 10px 0;
        color: white;
    }
    .result-wrong {
        background: linear-gradient(135deg, #eb3349 0%, #f45c43 100%);
        border-radius: 15px;
        padding: 20px;
        margin: 10px 0;
        color: white;
    }
    .score-card {
        background: linear-gradient(135deg, #f093fb 0%, #f5576c 100%);
        border-radius: 25px;
        padding: 40px;
        text-align: center;
        margin: 20px 0;
        box-shadow: 0 15px 50px rgba(0,0,0,0.3);
    }
    .score-number {
        font-size: 4rem;
        font-weight: bold;
        color: white;
    }
    .feedback-box {
        border-radius: 12px;
        padding: 15px 20px;
        margin: 15px 0;
        font-size: 1.1rem;
    }
    .feedback-correct {
        background: #d4edda;
        border-left: 5px solid #28a745;
        color: #155724;
    }
    .feedback-wrong {
        background: #f8d7da;
        border-left: 5px solid #dc3545;
        color: #721c24;
    }
    </style>
    """, unsafe_allow_html=True)

    st.markdown("## 🎯 Quiz Mode")
    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 available. Create some flashcards first!")
        return

    # Quiz state keys
    qkey = "quiz_state_"

    # Initialize state
    if qkey + "active" not in st.session_state:
        st.session_state[qkey + "active"] = False
        st.session_state[qkey + "questions"] = []
        st.session_state[qkey + "current"] = 0
        st.session_state[qkey + "answers"] = {}
        st.session_state[qkey + "show_feedback"] = False
        st.session_state[qkey + "last_correct"] = None

    # =====================================================
    # QUIZ SETUP (not active)
    # =====================================================
    if not st.session_state[qkey + "active"]:
        st.markdown("### Select Decks & Start")

        col1, col2 = st.columns([2, 1])

        with col1:
            selected = st.multiselect(
                "Choose decks to quiz from:",
                sorted(decks.keys()),
                default=list(decks.keys())[:1] if decks else []
            )

        with col2:
            num_q = st.slider("Number of questions", 3, 15, 5)

        if selected and st.button("🚀 Start Quiz", use_container_width=True):
            # Build question pool
            pool = []
            for name in selected:
                deck_cards = load_deck(decks[name])["cards"]
                for card in deck_cards:
                    pool.append({"front": card.get("front", ""), "back": card.get("back", "")})

            if len(pool) < 2:
                st.warning("Need at least 2 cards to create a quiz.")
                return

            # Generate questions
            questions = []
            used_prompts = set()

            for _ in range(min(num_q, len(pool))):
                # Pick a card not yet used
                available = [c for c in pool if c["front"] not in used_prompts]
                if not available:
                    break

                c = random.choice(available)
                used_prompts.add(c["front"])

                # Randomly choose question type
                qtype = random.choice(["mc", "mc", "fill"])  # More MC questions

                if qtype == "mc" and len(pool) >= 4:
                    # Multiple choice - get wrong options
                    wrong_pool = [x["back"] for x in pool if x["front"] != c["front"]]
                    wrong_opts = random.sample(wrong_pool, min(3, len(wrong_pool)))
                    opts = [c["back"]] + wrong_opts
                    random.shuffle(opts)
                    questions.append({
                        "type": "mc",
                        "prompt": c["front"],
                        "options": opts,
                        "answer": c["back"]
                    })
                else:
                    # Fill in the blank
                    questions.append({
                        "type": "fill",
                        "prompt": c["front"],
                        "answer": c["back"]
                    })

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

            # Set state
            st.session_state[qkey + "active"] = True
            st.session_state[qkey + "questions"] = questions
            st.session_state[qkey + "current"] = 0
            st.session_state[qkey + "answers"] = {}
            st.session_state[qkey + "show_feedback"] = False

        return

    # =====================================================
    # QUIZ IN PROGRESS
    # =====================================================
    questions = st.session_state[qkey + "questions"]
    current = st.session_state[qkey + "current"]
    answers = st.session_state[qkey + "answers"]
    show_feedback = st.session_state[qkey + "show_feedback"]

    # =====================================================
    # QUIZ COMPLETE - Show Results
    # =====================================================
    if current >= len(questions):
        correct_count = sum(1 for a in answers.values() if a["correct"])
        total = len(questions)
        percentage = (correct_count / total) * 100

        # Score display
        if percentage >= 80:
            grade_emoji = "🏆"
            grade_text = "Excellent!"
            grade_color = "#28a745"
        elif percentage >= 60:
            grade_emoji = "👍"
            grade_text = "Good job!"
            grade_color = "#17a2b8"
        elif percentage >= 40:
            grade_emoji = "📚"
            grade_text = "Keep practicing!"
            grade_color = "#ffc107"
        else:
            grade_emoji = "💪"
            grade_text = "Don't give up!"
            grade_color = "#dc3545"

        st.markdown(f"""
        <div class="score-card">
            <div style="font-size: 3rem;">{grade_emoji}</div>
            <div class="score-number">{correct_count}/{total}</div>
            <div style="font-size: 1.5rem; color: white; margin-top: 10px;">{percentage:.0f}% - {grade_text}</div>
        </div>
        """, unsafe_allow_html=True)

        # Detailed results
        st.markdown("### 📋 Review All Answers")

        for i, q in enumerate(questions):
            ans_data = answers.get(i, {})
            is_correct = ans_data.get("correct", False)
            user_ans = ans_data.get("given", "No answer")

            if is_correct:
                st.markdown(f"""
                <div class="feedback-box feedback-correct">
                    <strong>Q{i+1}:</strong> {q['prompt']}<br>
                    ✅ Your answer: <strong>{user_ans}</strong>
                </div>
                """, unsafe_allow_html=True)
            else:
                st.markdown(f"""
                <div class="feedback-box feedback-wrong">
                    <strong>Q{i+1}:</strong> {q['prompt']}<br>
                    ❌ Your answer: <strong>{user_ans}</strong><br>
                    ✅ Correct answer: <strong>{q['answer']}</strong>
                </div>
                """, unsafe_allow_html=True)

        st.markdown("---")

        if st.button("🔄 Start New Quiz", use_container_width=True):
            st.session_state[qkey + "active"] = False
            st.session_state[qkey + "questions"] = []
            st.session_state[qkey + "current"] = 0
            st.session_state[qkey + "answers"] = {}
            st.session_state[qkey + "show_feedback"] = False

        return

    # =====================================================
    # CURRENT QUESTION
    # =====================================================
    q = questions[current]

    # Progress bar
    progress = (current) / len(questions)
    st.progress(progress)
    st.markdown(f"**Question {current + 1} of {len(questions)}**")

    # Question card
    st.markdown(f"""
    <div class="quiz-question">
        <div class="quiz-prompt">{q['prompt']}</div>
    </div>
    """, unsafe_allow_html=True)

    # Show feedback from previous submission
    if show_feedback and st.session_state[qkey + "last_correct"] is not None:
        last_ans = answers.get(current, {})
        if st.session_state[qkey + "last_correct"]:
            st.markdown("""
            <div class="feedback-box feedback-correct">
                ✅ <strong>Correct!</strong> Well done!
            </div>
            """, unsafe_allow_html=True)
        else:
            st.markdown(f"""
            <div class="feedback-box feedback-wrong">
                ❌ <strong>Incorrect!</strong><br>
                The correct answer is: <strong>{q['answer']}</strong>
            </div>
            """, unsafe_allow_html=True)

        # Next button
        if st.button("Next Question →", use_container_width=True):
            st.session_state[qkey + "current"] = current + 1
            st.session_state[qkey + "show_feedback"] = False
            st.session_state[qkey + "last_correct"] = None

    else:
        # Answer input
        if q["type"] == "mc":
            st.markdown("**Choose the correct answer:**")
            choice = st.radio(
                "Options:",
                q["options"],
                key=f"quiz_mc_{current}",
                label_visibility="collapsed"
            )

            if st.button("✓ Submit Answer", use_container_width=True, key=f"quiz_sub_{current}"):
                is_correct = (choice == q["answer"])
                st.session_state[qkey + "answers"][current] = {
                    "given": choice,
                    "correct": is_correct,
                    "expected": q["answer"]
                }
                st.session_state[qkey + "show_feedback"] = True
                st.session_state[qkey + "last_correct"] = is_correct

        else:  # fill in blank
            st.markdown("**Type your answer:**")
            user_input = st.text_input(
                "Your answer:",
                key=f"quiz_fill_{current}",
                label_visibility="collapsed",
                placeholder="Type the translation..."
            )

            if st.button("✓ Submit Answer", use_container_width=True, key=f"quiz_sub_{current}"):
                is_correct = (user_input.strip().lower() == q["answer"].strip().lower())
                st.session_state[qkey + "answers"][current] = {
                    "given": user_input,
                    "correct": is_correct,
                    "expected": q["answer"]
                }
                st.session_state[qkey + "show_feedback"] = True
                st.session_state[qkey + "last_correct"] = is_correct

    # Quit button
    st.markdown("---")
    if st.button("🚪 Quit Quiz", use_container_width=True):
        st.session_state[qkey + "active"] = False
        st.session_state[qkey + "questions"] = []
        st.session_state[qkey + "current"] = 0
        st.session_state[qkey + "answers"] = {}
        st.session_state[qkey + "show_feedback"] = False


###############################################################
# 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.rerun()

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

    # ---------- TABS ----------
    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

    with tab_dash:
        dashboard_tab(username)

    with tab_conv:
        conversation_tab(username)

    with tab_ocr:
        ocr_tab(username)

    with tab_flash:
        flashcards_tab(username)

    with tab_quiz:
        quiz_tab(username)

    with tab_settings:
        settings_tab(username)


if __name__ == "__main__":
    main()