# ============================================================
# 📦 IMPORTS
# ============================================================

import re
import os
import json
import faiss
import gradio as gr
import numpy as np
import pandas as pd
import requests

from sentence_transformers import SentenceTransformer
from groq import Groq


# ============================================================
# 🔑 GROQ API KEY
# ============================================================

GROQ_API_KEY = os.getenv('GRAPI')

client = Groq(api_key=GROQ_API_KEY)


# ============================================================
# 📘 GOOGLE DOC LOADER
# ============================================================

DOC_ID = "1utErkC3Xa8hhiQul7tzil9SjRydGEKGEBpVgaCp6qXM"

def load_google_doc():

    url = f"https://docs.google.com/document/d/{DOC_ID}/export?format=txt"

    text = requests.get(url).text

    return text


raw_text = load_google_doc()


# ============================================================
# 📚 SPLIT CHAPTERS
# ============================================================

def split_units(text):

    units = {}

    splits = re.split(r"(UNIT\s+\d+.*?)\n", text, flags=re.I)

    for i in range(1, len(splits), 2):

        title = splits[i].strip()

        content = splits[i+1]

        units[title] = content

    return units


units_data = split_units(raw_text)

print("✅ Loaded Units:")
print(list(units_data.keys()))


# ============================================================
# 🔍 EMBEDDING MODEL
# ============================================================

embedding_model = SentenceTransformer(
    "all-MiniLM-L6-v2"
)


# ============================================================
# ✂️ CHUNKING
# ============================================================

def chunk_text(text, chunk_size=150):

    words = text.split()

    chunks = []

    for i in range(0, len(words), chunk_size):

        chunk = " ".join(words[i:i+chunk_size])

        chunks.append(chunk)

    return chunks


all_chunks = []
chunk_unit_map = []

for unit, text in units_data.items():

    chunks = chunk_text(text)

    all_chunks.extend(chunks)

    chunk_unit_map.extend([unit]*len(chunks))


# ============================================================
# 🧠 CREATE FAISS VECTOR DB
# ============================================================

embeddings = embedding_model.encode(all_chunks)

if len(embeddings.shape) == 1:

    embeddings = np.expand_dims(embeddings, axis=0)

dimension = embeddings.shape[1]

index = faiss.IndexFlatL2(dimension)

index.add(np.array(embeddings).astype("float32"))

print("✅ FAISS INDEX READY")


# ============================================================
# 🔍 RETRIEVE CONTEXT
# ============================================================

def retrieve_context(question, unit):

    q_embedding = embedding_model.encode([question])

    D, I = index.search(
        np.array(q_embedding).astype("float32"),
        5
    )

    retrieved = []

    for idx in I[0]:

        if chunk_unit_map[idx] == unit:

            retrieved.append(all_chunks[idx])

    return "\n".join(retrieved)


# ============================================================
# 🤖 GROQ HELPER
# ============================================================

def groq_call(prompt):

    completion = client.chat.completions.create(

        model="llama-3.3-70b-versatile",

        messages=[
            {
                "role": "user",
                "content": prompt
            }
        ],

        temperature=0.5
    )

    return completion.choices[0].message.content


# ============================================================
# 📚 VOCAB EXTRACTION
# ============================================================

def extract_vocabulary(unit):

    context = units_data[unit][:7000]

    prompt = f"""
You are an FSC English linguistic AI.

Extract ONLY 5 important contextual vocabulary words.

Return STRICT JSON:

[
 {{
   "word":"...",
   "meaning_en":"...",
   "meaning_ur":"...",
   "context_sentence":"..."
 }}
]

TEXT:
{context}
"""

    result = groq_call(prompt)

    try:

        data = json.loads(result)

        return data

    except:

        return []


# ============================================================
# 🧠 METACOGNITIVE ENGINE
# ============================================================

student_analytics = []


def build_mcq(word_data):

    word = word_data["word"]

    meaning = word_data["meaning_en"]

    prompt = f"""
Create one FSC-level MCQ.

Word: {word}
Meaning: {meaning}

Return STRICT JSON:

{{
 "question":"...",
 "options": {{
   "A":"...",
   "B":"...",
   "C":"..."
 }},
 "answer":"A/B/C"
}}
"""

    result = groq_call(prompt)

    try:

        return json.loads(result)

    except:

        return None


# ============================================================
# 🎯 PROCESS ANSWER
# ============================================================

def evaluate_answer(
    word,
    correct_answer,
    selected_answer,
    confidence
):

    correct = selected_answer == correct_answer

    feedback = ""

    autonomy_points = 0

    if correct:

        if confidence <= 2:

            feedback = (
                f"✅ Correct! "
                f"You were unsure, but contextual clues helped you."
            )

        elif confidence >= 4:

            feedback = (
                f"🔥 Excellent! "
                f"You understood the lexical context confidently."
            )

        autonomy_points += 10

    else:

        if confidence >= 4:

            feedback = (
                f"⚠️ High confidence but incorrect answer.\n"
                f"You may need stronger contextual inference."
            )

        else:

            feedback = (
                f"❌ Incorrect.\n"
                f"Try analyzing surrounding contextual clues."
            )

    student_analytics.append({

        "word": word,

        "correct": correct,

        "confidence": confidence
    })

    return feedback, autonomy_points


# ============================================================
# 🌉 BILINGUAL SCAFFOLD
# ============================================================

def scaffold(word_data, level):

    if level == 1:

        return f"""
🧠 English Scaffold:

{word_data['meaning_en']}
"""

    if level == 2:

        return f"""
🧠 Urdu Scaffold:

{word_data['meaning_ur']}
"""

    return ""


# ============================================================
# 📊 ANALYTICS DASHBOARD
# ============================================================

def generate_dashboard():

    if len(student_analytics) == 0:

        return pd.DataFrame()

    df = pd.DataFrame(student_analytics)

    return df


# ============================================================
# 🎨 CUSTOM CSS
# ============================================================

custom_css = """

/* =========================
   GLOBAL
========================= */

body{
    background:#F1F5F9 !important;
    font-family:'Inter',sans-serif;
    color:#111827 !important;
}

/* =========================
   MAIN APP CONTAINER
========================= */

.gradio-container{
    max-width:1200px !important;
    margin:auto !important;

    background:white !important;

    border-radius:24px;

    padding:35px !important;

    box-shadow:
    0px 8px 30px rgba(0,0,0,0.08);

    color:#111827 !important;
}


/* =========================
   TITLES
========================= */

.main-title{
    text-align:center;

    font-size:52px;

    font-weight:800;

    color:#0F172A !important;

    margin-bottom:8px;
}

.subtitle{
    text-align:center;

    font-size:20px;

    color:#475569 !important;

    margin-bottom:35px;
}


/* =========================
   CARDS
========================= */

.card{
    background:#FFFFFF !important;

    border:1px solid #E2E8F0 !important;

    border-radius:20px;

    padding:22px;

    margin-bottom:22px;

    box-shadow:
    0px 2px 12px rgba(0,0,0,0.05);
}


/* =========================
   TEXT COLORS
========================= */

h1,h2,h3,h4,h5,h6{
    color:#0F172A !important;
}

p,span,div,label{
    color:#111827 !important;
}


/* =========================
   TARGET WORD
========================= */

.target-word{
    font-size:40px;

    font-weight:700;

    color:#0284C7 !important;

    margin-bottom:10px;
}

.context-text{
    font-size:18px;

    line-height:1.8;

    color:#1E293B !important;
}


/* =========================
   SECTION LABELS
========================= */

.section-label{
    display:inline-block;

    background:#E0F2FE;

    color:#0369A1 !important;

    padding:8px 14px;

    border-radius:10px;

    font-size:15px;

    font-weight:700;

    margin-bottom:16px;
}


/* =========================
   INPUTS
========================= */

textarea,
input,
select{

    background:#FFFFFF !important;

    color:#111827 !important;

    border:1.5px solid #CBD5E1 !important;

    border-radius:12px !important;

    padding:12px !important;

    font-size:16px !important;
}


/* =========================
   DROPDOWN
========================= */

.gr-dropdown{
    color:#111827 !important;
}


/* =========================
   RADIO BUTTONS
========================= */

.gr-radio label{

    background:#F8FAFC !important;

    border:1px solid #CBD5E1 !important;

    border-radius:12px;

    padding:12px 16px;

    margin-right:10px;

    transition:0.2s;
}

.gr-radio label:hover{
    background:#E0F2FE !important;
}


/* =========================
   BUTTONS
========================= */

button{

    background:#0EA5E9 !important;

    color:white !important;

    border:none !important;

    border-radius:14px !important;

    font-size:17px !important;

    font-weight:700 !important;

    padding:14px 22px !important;

    transition:0.25s ease;
}

button:hover{

    background:#0284C7 !important;

    transform:translateY(-1px);
}


/* =========================
   SLIDER
========================= */

input[type="range"]{
    accent-color:#0EA5E9 !important;
}


/* =========================
   OUTPUT BOXES
========================= */

.output-box{

    background:#F8FAFC !important;

    border:1px solid #E2E8F0 !important;

    border-radius:16px;

    padding:18px;

    color:#111827 !important;
}


/* =========================
   CHATBOT / MARKDOWN
========================= */

.markdown-text{
    color:#111827 !important;
}


/* =========================
   MOBILE RESPONSIVE
========================= */

@media(max-width:768px){

    .main-title{
        font-size:34px;
    }

    .subtitle{
        font-size:16px;
    }

    .target-word{
        font-size:28px;
    }

    .context-text{
        font-size:16px;
    }

    .gradio-container{
        padding:18px !important;
    }
}

"""


# ============================================================
# 🚀 UI FUNCTIONS
# ============================================================

current_vocab = []
current_mcq = None
current_word_data = None


def load_chapter(unit):

    global current_vocab
    global current_word_data
    global current_mcq

    current_vocab = extract_vocabulary(unit)

    if len(current_vocab) == 0:

        return (
            "❌ Could not extract vocabulary.",
            "",
            gr.update(choices=[]),
            ""
        )

    current_word_data = current_vocab[0]

    current_mcq = build_mcq(current_word_data)

    vocab_text = f"""
# 📘 Target Word

### {current_word_data['word']}

### Context Sentence:
{current_word_data['context_sentence']}
"""

    return (

        vocab_text,

        current_mcq["question"],

        gr.update(
            choices=list(
                current_mcq["options"].values()
            )
        ),

        ""
    )


def submit_answer(selected, confidence):

    global current_mcq
    global current_word_data

    options = current_mcq["options"]

    reverse_map = {
        v:k for k,v in options.items()
    }

    selected_letter = reverse_map[selected]

    feedback, points = evaluate_answer(

        current_word_data["word"],

        current_mcq["answer"],

        selected_letter,

        confidence
    )

    return f"""
# 🧠 Feedback

{feedback}

⭐ Autonomy Points: {points}
"""


def show_english_scaffold():

    global current_word_data

    return scaffold(current_word_data, 1)


def show_urdu_scaffold():

    global current_word_data

    return scaffold(current_word_data, 2)


def show_dashboard():

    df = generate_dashboard()

    return df


# ============================================================
# 🌟 GRADIO UI
# ============================================================

with gr.Blocks(
    css=custom_css,
    theme=gr.themes.Soft()
) as demo:

    gr.HTML("""
    <div class='main-title'>
    🚀 LexiMetrica
    </div>

    <div class='subtitle'>
    AI Metacognitive Lexical Enhancement Framework
    </div>
    """)

    with gr.Row():

        chapter_dropdown = gr.Dropdown(

            choices=list(units_data.keys()),

            label="📚 Select Chapter"
        )

        load_btn = gr.Button(
            "Load Chapter",
            variant="primary"
        )

    with gr.Row():

        vocab_output = gr.Markdown()

    with gr.Row():

        mcq_question = gr.Markdown()

    with gr.Row():

        mcq_radio = gr.Radio(
            choices=[],
            label="Choose Meaning"
        )

    confidence_slider = gr.Slider(

        minimum=1,

        maximum=5,

        step=1,

        label="🧠 Confidence Level"
    )

    submit_btn = gr.Button(
        "Submit Answer",
        variant="primary"
    )

    feedback_output = gr.Markdown()

    gr.Markdown("---")

    gr.Markdown("## 🌉 Need a Bridge?")

    with gr.Row():

        english_btn = gr.Button(
            "English Scaffold"
        )

        urdu_btn = gr.Button(
            "Urdu Scaffold"
        )

    scaffold_output = gr.Markdown()

    gr.Markdown("---")

    dashboard_btn = gr.Button(
        "📊 Show Analytics Dashboard"
    )

    dashboard_output = gr.Dataframe()

    # ======================================
    # EVENTS
    # ======================================

    load_btn.click(

        fn=load_chapter,

        inputs=chapter_dropdown,

        outputs=[
            vocab_output,
            mcq_question,
            mcq_radio,
            scaffold_output
        ]
    )

    submit_btn.click(

        fn=submit_answer,

        inputs=[
            mcq_radio,
            confidence_slider
        ],

        outputs=feedback_output
    )

    english_btn.click(

        fn=show_english_scaffold,

        outputs=scaffold_output
    )

    urdu_btn.click(

        fn=show_urdu_scaffold,

        outputs=scaffold_output
    )

    dashboard_btn.click(

        fn=show_dashboard,

        outputs=dashboard_output
    )


# ============================================================
# 🚀 LAUNCH
# ============================================================

demo.launch(debug=True)