import os
import sys
import re
import json
import gradio as gr
from pathlib import Path
from groq import Groq

sys.path.insert(0, ".")

from src.data_loader import load_corpus, load_graph, link_corpus_to_graph, enrich_with_graph_topics
from src.retrieval.cgir_pipeline import CGIRPipeline
from src.retrieval.graph_navigator import mastery_to_bloom
from src.tracing.bkt import BayesianKnowledgeTracing, BKTParams, ObservationRecord

# ── Constants ─────────────────────────────────────────────────────────────────
MAX_QUESTIONS = 15

DEFAULT_BKT_PARAMS = BKTParams(
    p_init    = 0.10,
    p_transit = 0.05,
    p_slip    = 0.15,
    p_guess   = 0.45,
)

BLOOM_EMOJI  = {1:"🟢", 2:"🔵", 3:"🟡", 4:"🟠", 5:"🔴", 6:"🟣"}
BLOOM_LABELS = {
    1:"Mengingat", 2:"Memahami", 3:"Mengaplikasikan",
    4:"Menganalisis", 5:"Mengevaluasi", 6:"Mencipta",
}

# ── Load data + index ─────────────────────────────────────────────────────────
print("Loading corpus dan graph...")
documents     = load_corpus("data/Knowledge_Base_with_ID.csv")
graph         = load_graph("data/knowledge_graph.json")
doc_graph_map = link_corpus_to_graph(documents, graph)
documents     = enrich_with_graph_topics(documents, graph, doc_graph_map)

cgir = CGIRPipeline(graph=graph, doc_to_graph_id=doc_graph_map, bloom_window=1)

if Path("faiss_index/faiss.index").exists():
    print("Loading saved FAISS index...")
    cgir.load("faiss_index")
    for doc_idx, doc in enumerate(cgir.faiss.documents):
        graph_id = doc_graph_map.get(doc["id"])
        if graph_id:
            cgir._graph_to_docs.setdefault(graph_id, []).append(doc_idx)
else:
    print("Building FAISS index...")
    cgir.build(documents)
    Path("faiss_index").mkdir(parents=True, exist_ok=True)
    cgir.save("faiss_index")

# ── Groq judge — continuous 0.0–1.0 ──────────────────────────────────────────
_groq = Groq(api_key=os.environ.get("GROQ_API_KEY", ""))

JUDGE_PROMPT = """\
Kamu adalah penilai jawaban siswa untuk pelajaran IPS/Geografi SMA.
Gunakan konteks berikut sebagai referensi utama:
<konteks>
{context}
</konteks>
Soal:
{question}
Jawaban siswa:
{answer}
Nilai jawaban dalam skala 0.0 hingga 1.0:
  1.00 = benar sempurna dan lengkap
  0.75 = sebagian besar benar, ada yang kurang
  0.50 = setengah benar, inti ada tapi banyak yang hilang
  0.25 = sedikit benar, kebanyakan salah
  0.00 = salah total atau tidak relevan
Berikan penilaian dalam format JSON berikut (tanpa teks lain):
{{"score": 0.0 hingga 1.0, "feedback": "kalimat singkat 1-2 kalimat menjelaskan kenapa benar/salah"}}
"""

def judge_answer(question: str, answer: str, context: str = "") -> dict:
    prompt = JUDGE_PROMPT.format(
        context=context[:2000] if context else "Tidak ada konteks.",
        question=question,
        answer=answer,
    )
    try:
        resp = _groq.chat.completions.create(
            model="llama-3.3-70b-versatile",
            max_tokens=256,
            temperature=0.1,
            messages=[{"role": "user", "content": prompt}],
        )
        raw    = resp.choices[0].message.content.strip()
        raw    = re.sub(r"^```json\s*|```$", "", raw, flags=re.MULTILINE).strip()
        result = json.loads(raw)
        score  = float(result.get("score", 0.0))
        return {
            "score":    min(max(score, 0.0), 1.0),
            "feedback": str(result.get("feedback", "")),
        }
    except Exception as e:
        return {"score": 0.0, "feedback": f"[Error: {e}]"}


# ── Weighted BKT update — continuous score ────────────────────────────────────
def bkt_update_continuous(
    bkt_inst, concept: str, score: float,
    bloom_level: int = 1, question_id: str = ""
) -> float:
    params = bkt_inst.concept_params.get(concept, bkt_inst.default_params)
    p_l    = bkt_inst.get_mastery(concept)

    p_obs_c        = (1 - params.p_slip) * p_l + params.p_guess * (1 - p_l)
    p_l_if_correct = ((1 - params.p_slip) * p_l) / max(p_obs_c, 1e-12)
    p_next_correct = p_l_if_correct + (1 - p_l_if_correct) * params.p_transit

    p_obs_w       = params.p_slip * p_l + (1 - params.p_guess) * (1 - p_l)
    p_l_if_wrong  = (params.p_slip * p_l) / max(p_obs_w, 1e-12)
    p_next_wrong  = p_l_if_wrong + (1 - p_l_if_wrong) * params.p_transit

    p_l_next = score * p_next_correct + (1 - score) * p_next_wrong
    p_l_next = min(max(p_l_next, 0.0), 1.0)

    bkt_inst._mastery[concept] = p_l_next
    bkt_inst.history.append(ObservationRecord(
        concept=concept, correct=score >= 0.5,
        mastery_before=p_l, mastery_after=p_l_next,
        bloom_level=bloom_level, question_id=question_id,
    ))
    return p_l_next


# ── Session helpers ───────────────────────────────────────────────────────────
def create_session(keyword: str) -> dict:
    return {
        "bkt":               BayesianKnowledgeTracing(default_params=DEFAULT_BKT_PARAMS),
        "seen_ids":          set(),
        "consecutive_wrong": 0,
        "keyword":           keyword.strip(),
        "q_num":             0,
        "current_q":         None,
        "active":            True,
        "correct_count":     0,
        "wrong_count":       0,
    }


def get_mastery(session: dict) -> float:
    bkt = session.get("bkt")
    if not bkt:
        return 0.0
    return bkt.get_mastery(session.get("keyword", "").lower().strip())


def get_next_question(session: dict):
    keyword           = session["keyword"]
    mastery           = get_mastery(session)
    consecutive_wrong = session["consecutive_wrong"]

    effective = mastery
    if consecutive_wrong >= 2:
        effective = max(0.0, mastery - 0.15 * consecutive_wrong)

    target_bloom = mastery_to_bloom(effective)
    candidates   = cgir.retrieve_candidates(keyword, mastery=effective, top_k=20)
    unseen       = [c for c in candidates if c.question.get("id") not in session["seen_ids"]]

    if not unseen:
        return None

    # Prioritas 1: exact bloom match
    exact = [c for c in unseen if c.question.get("bloom_level") == target_bloom]
    if exact:
        return exact[0]

    # Prioritas 2: satu level di bawah
    below = [c for c in unseen if c.question.get("bloom_level") == target_bloom - 1]
    if below:
        return below[0]

    return unseen[0]


def add_msg(chat: list, role: str, content: str) -> list:
    return chat + [{"role": role, "content": content}]


# ── Stats panel ───────────────────────────────────────────────────────────────
INITIAL_STATS = """
<div style="height:280px;display:flex;flex-direction:column;align-items:center;
     justify-content:center;color:#aaa;font-size:13px;text-align:center;padding:24px">
  <div style="font-size:38px;margin-bottom:12px">📊</div>
  <div>Statistik kamu akan<br>muncul di sini</div>
</div>
"""

def build_stats_html(session: dict) -> str:
    if not session or session.get("q_num", 0) == 0:
        return INITIAL_STATS

    mastery  = get_mastery(session)
    q_num    = session.get("q_num", 0)
    correct  = session.get("correct_count", 0)
    wrong    = session.get("wrong_count", 0)
    answered = correct + wrong
    active   = session.get("active", False)
    acc_pct  = int(correct / answered * 100) if answered > 0 else 0

    target_bloom = mastery_to_bloom(mastery)
    b_emoji      = BLOOM_EMOJI.get(target_bloom, "")
    b_label      = BLOOM_LABELS.get(target_bloom, "")

    if mastery >= 0.7:
        bar_color = "#4CAF50"
    elif mastery >= 0.4:
        bar_color = "#2196F3"
    else:
        bar_color = "#90A4AE"

    if not active:
        badge_bg, badge_color, badge_text = "#E8F5E9", "#2E7D32", "🏁 Sesi Selesai"
    elif q_num >= MAX_QUESTIONS * 0.7:
        sisa = MAX_QUESTIONS - q_num + 1
        badge_bg, badge_color = "#FFF3E0", "#E65100"
        badge_text = f"⏳ {sisa} soal lagi"
    else:
        badge_bg, badge_color, badge_text = "#E3F2FD", "#1565C0", "📈 Terus semangat!"

    return f"""
<div style="font-family:sans-serif;padding:16px 14px;font-size:13px;line-height:1.6">

  <div style="font-weight:700;font-size:14px;margin-bottom:14px;color:#1a1a1a">
    📊 Progress Kamu
  </div>

  <!-- Mastery bar -->
  <div style="margin-bottom:14px">
    <div style="display:flex;justify-content:space-between;align-items:center;margin-bottom:5px">
      <span style="color:#555;font-size:12px;font-weight:600">Mastery</span>
      <span style="color:{bar_color};font-weight:800;font-size:22px">{mastery:.2f}</span>
    </div>
    <div style="background:#e8e8e8;border-radius:6px;height:14px">
      <div style="background:{bar_color};width:{min(mastery*100,100):.1f}%;
           height:100%;border-radius:6px;transition:width 0.5s ease"></div>
    </div>
  </div>

  <!-- Progress soal -->
  <div style="margin-bottom:14px">
    <div style="display:flex;justify-content:space-between;align-items:center;margin-bottom:5px">
      <span style="color:#555;font-size:12px;font-weight:600">Progress Soal</span>
      <span style="font-weight:700;color:#333">{q_num} / {MAX_QUESTIONS}</span>
    </div>
    <div style="background:#e8e8e8;border-radius:6px;height:10px">
      <div style="background:#90A4AE;width:{min(q_num/MAX_QUESTIONS*100,100):.1f}%;
           height:100%;border-radius:6px;transition:width 0.5s ease"></div>
    </div>
  </div>

  <!-- Grid benar / salah / akurasi -->
  <div style="display:grid;grid-template-columns:1fr 1fr 1fr;gap:7px;margin-bottom:14px">
    <div style="background:#E8F5E9;border-radius:8px;padding:8px 4px;text-align:center">
      <div style="font-size:22px;font-weight:800;color:#388E3C">{correct}</div>
      <div style="color:#777;font-size:10px">✅ Benar</div>
    </div>
    <div style="background:#FFEBEE;border-radius:8px;padding:8px 4px;text-align:center">
      <div style="font-size:22px;font-weight:800;color:#D32F2F">{wrong}</div>
      <div style="color:#777;font-size:10px">❌ Salah</div>
    </div>
    <div style="background:#E3F2FD;border-radius:8px;padding:8px 4px;text-align:center">
      <div style="font-size:22px;font-weight:800;color:#1565C0">{acc_pct}%</div>
      <div style="color:#777;font-size:10px">Akurasi</div>
    </div>
  </div>

  <!-- Level kognitif -->
  <div style="background:#EDE7F6;border-radius:8px;padding:10px;
       text-align:center;margin-bottom:10px">
    <div style="font-size:10px;color:#7B1FA2;margin-bottom:2px;font-weight:600">
      LEVEL KOGNITIF
    </div>
    <div style="font-size:14px;font-weight:700;color:#4A148C">
      {b_emoji} C{target_bloom} — {b_label}
    </div>
  </div>

  <!-- Status -->
  <div style="background:{badge_bg};border-radius:8px;padding:10px;text-align:center">
    <div style="font-size:13px;font-weight:600;color:{badge_color}">{badge_text}</div>
  </div>

</div>
"""


# ── Question message ──────────────────────────────────────────────────────────
def build_question_msg(result, q_num: int, mastery: float) -> str:
    q       = result.question
    bloom   = q.get("bloom_level", 1)
    context = q.get("context", "")

    lines = [
        f"### Soal {q_num} / {MAX_QUESTIONS}",
        (
            f"{BLOOM_EMOJI.get(bloom,'')} **C{bloom} — {BLOOM_LABELS.get(bloom,'')}** "
            f"&nbsp;|&nbsp; Mastery: **{mastery:.2f}**"
        ),
        f"*{q.get('topic', '')} — {q.get('concept', '')}*",
    ]
    if context:
        lines += ["", "---", "**📖 Bacaan:**", "", context, "", "---"]
    lines += ["", f"**{q.get('question', '')}**"]
    return "\n".join(lines)


# ── Event handlers ────────────────────────────────────────────────────────────
def start_quiz(keyword: str, session: dict):
    if not keyword.strip():
        return session, [], gr.update(interactive=False), "Masukkan topik terlebih dahulu.", INITIAL_STATS

    session = create_session(keyword)
    result  = get_next_question(session)

    if not result:
        session["active"] = False
        return session, [], gr.update(interactive=False), "Tidak ada soal untuk topik ini.", INITIAL_STATS

    session["q_num"]     = 1
    session["current_q"] = result
    session["seen_ids"].add(result.question.get("id"))

    mastery  = get_mastery(session)
    chat     = add_msg([], "assistant", build_question_msg(result, 1, mastery))
    progress = f"**{keyword}** &nbsp;|&nbsp; Soal 1/{MAX_QUESTIONS} &nbsp;|&nbsp; Mastery: {mastery:.2f}"

    return session, chat, gr.update(interactive=True), progress, build_stats_html(session)


def submit_answer(answer: str, session: dict, chat: list):
    if not answer.strip() or not session.get("active"):
        return session, chat, "", "", build_stats_html(session)

    keyword     = session["keyword"]
    concept_key = keyword.lower().strip()
    q           = session["current_q"].question
    context     = q.get("context", "")

    chat = add_msg(chat, "user", answer)

    # ── Judge ─────────────────────────────────────────────────────────────────
    verdict  = judge_answer(q.get("question", ""), answer, context)
    score    = verdict["score"]
    feedback = verdict["feedback"]

    if score >= 0.75:   label = "✅ BENAR"
    elif score >= 0.5:  label = "⚠️ SEBAGIAN BENAR"
    else:               label = "❌ SALAH"

    # ── Update counters ───────────────────────────────────────────────────────
    if score >= 0.5:
        session["correct_count"]     = session.get("correct_count", 0) + 1
        session["consecutive_wrong"] = 0
    else:
        session["wrong_count"]       = session.get("wrong_count", 0) + 1
        session["consecutive_wrong"] += 1

    # ── Weighted BKT update ───────────────────────────────────────────────────
    old_mastery = get_mastery(session)
    bkt_update_continuous(
        session["bkt"], concept_key,
        score=score,
        bloom_level=q.get("bloom_level", 1),
        question_id=q.get("id", ""),
    )
    new_mastery = get_mastery(session)

    verdict_msg = (
        f"{label} **(score: {score:.2f})**\n\n"
        f"💬 {feedback}\n\n"
        f"📈 Mastery: {old_mastery:.3f} → **{new_mastery:.3f}**"
    )
    chat = add_msg(chat, "assistant", verdict_msg)

    # ── Soal terakhir → sesi selesai ─────────────────────────────────────────
    if session["q_num"] >= MAX_QUESTIONS:
        answered = session["correct_count"] + session["wrong_count"]
        acc      = int(session["correct_count"] / answered * 100) if answered > 0 else 0
        chat = add_msg(chat, "assistant",
            f"🎉 **Sesi selesai!**\n\n"
            f"Mastery akhir: **{new_mastery:.3f}** &nbsp;|&nbsp; "
            f"{session['correct_count']}/{answered} benar ({acc}% akurasi)"
        )
        session["active"] = False
        return session, chat, "", f"✅ Selesai | {keyword} | Mastery: {new_mastery:.3f}", build_stats_html(session)

    if session["consecutive_wrong"] >= 2:
        chat = add_msg(chat, "assistant", "🔽 *Mencari soal lebih mudah...*")

    session["q_num"] += 1
    next_result = get_next_question(session)

    if not next_result:
        chat = add_msg(chat, "assistant", "📚 Semua soal tersedia sudah diberikan.")
        session["active"] = False
        return session, chat, "", f"Selesai | Mastery: {new_mastery:.3f}", build_stats_html(session)

    session["current_q"] = next_result
    session["seen_ids"].add(next_result.question.get("id"))
    chat = add_msg(chat, "assistant", build_question_msg(next_result, session["q_num"], new_mastery))

    progress = f"**{keyword}** &nbsp;|&nbsp; Soal {session['q_num']}/{MAX_QUESTIONS} &nbsp;|&nbsp; Mastery: {new_mastery:.3f}"
    return session, chat, "", progress, build_stats_html(session)


def reset_quiz():
    return {}, [], gr.update(interactive=False), "Masukkan topik untuk mulai.", INITIAL_STATS


# ── UI ────────────────────────────────────────────────────────────────────────
with gr.Blocks(
    theme=gr.themes.Soft(primary_hue="blue", secondary_hue="indigo"),
    title="CG-IR Adaptive Quiz",
) as demo:

    gr.Markdown(
        "# 🎓 CG-IR Adaptive Quiz\n"
        "**IPS / Geografi SMA** — Sistem kuis adaptif berbasis "
        "Knowledge Graph · FAISS · Bayesian Knowledge Tracing"
    )

    session_state = gr.State({})

    with gr.Row():
        with gr.Column(scale=4):
            with gr.Row():
                keyword_input = gr.Textbox(
                    label="Topik",
                    placeholder="contoh: pancasila, kearifan lokal, perang dunia...",
                    lines=1,
                    scale=5,
                )
                start_btn = gr.Button("▶️ Mulai", variant="primary", scale=1)
                reset_btn = gr.Button("🔄", scale=0, min_width=52)

            progress_display = gr.Markdown("Masukkan topik untuk mulai.")

            chatbot = gr.Chatbot(label="Sesi Quiz", height=500)

            with gr.Row():
                answer_input = gr.Textbox(
                    label="Jawaban",
                    placeholder="Ketik jawaban lalu tekan Enter atau klik Submit...",
                    lines=3,
                    scale=5,
                    interactive=False,
                )
                submit_btn = gr.Button("📨 Submit", variant="primary", scale=1)

        with gr.Column(scale=1, min_width=220):
            stats_panel = gr.HTML(value=INITIAL_STATS)

    start_btn.click(
        fn=start_quiz,
        inputs=[keyword_input, session_state],
        outputs=[session_state, chatbot, answer_input, progress_display, stats_panel],
    )
    submit_btn.click(
        fn=submit_answer,
        inputs=[answer_input, session_state, chatbot],
        outputs=[session_state, chatbot, answer_input, progress_display, stats_panel],
    )
    answer_input.submit(
        fn=submit_answer,
        inputs=[answer_input, session_state, chatbot],
        outputs=[session_state, chatbot, answer_input, progress_display, stats_panel],
    )
    reset_btn.click(
        fn=reset_quiz,
        inputs=[],
        outputs=[session_state, chatbot, answer_input, progress_display, stats_panel],
    )

if __name__ == "__main__":
    demo.launch()