import base64
import io
import json
import os
import re
import tempfile
import threading
import time
from pathlib import Path
from typing import Any, Dict, List, Optional, Tuple

import gradio as gr
import nbformat
import requests
from docx import Document
from docx.enum.text import WD_ALIGN_PARAGRAPH
from docx.shared import Inches, Pt
from openai import OpenAI
from PIL import Image

# ==========================================================
# KONFIGURASI NVIDIA API
# ==========================================================
NVIDIA_API_KEY = os.getenv("NVIDIA_API_KEY", "")
NVIDIA_MODEL = os.getenv("NVIDIA_MODEL", "qwen/qwen3-coder-480b-a35b-instruct")
NVIDIA_BASE_URL = os.getenv("NVIDIA_BASE_URL", "https://integrate.api.nvidia.com/v1")

client = OpenAI(
    api_key=NVIDIA_API_KEY,
    base_url=NVIDIA_BASE_URL,
)

# ==========================================================
# KONFIGURASI TELEGRAM BOT
# ==========================================================
TELEGRAM_BOT_TOKEN = os.getenv("TELEGRAM_BOT_TOKEN", "")
TELEGRAM_ENABLED = os.getenv("TELEGRAM_ENABLED", "true").strip().lower() in ["1", "true", "yes", "on"]
TELEGRAM_ALLOWED_USER_IDS_RAW = os.getenv("TELEGRAM_ALLOWED_USER_IDS", "").strip()
TELEGRAM_MAX_HISTORY = int(os.getenv("TELEGRAM_MAX_HISTORY", "8"))
TELEGRAM_REQUEST_TIMEOUT = int(os.getenv("TELEGRAM_REQUEST_TIMEOUT", "60"))
TELEGRAM_RETRY_COUNT = int(os.getenv("TELEGRAM_RETRY_COUNT", "3"))

telegram_histories: Dict[int, List[Dict[str, str]]] = {}
telegram_bot_status: Dict[str, Any] = {
    "running": False,
    "last_error": "",
    "last_update": "",
}
telegram_thread_started = False

# Sesi laporan Telegram disimpan sementara di memori Space.
# Jika Space restart/sleep, data sesi akan hilang dan user perlu upload ulang.
telegram_report_sessions: Dict[int, Dict[str, Any]] = {}

SYSTEM_PROMPT = """
Kamu adalah AI Assistant Mahasiswa.

Tugas utamamu membantu mahasiswa dalam kegiatan akademik, terutama:
- menjelaskan materi kuliah,
- membantu coding dasar,
- membantu memahami error,
- membantu membuat laporan akademik,
- membantu membuat laporan Deep Learning dari kode notebook.

Gunakan bahasa Indonesia yang jelas, sopan, dan mudah dipahami.
Untuk kebutuhan laporan, gunakan bahasa formal seperti laporan mahasiswa.

Aturan penting:
- Jangan mengarang data, akurasi, loss, jumlah dataset, nama kelas, hasil evaluasi, atau referensi.
- Jika angka atau hasil tidak ada di notebook, jangan dibuat-buat.
- Jika memberi contoh, beri label sebagai contoh.
- Bantu mahasiswa memahami isi, bukan sekadar memberi jawaban untuk disalin mentah-mentah.
"""

REPORT_STYLE_PROMPT = """
Kamu membuat penjelasan laporan praktikum Deep Learning dari satu bagian kode notebook.

Format jawaban WAJIB JSON valid:
{
  "judul": "Judul Bagian Singkat",
  "penjelasan": "Penjelasannya: ..."
}

Gaya laporan:
- Bahasa Indonesia formal, seperti laporan mahasiswa.
- Judul bagian singkat dan relevan, misalnya Import Library, Pembacaan Dataset, Split Dataset, Arsitektur Model, Training Model, Evaluasi Model, Confusion Matrix, Deployment Gradio.
- Penjelasan diawali persis dengan: Penjelasannya:
- Buat penjelasan cukup lengkap, idealnya 2 sampai 4 paragraf pendek.
- Jelaskan fungsi kode, tujuan tahap, alur proses, library/fungsi penting, dan arti output jika output tersedia.
- Jangan terlalu singkat. Hindari penjelasan satu kalimat kecuali kode memang sangat sederhana.
- Jangan mengarang angka, akurasi, loss, dataset, nama kelas, atau hasil evaluasi yang tidak ada pada kode/output.
- Jika output kosong, jelaskan fungsi kode, tujuan tahap, serta hubungannya dengan proses Deep Learning.
- Jangan membuat markdown, jangan membuat bullet panjang, dan jangan menambahkan teks di luar JSON.
"""

# ==========================================================
# HELPER UMUM
# ==========================================================
def get_uploaded_path(file_obj: Any) -> str:
    """Mengambil path file dari komponen gr.File."""
    if file_obj is None:
        raise gr.Error("File belum di-upload.")
    if isinstance(file_obj, str):
        return file_obj
    if hasattr(file_obj, "name"):
        return file_obj.name
    if isinstance(file_obj, dict) and "path" in file_obj:
        return file_obj["path"]
    raise gr.Error("Format file upload tidak dikenali.")


def safe_filename(text: str) -> str:
    text = str(text or "").strip().lower()
    text = re.sub(r"[^a-z0-9A-Z_-]+", "_", text)
    text = re.sub(r"_+", "_", text).strip("_")
    return text or "laporan_deep_learning"


def clean_text_for_docx(text: Any) -> str:
    """
    Membersihkan teks agar aman dimasukkan ke file DOCX/XML.
    Error yang diperbaiki:
    ValueError: All strings must be XML compatible: Unicode or ASCII, no NULL bytes or control characters
    """
    if text is None:
        return ""

    text = str(text)

    # Hapus ANSI escape sequence dari output terminal/notebook.
    text = re.sub(r"\x1B(?:[@-Z\\-_]|\[[0-?]*[ -/]*[@-~])", "", text)

    # Normalisasi newline.
    text = text.replace("\r\n", "\n").replace("\r", "\n")

    # Hapus karakter ilegal XML 1.0, kecuali tab, newline, carriage return.
    text = re.sub(r"[\x00-\x08\x0B\x0C\x0E-\x1F]", "", text)

    # Hapus karakter non-XML/surrogate yang kadang muncul dari output notebook.
    text = re.sub(r"[\uD800-\uDFFF\uFFFE\uFFFF]", "", text)

    return text


def split_text_chunks(text: str, chunk_size: int = 6000) -> List[str]:
    """
    Memecah teks panjang agar aman dimasukkan ke run DOCX.
    Tidak memotong isi; hanya membagi ke beberapa paragraph/run.
    """
    text = clean_text_for_docx(text)
    if not text:
        return [""]

    chunks: List[str] = []
    current: List[str] = []
    current_len = 0

    for line in text.splitlines(True):
        if current_len + len(line) > chunk_size and current:
            chunks.append("".join(current))
            current = [line]
            current_len = len(line)
        else:
            current.append(line)
            current_len += len(line)

    if current:
        chunks.append("".join(current))

    return chunks or [""]


def truncate_text(text: str, limit: int) -> str:
    """
    Dipakai hanya untuk prompt ke model agar tidak terlalu berat.
    Isi draft dan DOCX tidak lagi memakai fungsi ini untuk memotong kode/output.
    """
    text = clean_text_for_docx(text)
    if not text:
        return ""
    if len(text) <= limit:
        return text
    return text[:limit] + "\n... [hanya bagian prompt AI yang dipersingkat; kode/output lengkap tetap masuk draft/DOCX]"


def ask_nvidia(messages: List[Dict[str, str]], temperature: float = 0.4, max_tokens: int = 1200) -> str:
    if not NVIDIA_API_KEY:
        raise gr.Error("NVIDIA_API_KEY belum diset di Hugging Face Secrets.")

    response = client.chat.completions.create(
        model=NVIDIA_MODEL,
        messages=messages,
        temperature=temperature,
        max_tokens=max_tokens,
    )
    return response.choices[0].message.content or ""


def parse_json_from_model(text: str) -> Optional[Dict[str, Any]]:
    """Mencoba mengambil JSON dari jawaban model."""
    try:
        return json.loads(text)
    except Exception:
        pass

    match = re.search(r"\{.*\}", text, flags=re.DOTALL)
    if not match:
        return None
    try:
        return json.loads(match.group(0))
    except Exception:
        return None


# ==========================================================
# CHAT UMUM
# ==========================================================
def normal_chat(message: str, history: Optional[List[Any]]):
    if not NVIDIA_API_KEY:
        return "Error: NVIDIA_API_KEY belum diset di Hugging Face Secrets."

    messages = [{"role": "system", "content": SYSTEM_PROMPT}]

    # Gradio 4 biasanya memberi history list tuple: [(user, assistant), ...]
    for item in history or []:
        if isinstance(item, (list, tuple)) and len(item) >= 2:
            user_msg, assistant_msg = item[0], item[1]
            if user_msg:
                messages.append({"role": "user", "content": str(user_msg)})
            if assistant_msg:
                messages.append({"role": "assistant", "content": str(assistant_msg)})

    messages.append({"role": "user", "content": message})

    try:
        return ask_nvidia(messages, temperature=0.5, max_tokens=1500)
    except Exception as e:
        return f"Terjadi error saat menghubungi NVIDIA API: {str(e)}"


def test_api_connection():
    if not NVIDIA_API_KEY:
        return "❌ NVIDIA_API_KEY belum diset di Hugging Face Secrets."
    try:
        result = ask_nvidia(
            [
                {"role": "system", "content": "Jawab singkat."},
                {"role": "user", "content": "Balas hanya dengan kata: OK"},
            ],
            temperature=0,
            max_tokens=20,
        )
        return f"✅ API aktif. Model: {NVIDIA_MODEL}. Respons: {result.strip()}"
    except Exception as e:
        return f"❌ API error: {str(e)}"


# ==========================================================
# TELEGRAM BOT INTEGRATION
# ==========================================================
def parse_allowed_telegram_ids(raw: str) -> set:
    ids = set()
    for item in (raw or "").replace(";", ",").split(","):
        item = item.strip()
        if not item:
            continue
        try:
            ids.add(int(item))
        except ValueError:
            pass
    return ids


TELEGRAM_ALLOWED_USER_IDS = parse_allowed_telegram_ids(TELEGRAM_ALLOWED_USER_IDS_RAW)


def telegram_api_url(method: str) -> str:
    return f"https://api.telegram.org/bot{TELEGRAM_BOT_TOKEN}/{method}"


def telegram_file_url(file_path: str) -> str:
    return f"https://api.telegram.org/file/bot{TELEGRAM_BOT_TOKEN}/{file_path}"


def telegram_api_call(method: str, payload: Optional[Dict[str, Any]] = None, timeout: int = 35) -> Dict[str, Any]:
    """
    Memanggil Telegram API dengan retry.
    Ini membantu jika Hugging Face/Telegram sedang lambat dan muncul error:
    Read timed out.
    """
    if not TELEGRAM_BOT_TOKEN:
        raise RuntimeError("TELEGRAM_BOT_TOKEN belum diset.")

    last_error = None
    effective_timeout = max(timeout, TELEGRAM_REQUEST_TIMEOUT)

    for attempt in range(1, TELEGRAM_RETRY_COUNT + 1):
        try:
            response = requests.post(
                telegram_api_url(method),
                json=payload or {},
                timeout=(10, effective_timeout),  # connect timeout, read timeout
            )
            response.raise_for_status()
            data = response.json()
            if not data.get("ok"):
                raise RuntimeError(data.get("description", "Telegram API error"))
            return data

        except requests.exceptions.ReadTimeout as e:
            last_error = e
            telegram_bot_status["last_error"] = f"Telegram read timeout pada {method}, percobaan {attempt}/{TELEGRAM_RETRY_COUNT}"
            time.sleep(min(2 * attempt, 8))

        except requests.exceptions.ConnectTimeout as e:
            last_error = e
            telegram_bot_status["last_error"] = f"Telegram connect timeout pada {method}, percobaan {attempt}/{TELEGRAM_RETRY_COUNT}"
            time.sleep(min(2 * attempt, 8))

        except requests.exceptions.ConnectionError as e:
            last_error = e
            telegram_bot_status["last_error"] = f"Telegram connection error pada {method}, percobaan {attempt}/{TELEGRAM_RETRY_COUNT}"
            time.sleep(min(2 * attempt, 8))

        except Exception as e:
            # Error non-timeout seperti token salah tetap langsung dilempar agar jelas.
            raise e

    raise RuntimeError(
        f"Telegram API timeout saat memanggil {method}. "
        f"Coba Restart Space atau naikkan TELEGRAM_REQUEST_TIMEOUT. Detail: {last_error}"
    )


def telegram_default_keyboard() -> Dict[str, Any]:
    return {
        "keyboard": [
            [{"text": "👋 Menu"}, {"text": "💬 Chat AI"}],
            [{"text": "📄 Laporan IPYNB"}, {"text": "📝 Perintah Tugas"}],
            [{"text": "✅ Buat Draft"}, {"text": "➕ Tambah Bagian"}],
            [{"text": "♻️ Ubah Draft"}, {"text": "📥 Export DOCX"}],
            [{"text": "🧾 Set Cover"}, {"text": "🧹 Reset"}],
        ],
        "resize_keyboard": True,
        "one_time_keyboard": False,
        "is_persistent": True,
    }


def telegram_send_message(chat_id: int, text: str, reply_markup: Optional[Dict[str, Any]] = None):
    for chunk in split_telegram_message(text):
        payload = {
            "chat_id": chat_id,
            "text": chunk,
            "disable_web_page_preview": True,
        }
        if reply_markup:
            payload["reply_markup"] = reply_markup
        telegram_api_call("sendMessage", payload, timeout=60)


def telegram_send_menu(chat_id: int, text: str):
    telegram_send_message(chat_id, text, reply_markup=telegram_default_keyboard())


def telegram_send_document(chat_id: int, file_path: str, caption: str = ""):
    if not TELEGRAM_BOT_TOKEN:
        raise RuntimeError("TELEGRAM_BOT_TOKEN belum diset.")

    with open(file_path, "rb") as f:
        response = requests.post(
            telegram_api_url("sendDocument"),
            data={"chat_id": chat_id, "caption": caption[:1000]},
            files={"document": f},
            timeout=60,
        )
    response.raise_for_status()
    data = response.json()
    if not data.get("ok"):
        raise RuntimeError(data.get("description", "Gagal mengirim dokumen Telegram"))


def telegram_send_text_as_document(chat_id: int, text: str, filename: str, caption: str = ""):
    tmp_path = str(Path(tempfile.gettempdir()) / safe_filename(filename.replace(".txt", "")))
    if not tmp_path.endswith(".txt"):
        tmp_path += ".txt"
    Path(tmp_path).write_text(clean_text_for_docx(text), encoding="utf-8")
    telegram_send_document(chat_id, tmp_path, caption=caption)


def split_telegram_message(text: str, limit: int = 3500) -> List[str]:
    text = clean_text_for_docx(text).strip()
    if not text:
        return [""]

    chunks = []
    while len(text) > limit:
        split_at = text.rfind("\n", 0, limit)
        if split_at < 500:
            split_at = text.rfind(" ", 0, limit)
        if split_at < 500:
            split_at = limit
        chunks.append(text[:split_at].strip())
        text = text[split_at:].strip()
    if text:
        chunks.append(text)
    return chunks


def telegram_user_allowed(user_id: Optional[int]) -> bool:
    if not TELEGRAM_ALLOWED_USER_IDS:
        return True
    return user_id in TELEGRAM_ALLOWED_USER_IDS


def telegram_get_session(chat_id: int) -> Dict[str, Any]:
    if chat_id not in telegram_report_sessions:
        telegram_report_sessions[chat_id] = {
            "mode": "chat",
            "expecting": "",
            "pending_edit_mode": "Tambah bagian baru di akhir",
            "notebook_paths": [],
            "notebook_names": [],
            "assignment_instructions": "",
            "draft_text": "",
            "report_state": {},
            "cover": telegram_default_cover(),
        }
    return telegram_report_sessions[chat_id]


def telegram_default_cover() -> Dict[str, str]:
    return {
        "judul_laporan": os.getenv("TELEGRAM_REPORT_TITLE", "Laporan Deep Learning"),
        "nama": os.getenv("TELEGRAM_REPORT_NAMA", ""),
        "nim": os.getenv("TELEGRAM_REPORT_NIM", ""),
        "dosen": os.getenv("TELEGRAM_REPORT_DOSEN", ""),
        "kelas": os.getenv("TELEGRAM_REPORT_KELAS", ""),
        "anggota": os.getenv("TELEGRAM_REPORT_ANGGOTA", ""),
        "prodi": os.getenv("TELEGRAM_REPORT_PRODI", "Program Studi Teknik Informatika"),
        "kampus": os.getenv("TELEGRAM_REPORT_KAMPUS", ""),
        "tahun": os.getenv("TELEGRAM_REPORT_TAHUN", "2025"),
    }


def telegram_cover_summary(cover: Dict[str, str]) -> str:
    return (
        "Data cover saat ini:\n"
        f"Judul: {cover.get('judul_laporan', '') or '-'}\n"
        f"Nama: {cover.get('nama', '') or '-'}\n"
        f"NIM: {cover.get('nim', '') or '-'}\n"
        f"Dosen: {cover.get('dosen', '') or '-'}\n"
        f"Kelas: {cover.get('kelas', '') or '-'}\n"
        f"Anggota: {cover.get('anggota', '') or '-'}\n"
        f"Prodi: {cover.get('prodi', '') or '-'}\n"
        f"Kampus: {cover.get('kampus', '') or '-'}\n"
        f"Tahun: {cover.get('tahun', '') or '-'}"
    )


def telegram_parse_cover_text(text: str, current_cover: Dict[str, str]) -> Dict[str, str]:
    cover = dict(current_cover or telegram_default_cover())
    key_map = {
        "judul": "judul_laporan",
        "judul laporan": "judul_laporan",
        "nama": "nama",
        "nim": "nim",
        "dosen": "dosen",
        "dosen mata kuliah": "dosen",
        "kelas": "kelas",
        "anggota": "anggota",
        "partner": "anggota",
        "prodi": "prodi",
        "program studi": "prodi",
        "kampus": "kampus",
        "tahun": "tahun",
    }

    for raw_line in clean_text_for_docx(text).splitlines():
        if ":" not in raw_line:
            continue
        key, value = raw_line.split(":", 1)
        normalized = key.strip().lower()
        mapped = key_map.get(normalized)
        if mapped:
            cover[mapped] = value.strip()

    return cover


def telegram_welcome_text() -> str:
    return (
        "Halo! 👋\n"
        "Selamat datang di AI Assistant Mahasiswa.\n\n"
        "Kamu bisa pakai tombol di bawah tanpa mengetik perintah garis miring.\n\n"
        "Fitur Telegram:\n"
        "💬 Chat AI untuk tanya materi, coding, dan tugas.\n"
        "📄 Laporan IPYNB untuk upload notebook dan membuat draft laporan.\n"
        "📝 Perintah Tugas untuk memasukkan instruksi dosen.\n"
        "✅ Buat Draft untuk membuat draft laporan dari file IPYNB.\n"
        "➕ Tambah Bagian / ♻️ Ubah Draft untuk minta AI memperbaiki draft.\n"
        "📥 Export DOCX untuk mengirim hasil Word ke Telegram.\n\n"
        "Catatan: data sesi Telegram disimpan sementara. Jika Space restart, upload file perlu diulang."
    )


def telegram_build_messages(chat_id: int, user_text: str) -> List[Dict[str, str]]:
    history = telegram_histories.get(chat_id, [])[-TELEGRAM_MAX_HISTORY:]
    messages = [
        {
            "role": "system",
            "content": (
                SYSTEM_PROMPT
                + "\n\nKamu sedang menjawab lewat Telegram. Jawab ringkas, jelas, dan tetap sopan. "
                  "Jika user ingin membuat laporan dari file .ipynb, arahkan memakai tombol 📄 Laporan IPYNB."
            ),
        }
    ]
    messages.extend(history)
    messages.append({"role": "user", "content": user_text})
    return messages


def telegram_save_history(chat_id: int, user_text: str, assistant_text: str):
    history = telegram_histories.get(chat_id, [])
    history.append({"role": "user", "content": user_text})
    history.append({"role": "assistant", "content": assistant_text})
    telegram_histories[chat_id] = history[-(TELEGRAM_MAX_HISTORY * 2):]


def telegram_download_document(document: Dict[str, Any]) -> str:
    file_name = document.get("file_name") or "notebook.ipynb"
    if not file_name.lower().endswith(".ipynb"):
        raise RuntimeError("File harus berformat .ipynb")

    file_id = document.get("file_id")
    if not file_id:
        raise RuntimeError("file_id Telegram tidak ditemukan.")

    data = telegram_api_call("getFile", {"file_id": file_id}, timeout=20)
    file_path = data.get("result", {}).get("file_path")
    if not file_path:
        raise RuntimeError("file_path Telegram tidak ditemukan.")

    safe_name = safe_filename(Path(file_name).stem) + ".ipynb"
    local_path = str(Path(tempfile.gettempdir()) / f"telegram_{int(time.time())}_{safe_name}")

    response = requests.get(telegram_file_url(file_path), timeout=90)
    response.raise_for_status()
    Path(local_path).write_bytes(response.content)
    return local_path


def telegram_handle_document(message: Dict[str, Any]):
    chat = message.get("chat", {}) or {}
    user = message.get("from", {}) or {}
    chat_id = chat.get("id")
    user_id = user.get("id")

    if not chat_id:
        return

    if not telegram_user_allowed(user_id):
        telegram_send_menu(chat_id, "Maaf, bot ini dibatasi untuk user tertentu.")
        return

    document = message.get("document") or {}
    session = telegram_get_session(chat_id)

    try:
        local_path = telegram_download_document(document)
        session["notebook_paths"].append(local_path)
        session["notebook_names"].append(document.get("file_name") or Path(local_path).name)
        session["mode"] = "report"

        telegram_send_menu(
            chat_id,
            "✅ File IPYNB berhasil diterima.\n\n"
            f"Total notebook tersimpan: {len(session['notebook_paths'])}\n"
            "Kamu bisa upload file IPYNB lain, isi Perintah Tugas, atau tekan ✅ Buat Draft."
        )
    except Exception as e:
        telegram_send_menu(chat_id, f"❌ Gagal menerima file: {str(e)}")


def telegram_make_report_draft(chat_id: int):
    session = telegram_get_session(chat_id)
    notebook_paths = session.get("notebook_paths", [])

    if not notebook_paths:
        telegram_send_menu(
            chat_id,
            "Belum ada file IPYNB. Tekan 📄 Laporan IPYNB lalu upload satu atau beberapa file `.ipynb`."
        )
        return

    telegram_send_message(chat_id, "⏳ Sedang membaca notebook dan membuat draft laporan. Ini bisa memakan waktu...")

    all_code_sections: List[Dict[str, Any]] = []
    notebook_names: List[str] = []

    for ipynb_path in notebook_paths:
        notebook_name = Path(ipynb_path).name
        notebook_names.append(notebook_name)

        sections = read_ipynb(ipynb_path)
        code_sections = [item for item in sections if item.get("type") == "code"]

        for item in code_sections:
            item["context"] = f"Notebook: {notebook_name}\n" + str(item.get("context", ""))

        all_code_sections.extend(code_sections)

    if not all_code_sections:
        telegram_send_menu(chat_id, "❌ Notebook tidak memiliki cell kode yang bisa dibuat menjadi laporan.")
        return

    multi_file_note = ""
    if len(notebook_names) > 1:
        multi_file_note = (
            "Laporan ini dibuat dari beberapa file notebook berikut:\n"
            + "\n".join(f"- {name}" for name in notebook_names)
        )

    assignment = session.get("assignment_instructions", "").strip()
    combined_assignment = assignment
    if multi_file_note:
        combined_assignment = (multi_file_note + "\n\n" + assignment).strip()

    draft_text, image_map = build_report_draft_text(
        code_sections=all_code_sections,
        assignment_instructions=combined_assignment,
        include_code=True,
        include_output=True,
        include_images=True,
    )

    cover = session.get("cover") or telegram_default_cover()
    report_state = {
        "image_map": image_map,
        "judul_laporan": cover.get("judul_laporan", ""),
        "nama": cover.get("nama", ""),
        "nim": cover.get("nim", ""),
        "dosen": cover.get("dosen", ""),
        "kelas": cover.get("kelas", ""),
        "anggota": cover.get("anggota", ""),
        "prodi": cover.get("prodi", ""),
        "kampus": cover.get("kampus", ""),
        "tahun": cover.get("tahun", ""),
        "assignment_instructions": combined_assignment,
        "notebook_files": notebook_names,
    }

    session["draft_text"] = draft_text
    session["report_state"] = report_state

    preview = truncate_text(draft_text, 2500)
    telegram_send_menu(
        chat_id,
        "✅ Draft laporan berhasil dibuat.\n\n"
        f"Jumlah notebook: {len(notebook_names)}\n"
        f"Jumlah bagian kode: {len(all_code_sections)}\n\n"
        "Preview awal:\n\n"
        f"{preview}\n\n"
        "Saya juga kirim draft lengkap sebagai file TXT. Kamu bisa tekan ➕ Tambah Bagian, ♻️ Ubah Draft, atau 📥 Export DOCX."
    )

    telegram_send_text_as_document(
        chat_id,
        draft_text,
        "draft_laporan.txt",
        caption="Draft lengkap laporan dalam format TXT untuk review."
    )


def telegram_export_docx(chat_id: int):
    session = telegram_get_session(chat_id)
    draft = session.get("draft_text", "")
    report_state = session.get("report_state", {})
    cover = session.get("cover") or telegram_default_cover()

    if not draft.strip():
        telegram_send_menu(chat_id, "Draft masih kosong. Tekan ✅ Buat Draft terlebih dahulu.")
        return

    telegram_send_message(chat_id, "⏳ Sedang membuat file DOCX...")

    try:
        docx_path = export_review_to_docx(
            edited_draft=draft,
            report_state=report_state,
            judul_laporan=cover.get("judul_laporan", "Laporan Deep Learning"),
            nama=cover.get("nama", ""),
            nim=cover.get("nim", ""),
            dosen=cover.get("dosen", ""),
            kelas=cover.get("kelas", ""),
            anggota=cover.get("anggota", ""),
            prodi=cover.get("prodi", ""),
            kampus=cover.get("kampus", ""),
            tahun=cover.get("tahun", "2025"),
        )
        telegram_send_document(chat_id, docx_path, caption="✅ Laporan DOCX berhasil dibuat.")
        telegram_send_menu(chat_id, "Selesai. Kamu masih bisa minta AI ubah draft lalu export ulang.")
    except Exception as e:
        telegram_send_menu(chat_id, f"❌ Gagal export DOCX: {str(e)}")


def telegram_apply_ai_edit(chat_id: int, instruction: str):
    session = telegram_get_session(chat_id)
    draft = session.get("draft_text", "")
    report_state = session.get("report_state", {})
    edit_mode = session.get("pending_edit_mode", "Tambah bagian baru di akhir")

    if not draft.strip():
        telegram_send_menu(chat_id, "Draft masih kosong. Buat draft dulu dengan tombol ✅ Buat Draft.")
        return

    telegram_send_message(chat_id, "⏳ AI sedang mengubah draft sesuai instruksi kamu...")

    try:
        updated_draft, status = ai_update_review_draft(
            edited_draft=draft,
            report_state=report_state,
            edit_instruction=instruction,
            edit_mode=edit_mode,
        )
        session["draft_text"] = clean_text_for_docx(updated_draft)
        session["expecting"] = ""

        preview = truncate_text(session["draft_text"], 2500)
        telegram_send_menu(chat_id, f"{status}\n\nPreview:\n\n{preview}")
        telegram_send_text_as_document(
            chat_id,
            session["draft_text"],
            "draft_laporan_revisi.txt",
            caption="Draft revisi terbaru dalam format TXT."
        )
    except Exception as e:
        session["expecting"] = ""
        telegram_send_menu(chat_id, f"❌ Gagal mengubah draft: {str(e)}")


def telegram_handle_text(message: Dict[str, Any]):
    chat = message.get("chat", {}) or {}
    user = message.get("from", {}) or {}
    chat_id = chat.get("id")
    user_id = user.get("id")
    text = clean_text_for_docx((message.get("text") or "").strip())

    if not chat_id:
        return

    if not telegram_user_allowed(user_id):
        telegram_send_menu(
            chat_id,
            "Maaf, bot ini dibatasi untuk user tertentu. Tambahkan Telegram user ID kamu ke TELEGRAM_ALLOWED_USER_IDS di Hugging Face.",
        )
        return

    session = telegram_get_session(chat_id)
    lower = text.lower()

    if lower in ["/start", "/help", "start", "help", "👋 menu", "menu"]:
        telegram_send_menu(chat_id, telegram_welcome_text())
        return

    if text == "💬 Chat AI":
        session["mode"] = "chat"
        session["expecting"] = ""
        telegram_send_menu(chat_id, "Silakan ketik pertanyaan kamu. Saya akan jawab sebagai AI Assistant Mahasiswa.")
        return

    if text == "📄 Laporan IPYNB":
        session["mode"] = "report"
        session["expecting"] = ""
        telegram_send_menu(
            chat_id,
            "Mode laporan aktif.\n\n"
            "Silakan upload satu atau beberapa file `.ipynb` langsung di Telegram.\n"
            "Setelah upload, kamu bisa tekan 📝 Perintah Tugas untuk menambahkan instruksi dosen, lalu tekan ✅ Buat Draft."
        )
        return

    if text == "📝 Perintah Tugas":
        session["mode"] = "report"
        session["expecting"] = "task_instruction"
        telegram_send_menu(
            chat_id,
            "Kirim perintah tugas dari dosen dalam satu pesan.\n\n"
            "Contoh:\n"
            "Buat laporan yang menjelaskan preprocessing, arsitektur model, training, evaluasi, confusion matrix, dan kesimpulan."
        )
        return

    if text == "🧾 Set Cover":
        session["expecting"] = "cover"
        telegram_send_menu(
            chat_id,
            telegram_cover_summary(session.get("cover") or telegram_default_cover())
            + "\n\nKirim data cover dengan format seperti ini:\n"
              "Judul: Klasifikasi Penyakit Mata\n"
              "Nama: Nama Kamu\n"
              "NIM: 123456\n"
              "Dosen: Nama Dosen\n"
              "Kelas: ILB\n"
              "Anggota: Nama Partner\n"
              "Prodi: Teknik Informatika\n"
              "Kampus: Politeknik Caltex Riau\n"
              "Tahun: 2025"
        )
        return

    if text == "✅ Buat Draft":
        try:
            telegram_make_report_draft(chat_id)
        except Exception as e:
            telegram_send_menu(chat_id, f"❌ Gagal membuat draft: {str(e)}")
        return

    if text == "➕ Tambah Bagian":
        session["expecting"] = "edit_instruction"
        session["pending_edit_mode"] = "Tambah bagian baru di akhir"
        telegram_send_menu(
            chat_id,
            "Kirim instruksi bagian yang ingin ditambahkan.\n\n"
            "Contoh:\nTambahkan bagian kesimpulan dan saran berdasarkan isi laporan."
        )
        return

    if text == "♻️ Ubah Draft":
        session["expecting"] = "edit_instruction"
        session["pending_edit_mode"] = "Ubah draft penuh sesuai instruksi"
        telegram_send_menu(
            chat_id,
            "Kirim instruksi perubahan draft.\n\n"
            "Contoh:\nRapikan bahasa laporan agar lebih formal dan perjelas bagian evaluasi model."
        )
        return

    if text == "📥 Export DOCX":
        telegram_export_docx(chat_id)
        return

    if text == "ℹ️ Status":
        telegram_send_menu(chat_id, telegram_status_text())
        return

    if text == "🧹 Reset":
        telegram_histories.pop(chat_id, None)
        telegram_report_sessions.pop(chat_id, None)
        telegram_send_menu(chat_id, "✅ Sesi Telegram sudah direset. Kamu bisa mulai lagi dari menu.")
        return

    # Proses input lanjutan berdasarkan state
    expecting = session.get("expecting", "")

    if expecting == "task_instruction":
        session["assignment_instructions"] = text
        session["expecting"] = ""
        telegram_send_menu(
            chat_id,
            "✅ Perintah tugas sudah disimpan.\n\n"
            "Sekarang upload file `.ipynb` atau tekan ✅ Buat Draft jika file sudah diupload."
        )
        return

    if expecting == "cover":
        session["cover"] = telegram_parse_cover_text(text, session.get("cover") or telegram_default_cover())
        session["expecting"] = ""
        telegram_send_menu(chat_id, "✅ Data cover diperbarui.\n\n" + telegram_cover_summary(session["cover"]))
        return

    if expecting == "edit_instruction":
        telegram_apply_ai_edit(chat_id, text)
        return

    # Default: chat AI tanpa perlu tombol slash command
    try:
        telegram_send_message(chat_id, "⏳ Sedang diproses...")
        messages = telegram_build_messages(chat_id, text)
        answer = ask_nvidia(messages, temperature=0.5, max_tokens=1800)
        answer = clean_text_for_docx(answer).strip() or "Maaf, AI tidak memberi jawaban."
        telegram_save_history(chat_id, text, answer)
        telegram_send_menu(chat_id, answer)
    except Exception as e:
        telegram_bot_status["last_error"] = str(e)
        telegram_send_menu(chat_id, f"❌ Terjadi error: {str(e)}")


def telegram_handle_message(message: Dict[str, Any]):
    if message.get("document"):
        telegram_handle_document(message)
    else:
        telegram_handle_text(message)


def telegram_polling_loop():
    telegram_bot_status["running"] = True
    offset = None

    try:
        # Jika sebelumnya bot pernah memakai webhook, long polling tidak akan berjalan sampai webhook dihapus.
        telegram_api_call("deleteWebhook", {"drop_pending_updates": False}, timeout=60)
    except Exception as e:
        telegram_bot_status["last_error"] = f"deleteWebhook error: {str(e)}"

    while TELEGRAM_ENABLED and TELEGRAM_BOT_TOKEN:
        try:
            payload = {
                "timeout": 25,
                "allowed_updates": ["message"],
            }
            if offset is not None:
                payload["offset"] = offset

            data = telegram_api_call("getUpdates", payload, timeout=75)
            updates = data.get("result", [])

            for update in updates:
                offset = update.get("update_id", 0) + 1
                telegram_bot_status["last_update"] = str(update.get("update_id", ""))
                message = update.get("message")
                if message:
                    telegram_handle_message(message)

            telegram_bot_status["running"] = True
            telegram_bot_status["last_error"] = ""
        except Exception as e:
            telegram_bot_status["running"] = False
            telegram_bot_status["last_error"] = str(e)
            time.sleep(5)


def start_telegram_bot_if_configured():
    global telegram_thread_started

    if telegram_thread_started:
        return

    if not TELEGRAM_ENABLED or not TELEGRAM_BOT_TOKEN:
        return

    telegram_thread_started = True
    thread = threading.Thread(target=telegram_polling_loop, daemon=True)
    thread.start()


def telegram_status_text():
    if not TELEGRAM_BOT_TOKEN:
        return (
            "❌ TELEGRAM_BOT_TOKEN belum diset.\n\n"
            "Buat bot lewat @BotFather, lalu simpan token di Hugging Face sebagai Secret bernama TELEGRAM_BOT_TOKEN."
        )

    try:
        data = telegram_api_call("getMe", timeout=60)
        bot = data.get("result", {})
        allowed = "Semua user boleh memakai bot." if not TELEGRAM_ALLOWED_USER_IDS else f"Dibatasi untuk user ID: {sorted(TELEGRAM_ALLOWED_USER_IDS)}"
        return (
            "✅ Telegram Bot terhubung.\n"
            f"Username bot: @{bot.get('username', '-')}\n"
            f"Nama bot: {bot.get('first_name', '-')}\n"
            f"Polling aktif: {telegram_bot_status.get('running')}\n"
            f"Model AI: {NVIDIA_MODEL}\n"
            f"Akses: {allowed}\n"
            f"Last update: {telegram_bot_status.get('last_update') or '-'}\n"
            f"Last error: {telegram_bot_status.get('last_error') or '-'}\n\n"
            "Tombol menu aktif. User tidak perlu mengetik command /."
        )
    except Exception as e:
        return f"❌ Telegram error: {str(e)}"


# ==========================================================
# PEMBACA IPYNB
# ==========================================================
def save_notebook_image(data_b64: str, suffix: str = ".png") -> Optional[str]:
    """Menyimpan image base64 dari output notebook menjadi file sementara."""
    try:
        if isinstance(data_b64, list):
            data_b64 = "".join(data_b64)
        image_bytes = base64.b64decode(data_b64)
        image = Image.open(io.BytesIO(image_bytes))
        tmp = tempfile.NamedTemporaryFile(delete=False, suffix=suffix)
        image.save(tmp.name)
        return tmp.name
    except Exception:
        return None


def read_ipynb(ipynb_path: str) -> List[Dict[str, Any]]:
    """Membaca cell notebook, output teks, dan gambar output."""
    notebook = nbformat.read(ipynb_path, as_version=4)
    sections: List[Dict[str, Any]] = []
    last_markdown = ""

    for cell in notebook.cells:
        if cell.cell_type == "markdown":
            text = (cell.source or "").strip()
            if text:
                last_markdown = text
                sections.append(
                    {
                        "type": "markdown",
                        "source": text,
                        "context": text,
                        "output_text": "",
                        "images": [],
                    }
                )
            continue

        if cell.cell_type != "code":
            continue

        code = (cell.source or "").strip()
        if not code:
            continue

        output_texts: List[str] = []
        image_paths: List[str] = []

        for output in cell.get("outputs", []):
            output_type = output.get("output_type")

            if output_type == "stream":
                output_texts.append(clean_text_for_docx(output.get("text", "")))

            elif output_type in ["execute_result", "display_data"]:
                data = output.get("data", {})

                if "text/plain" in data:
                    text_plain = data["text/plain"]
                    if isinstance(text_plain, list):
                        text_plain = "".join(text_plain)
                    output_texts.append(clean_text_for_docx(text_plain))

                if "image/png" in data:
                    path = save_notebook_image(data["image/png"], ".png")
                    if path:
                        image_paths.append(path)

                if "image/jpeg" in data:
                    path = save_notebook_image(data["image/jpeg"], ".jpg")
                    if path:
                        image_paths.append(path)

            elif output_type == "error":
                ename = output.get("ename", "Error")
                evalue = output.get("evalue", "")
                traceback = output.get("traceback", [])
                output_texts.append(clean_text_for_docx(f"{ename}: {evalue}\n" + "\n".join(traceback)))

        sections.append(
            {
                "type": "code",
                "source": clean_text_for_docx(code),
                "context": last_markdown,
                "output_text": clean_text_for_docx("\n".join(output_texts)).strip(),
                "images": image_paths,
            }
        )

    return sections


def fallback_title_from_code(code: str, context: str = "") -> str:
    text = f"{context}\n{code}".lower()

    rules = [
        ("import ", "Import Library"),
        ("mount", "Mount Google Drive"),
        ("zipfile", "Ekstrak Dataset"),
        ("os.listdir", "Pembacaan Dataset"),
        ("dataframe", "Memuat Data ke DataFrame"),
        ("df.info", "Cek Informasi Data"),
        ("train_test_split", "Split Dataset"),
        ("compute_class_weight", "Cek Distribusi & Class Weights"),
        ("imagedatagenerator", "Generator Inputan Data Augmentation"),
        ("flow_from_dataframe", "Flow Generator"),
        ("mobilenet", "Arsitektur Model MobileNet"),
        ("sequential", "Arsitektur Model CNN"),
        ("model.compile", "Kompilasi Model"),
        ("earlystopping", "Kompilasi & Callback Model"),
        ("model.fit", "Training Model"),
        ("classification_report", "Classification Report"),
        ("confusion_matrix", "Confusion Matrix"),
        ("model.evaluate", "Evaluasi Model"),
        ("plt.plot", "Visualisasi Loss & Accuracy"),
        ("model.save", "Simpan Model"),
        ("gradio", "Deployment dengan Gradio"),
    ]

    for keyword, title in rules:
        if keyword in text:
            return title
    return "Bagian Notebook"


def summarize_code_section(code: str, output_text: str, context: str = "", assignment_instructions: str = "") -> Tuple[str, str]:
    prompt = f"""
{REPORT_STYLE_PROMPT}

Konteks markdown sebelumnya:
```txt
{truncate_text(context, 4000)}
```

Perintah tugas tambahan dari dosen/user:
```txt
{truncate_text(assignment_instructions, 6000)}
```

Catatan:
- Jika perintah tugas tambahan kosong, cukup buat laporan berdasarkan notebook.
- Jika perintah tugas tambahan berisi instruksi format, fokus pembahasan, atau soal tugas, ikuti instruksi tersebut selama datanya ada di notebook.
- Jangan mengarang hasil yang tidak ada di kode/output.

Kode notebook:
```python
{truncate_text(code, 18000)}
```

Output notebook:
```txt
{truncate_text(output_text, 9000)}
```
"""

    try:
        result = ask_nvidia(
            [
                {"role": "system", "content": SYSTEM_PROMPT},
                {"role": "user", "content": prompt},
            ],
            temperature=0.25,
            max_tokens=1800,
        )
        data = parse_json_from_model(result)
        if data:
            judul = str(data.get("judul") or fallback_title_from_code(code, context)).strip()
            penjelasan = str(data.get("penjelasan") or "").strip()
            if not penjelasan.lower().startswith("penjelasannya"):
                penjelasan = "Penjelasannya: " + penjelasan
            return judul, penjelasan
        return fallback_title_from_code(code, context), result.strip()
    except Exception:
        judul = fallback_title_from_code(code, context)
        penjelasan = (
            "Penjelasannya: Bagian kode ini digunakan dalam proses pembuatan model Deep Learning. "
            "Kode perlu dibaca bersama output notebook untuk mengetahui hasil yang diperoleh secara lebih detail."
        )
        return judul, penjelasan


# ==========================================================
# PEMBUAT DOCX
# ==========================================================
def set_document_style(doc: Document):
    style = doc.styles["Normal"]
    style.font.name = "Times New Roman"
    style.font.size = Pt(12)

    for section in doc.sections:
        section.top_margin = Inches(1)
        section.bottom_margin = Inches(0.8)
        section.left_margin = Inches(1.2)
        section.right_margin = Inches(1)


def add_center_paragraph(doc: Document, text: str, bold: bool = True, size: int = 12):
    p = doc.add_paragraph()
    p.alignment = WD_ALIGN_PARAGRAPH.CENTER
    run = p.add_run(clean_text_for_docx(text))
    run.bold = bold
    run.font.name = "Times New Roman"
    run.font.size = Pt(size)
    return p


def add_code_block(doc: Document, code: str):
    """
    Memasukkan kode/output ke DOCX tanpa dipangkas.
    Teks dibersihkan dari karakter ilegal XML dan dipecah ke beberapa paragraf agar aman.
    """
    code = clean_text_for_docx(code)
    if not code.strip():
        return

    table = doc.add_table(rows=1, cols=1)
    table.style = "Table Grid"
    cell = table.cell(0, 0)

    chunks = split_text_chunks(code, chunk_size=5000)

    for idx, chunk in enumerate(chunks):
        paragraph = cell.paragraphs[0] if idx == 0 else cell.add_paragraph()
        run = paragraph.add_run(chunk)
        run.font.name = "Courier New"
        run.font.size = Pt(8)


def add_output_block(doc: Document, output_text: str):
    if not output_text.strip():
        return
    p = doc.add_paragraph()
    r = p.add_run("Output:")
    r.bold = True
    add_code_block(doc, output_text)


def add_image_to_doc(doc: Document, image_path: str):
    try:
        p = doc.add_paragraph()
        p.alignment = WD_ALIGN_PARAGRAPH.CENTER
        run = p.add_run()
        run.add_picture(image_path, width=Inches(5.6))
    except Exception:
        # Jika gambar gagal dimasukkan, lewati agar pembuatan laporan tetap jalan.
        pass


def create_cover(
    doc: Document,
    judul_laporan: str,
    nama: str,
    nim: str,
    dosen: str,
    kelas: str,
    anggota: str,
    prodi: str,
    kampus: str,
    tahun: str,
):
    add_center_paragraph(doc, "Laporan Deep Learning", bold=True, size=16)
    add_center_paragraph(doc, judul_laporan.upper(), bold=True, size=14)

    doc.add_paragraph("\n\n\n")

    add_center_paragraph(doc, nama, bold=True)
    add_center_paragraph(doc, f"NIM. {nim}", bold=True)

    doc.add_paragraph("")
    add_center_paragraph(doc, "Dosen Mata kuliah", bold=True)
    add_center_paragraph(doc, dosen, bold=True)
    add_center_paragraph(doc, kelas, bold=True)

    if anggota.strip():
        doc.add_paragraph("")
        add_center_paragraph(doc, anggota.upper(), bold=True)

    doc.add_paragraph("\n\n")

    add_center_paragraph(doc, prodi.upper(), bold=True)
    add_center_paragraph(doc, kampus.upper(), bold=True)
    add_center_paragraph(doc, tahun, bold=True)

    doc.add_page_break()


def add_footer(doc: Document, kampus: str):
    for section in doc.sections:
        footer = section.footer
        p = footer.paragraphs[0]
        p.alignment = WD_ALIGN_PARAGRAPH.RIGHT
        run = p.add_run(clean_text_for_docx(kampus))
        run.italic = True
        run.font.size = Pt(9)


def create_report_docx(
    ipynb_file: Any,
    judul_laporan: str,
    nama: str,
    nim: str,
    dosen: str,
    kelas: str,
    anggota: str,
    prodi: str,
    kampus: str,
    tahun: str,
):
    if not NVIDIA_API_KEY:
        raise gr.Error("NVIDIA_API_KEY belum diset di Hugging Face Secrets.")

    ipynb_path = get_uploaded_path(ipynb_file)
    if not ipynb_path.endswith(".ipynb"):
        raise gr.Error("File harus berformat .ipynb")

    sections = read_ipynb(ipynb_path)
    code_sections = [item for item in sections if item.get("type") == "code"]

    if not code_sections:
        raise gr.Error("Notebook tidak memiliki cell kode yang bisa dibuat menjadi laporan.")

    doc = Document()
    set_document_style(doc)
    add_footer(doc, kampus)

    create_cover(
        doc=doc,
        judul_laporan=judul_laporan,
        nama=nama,
        nim=nim,
        dosen=dosen,
        kelas=kelas,
        anggota=anggota,
        prodi=prodi,
        kampus=kampus,
        tahun=tahun,
    )

    number = 1
    total = len(code_sections)

    for item in code_sections:
        code = item.get("source", "")
        output_text = item.get("output_text", "")
        context = item.get("context", "")
        images = item.get("images", [])

        judul, penjelasan = summarize_code_section(code, output_text, context)

        heading = doc.add_paragraph()
        heading_run = heading.add_run(clean_text_for_docx(f"{number}. {judul}"))
        heading_run.bold = True
        heading_run.font.name = "Times New Roman"
        heading_run.font.size = Pt(13)

        add_code_block(doc, code)
        add_output_block(doc, output_text)

        for image_path in images[:4]:
            add_image_to_doc(doc, image_path)

        p = doc.add_paragraph()
        p.alignment = WD_ALIGN_PARAGRAPH.JUSTIFY
        r = p.add_run(clean_text_for_docx(penjelasan))
        r.font.name = "Times New Roman"
        r.font.size = Pt(12)

        if number < total:
            doc.add_paragraph("")
        number += 1

    output_name = safe_filename(judul_laporan)
    output_path = str(Path(tempfile.gettempdir()) / f"{output_name}.docx")
    doc.save(output_path)
    return output_path




def get_uploaded_paths(uploaded_files: Any) -> List[str]:
    """
    Mendukung input dari gr.File atau gr.Files:
    - single file object
    - list file object
    - string path
    """
    if uploaded_files is None:
        return []

    if isinstance(uploaded_files, (list, tuple)):
        result = []
        for item in uploaded_files:
            path = get_uploaded_path(item)
            if path:
                result.append(path)
        return result

    path = get_uploaded_path(uploaded_files)
    return [path] if path else []


# ==========================================================
# WORKFLOW REVIEW & EDIT SEBELUM DOCX
# ==========================================================
def build_report_draft_text(
    code_sections: List[Dict[str, Any]],
    assignment_instructions: str = "",
    include_code: bool = True,
    include_output: bool = True,
    include_images: bool = True,
) -> Tuple[str, Dict[str, str]]:
    """
    Membuat draft laporan dalam bentuk teks editable.
    Placeholder gambar seperti [GAMBAR_OUTPUT_1_1] akan diganti menjadi gambar asli saat export DOCX.
    """
    parts: List[str] = []
    image_map: Dict[str, str] = {}

    if assignment_instructions.strip():
        parts.append("Perintah Tugas")
        parts.append(assignment_instructions.strip())
        parts.append("")

    number = 1
    for item in code_sections:
        code = item.get("source", "")
        output_text = item.get("output_text", "")
        context = item.get("context", "")
        images = item.get("images", [])

        judul, penjelasan = summarize_code_section(
            code=code,
            output_text=output_text,
            context=context,
            assignment_instructions=assignment_instructions,
        )

        parts.append(f"{number}. {judul}")

        if include_code:
            parts.append("Kode:")
            parts.append("```python")
            parts.append(clean_text_for_docx(code))
            parts.append("```")

        if include_output and output_text.strip():
            parts.append("Output:")
            parts.append("```txt")
            parts.append(clean_text_for_docx(output_text))
            parts.append("```")

        if include_images:
            for image_index, image_path in enumerate(images[:4], start=1):
                placeholder = f"[GAMBAR_OUTPUT_{number}_{image_index}]"
                image_map[placeholder] = image_path
                parts.append(placeholder)

        parts.append(penjelasan.strip())
        parts.append("")
        number += 1

    return "\n".join(parts).strip(), image_map


def generate_report_review(
    ipynb_file: Any,
    judul_laporan: str,
    nama: str,
    nim: str,
    dosen: str,
    kelas: str,
    anggota: str,
    prodi: str,
    kampus: str,
    tahun: str,
    assignment_instructions: str,
    include_code: bool,
    include_output: bool,
    include_images: bool,
):
    """
    Tahap 1:
    Upload satu atau beberapa IPYNB -> AI membuat draft -> user bisa review dan edit dulu.
    """
    if not NVIDIA_API_KEY:
        raise gr.Error("NVIDIA_API_KEY belum diset di Hugging Face Secrets.")

    ipynb_paths = get_uploaded_paths(ipynb_file)

    if not ipynb_paths:
        raise gr.Error("Upload minimal satu file .ipynb.")

    invalid_files = [path for path in ipynb_paths if not str(path).endswith(".ipynb")]
    if invalid_files:
        raise gr.Error("Semua file harus berformat .ipynb.")

    all_code_sections: List[Dict[str, Any]] = []
    notebook_names: List[str] = []

    for ipynb_path in ipynb_paths:
        notebook_name = Path(ipynb_path).name
        notebook_names.append(notebook_name)

        sections = read_ipynb(ipynb_path)
        code_sections = [item for item in sections if item.get("type") == "code"]

        for item in code_sections:
            item["context"] = f"Notebook: {notebook_name}\n" + str(item.get("context", ""))

        all_code_sections.extend(code_sections)

    if not all_code_sections:
        raise gr.Error("Notebook tidak memiliki cell kode yang bisa dibuat menjadi laporan.")

    multi_file_note = ""
    if len(notebook_names) > 1:
        multi_file_note = (
            "Laporan ini dibuat dari beberapa file notebook berikut:\n"
            + "\n".join(f"- {name}" for name in notebook_names)
        )

    combined_assignment = assignment_instructions.strip()
    if multi_file_note:
        combined_assignment = (multi_file_note + "\n\n" + combined_assignment).strip()

    draft_text, image_map = build_report_draft_text(
        code_sections=all_code_sections,
        assignment_instructions=combined_assignment,
        include_code=include_code,
        include_output=include_output,
        include_images=include_images,
    )

    state = {
        "image_map": image_map,
        "judul_laporan": judul_laporan,
        "nama": nama,
        "nim": nim,
        "dosen": dosen,
        "kelas": kelas,
        "anggota": anggota,
        "prodi": prodi,
        "kampus": kampus,
        "tahun": tahun,
        "assignment_instructions": combined_assignment,
        "notebook_files": notebook_names,
    }

    status = (
        f"✅ Draft review berhasil dibuat dari {len(notebook_names)} file notebook. "
        f"Terdapat {len(all_code_sections)} bagian kode. "
        "Silakan edit teks di kotak review. Kode dan output tidak dipangkas di draft. "
        "Setelah sesuai, klik Export ke DOCX."
    )
    return draft_text, state, status


def add_heading_like_report(doc: Document, text: str, level: int = 1):
    text = clean_text_for_docx(text).strip()
    if not text:
        return None
    p = doc.add_paragraph()
    run = p.add_run(text)
    run.bold = True
    run.font.name = "Times New Roman"
    run.font.size = Pt(13 if level <= 1 else 12)
    return p


def add_normal_report_paragraph(doc: Document, text: str):
    text = clean_text_for_docx(text).strip()
    if not text:
        return
    p = doc.add_paragraph()
    p.alignment = WD_ALIGN_PARAGRAPH.JUSTIFY
    run = p.add_run(text)
    run.font.name = "Times New Roman"
    run.font.size = Pt(12)


def flush_paragraph_buffer(doc: Document, paragraph_buffer: List[str]):
    if not paragraph_buffer:
        return
    text = " ".join(line.strip() for line in paragraph_buffer if line.strip())
    paragraph_buffer.clear()
    if text:
        add_normal_report_paragraph(doc, text)


def add_edited_draft_to_docx(doc: Document, edited_draft: str, image_map: Dict[str, str]):
    """
    Mengubah draft teks editable menjadi isi DOCX.
    Mendukung:
    - heading seperti "1. Import Library" atau "# Judul"
    - code fence ```python / ```txt
    - placeholder gambar [GAMBAR_OUTPUT_1_1]
    - paragraf biasa
    """
    lines = edited_draft.splitlines()
    in_code = False
    code_lines: List[str] = []
    paragraph_buffer: List[str] = []

    for raw_line in lines:
        line = clean_text_for_docx(raw_line.rstrip("\n"))
        stripped = line.strip()

        if stripped.startswith("```"):
            if not in_code:
                flush_paragraph_buffer(doc, paragraph_buffer)
                in_code = True
                code_lines = []
            else:
                add_code_block(doc, "\n".join(code_lines).strip())
                in_code = False
                code_lines = []
            continue

        if in_code:
            code_lines.append(line)
            continue

        if not stripped:
            flush_paragraph_buffer(doc, paragraph_buffer)
            continue

        if re.fullmatch(r"\[GAMBAR_OUTPUT_\d+_\d+\]", stripped):
            flush_paragraph_buffer(doc, paragraph_buffer)
            image_path = image_map.get(stripped)
            if image_path:
                add_image_to_doc(doc, image_path)
            else:
                add_normal_report_paragraph(doc, stripped)
            continue

        if stripped.lower() in ["kode:", "output:"]:
            flush_paragraph_buffer(doc, paragraph_buffer)
            p = doc.add_paragraph()
            r = p.add_run(stripped)
            r.bold = True
            r.font.name = "Times New Roman"
            r.font.size = Pt(12)
            continue

        if stripped.startswith("#"):
            flush_paragraph_buffer(doc, paragraph_buffer)
            heading_text = stripped.lstrip("#").strip()
            if heading_text:
                add_heading_like_report(doc, heading_text)
            continue

        if re.match(r"^\d+(\.\d+)?\.?\s+.+", stripped) and len(stripped) <= 140:
            flush_paragraph_buffer(doc, paragraph_buffer)
            add_heading_like_report(doc, stripped)
            continue

        # "Perintah Tugas" dibuat heading agar rapi
        if stripped.lower() == "perintah tugas":
            flush_paragraph_buffer(doc, paragraph_buffer)
            add_heading_like_report(doc, stripped)
            continue

        paragraph_buffer.append(stripped)

    if in_code and code_lines:
        add_code_block(doc, "\n".join(code_lines).strip())

    flush_paragraph_buffer(doc, paragraph_buffer)


def export_review_to_docx(
    edited_draft: str,
    report_state: Dict[str, Any],
    judul_laporan: str,
    nama: str,
    nim: str,
    dosen: str,
    kelas: str,
    anggota: str,
    prodi: str,
    kampus: str,
    tahun: str,
):
    """
    Tahap 2:
    User sudah edit draft -> export menjadi DOCX.
    """
    if not edited_draft or not edited_draft.strip():
        raise gr.Error("Draft review masih kosong. Buat draft dulu atau isi teks laporan.")

    image_map = {}
    if isinstance(report_state, dict):
        image_map = report_state.get("image_map", {}) or {}

    doc = Document()
    set_document_style(doc)
    add_footer(doc, kampus)

    create_cover(
        doc=doc,
        judul_laporan=judul_laporan,
        nama=nama,
        nim=nim,
        dosen=dosen,
        kelas=kelas,
        anggota=anggota,
        prodi=prodi,
        kampus=kampus,
        tahun=tahun,
    )

    add_edited_draft_to_docx(doc, edited_draft, image_map)

    output_name = safe_filename(judul_laporan)
    output_path = str(Path(tempfile.gettempdir()) / f"{output_name}_final.docx")
    doc.save(output_path)
    return output_path



# ==========================================================
# AI EDIT DRAFT SEBELUM EXPORT DOCX
# ==========================================================
def ai_update_review_draft(
    edited_draft: str,
    report_state: Dict[str, Any],
    edit_instruction: str,
    edit_mode: str,
):
    """
    Memungkinkan user menyuruh AI menambah atau mengubah draft setelah tahap review.
    Output tetap kembali ke kotak draft, jadi user bisa cek lagi sebelum export DOCX.
    """
    if not NVIDIA_API_KEY:
        raise gr.Error("NVIDIA_API_KEY belum diset di Hugging Face Secrets.")

    edited_draft = clean_text_for_docx(edited_draft)
    edit_instruction = clean_text_for_docx(edit_instruction)

    if not edited_draft.strip():
        raise gr.Error("Draft masih kosong. Buat draft review terlebih dahulu.")

    if not edit_instruction.strip():
        raise gr.Error("Isi dulu perintah edit AI, misalnya: tambahkan kesimpulan dan saran.")

    assignment_context = ""
    notebook_files = []
    if isinstance(report_state, dict):
        assignment_context = report_state.get("assignment_instructions", "") or ""
        notebook_files = report_state.get("notebook_files", []) or []

    context_note = ""
    if notebook_files:
        context_note += "File notebook sumber:\n" + "\n".join(f"- {x}" for x in notebook_files) + "\n\n"
    if assignment_context:
        context_note += "Perintah tugas awal:\n" + assignment_context + "\n\n"

    preserve_rules = """
Aturan edit:
- Gunakan bahasa Indonesia formal seperti laporan mahasiswa.
- Jangan mengarang angka, akurasi, loss, jumlah dataset, nama kelas, atau hasil evaluasi yang tidak ada.
- Jangan menghapus placeholder gambar seperti [GAMBAR_OUTPUT_1_1], kecuali user secara jelas meminta menghapusnya.
- Jangan menghapus blok kode/output yang sudah ada, kecuali user secara jelas meminta menghapus atau meringkasnya.
- Pertahankan format judul bernomor dan pola "Penjelasannya:".
- Balas hanya dengan isi draft, tanpa pembuka seperti "Berikut hasilnya".
"""

    mode = (edit_mode or "").lower()

    # Mode tambah: AI membuat bagian tambahan saja, lalu sistem append ke draft.
    if "tambah" in mode:
        prompt = f"""
Kamu diminta MENAMBAHKAN bagian baru ke draft laporan.

{preserve_rules}

Konteks laporan:
```txt
{truncate_text(context_note, 6000)}
```

Perintah user untuk tambahan:
```txt
{edit_instruction}
```

Draft laporan saat ini hanya untuk konteks, JANGAN diulang:
```txt
{truncate_text(edited_draft, 25000)}
```

Tugas:
- Buat hanya bagian tambahan yang perlu dimasukkan.
- Gunakan format laporan, boleh memakai judul bernomor jika cocok.
- Jika user meminta kesimpulan/saran/abstrak/ringkasan, buat bagian tersebut.
- Jangan ulangi seluruh draft.
"""

        addition = ask_nvidia(
            [
                {"role": "system", "content": SYSTEM_PROMPT},
                {"role": "user", "content": prompt},
            ],
            temperature=0.35,
            max_tokens=3000,
        )

        addition = clean_text_for_docx(addition).strip()
        if not addition:
            raise gr.Error("AI tidak menghasilkan tambahan. Coba tulis instruksi lebih jelas.")

        updated = edited_draft.rstrip() + "\n\n" + addition + "\n"
        status = "✅ AI berhasil menambahkan bagian baru ke draft. Silakan review lagi sebelum export DOCX."
        return updated, status

    # Mode ubah penuh: AI mengembalikan draft revisi.
    # Batas dibuat untuk mencegah request terlalu besar.
    draft_limit = 90000
    draft_for_model = edited_draft
    warning = ""

    if len(edited_draft) > draft_limit:
        draft_for_model = edited_draft[:draft_limit]
        warning = (
            "⚠️ Draft sangat panjang, jadi AI hanya menerima bagian awal draft untuk proses ubah penuh. "
            "Untuk draft besar, lebih aman pakai mode 'Tambah bagian baru di akhir' atau edit manual bagian tertentu."
        )

    prompt = f"""
Kamu diminta MENGUBAH draft laporan yang sudah ada sesuai instruksi user.

{preserve_rules}

Konteks laporan:
```txt
{truncate_text(context_note, 6000)}
```

Instruksi perubahan dari user:
```txt
{edit_instruction}
```

Draft laporan yang harus direvisi:
```txt
{draft_for_model}
```

Tugas:
- Kembalikan draft hasil revisi secara utuh.
- Terapkan instruksi user.
- Jika instruksi user hanya meminta mengubah bagian tertentu, ubah bagian itu saja dan pertahankan bagian lain.
- Jangan menghapus kode, output, atau placeholder gambar kecuali diminta.
- Jangan menambahkan data palsu.
"""

    revised = ask_nvidia(
        [
            {"role": "system", "content": SYSTEM_PROMPT},
            {"role": "user", "content": prompt},
        ],
        temperature=0.25,
        max_tokens=12000,
    )

    revised = clean_text_for_docx(revised).strip()
    if not revised:
        raise gr.Error("AI tidak menghasilkan revisi. Coba instruksi yang lebih jelas.")

    status = "✅ AI berhasil mengubah draft. Silakan review lagi sebelum export DOCX."
    if warning:
        status = warning + "\n" + status

    return revised, status


# ==========================================================
# UI GRADIO
# ==========================================================
with gr.Blocks(title="AI Assistant Mahasiswa NVIDIA") as demo:
    gr.Markdown("# 🎓 AI Assistant Mahasiswa NVIDIA")
    gr.Markdown(
        "Asisten AI untuk chat umum mahasiswa dan pembuatan laporan Deep Learning otomatis dari file `.ipynb` ke `.docx`."
    )

    with gr.Tab("Chat Umum"):
        gr.Markdown("Gunakan tab ini untuk bertanya tentang materi, tugas, coding, laporan, atau pertanyaan umum.")
        gr.ChatInterface(
            fn=normal_chat,
            title="Chat AI Mahasiswa",
            description="Tanya apa saja seputar kuliah, coding, laporan, Deep Learning, atau pertanyaan umum.",
        )

    with gr.Tab("Buat Laporan dari IPYNB"):
        gr.Markdown(
            "Upload satu atau beberapa file `.ipynb`, isi data cover, tambahkan **perintah tugas** bila ada, "
            "lalu buat draft review. Setelah draft diedit, baru export menjadi file Word `.docx`."
        )

        report_state = gr.State({})

        with gr.Row():
            with gr.Column(scale=1):
                ipynb_file = gr.Files(label="Upload File IPYNB Bisa Lebih dari Satu", file_types=[".ipynb"])
                judul_laporan = gr.Textbox(label="Judul Laporan", value="Klasifikasi Penyakit Mata")
                nama = gr.Textbox(label="Nama", value="Hadid Zarid Nawfal")
                nim = gr.Textbox(label="NIM", value="2355301079")
                dosen = gr.Textbox(label="Dosen Mata Kuliah", value="Dr. Juni Nurma Sari, S.Kom., M.MT.")
                kelas = gr.Textbox(label="Kelas", value="ILB")
                anggota = gr.Textbox(label="Nama Anggota / Partner Opsional", value="")
                prodi = gr.Textbox(label="Program Studi", value="Program Studi Teknik Informatika")
                kampus = gr.Textbox(label="Kampus", value="Politeknik Caltex Riau")
                tahun = gr.Textbox(label="Tahun", value="2025")

                assignment_instructions = gr.Textbox(
                    label="Perintah Tugas Tambahan",
                    placeholder=(
                        "Contoh: Buat laporan sesuai instruksi dosen. Jelaskan preprocessing, arsitektur model, "
                        "hasil training, evaluasi, confusion matrix, dan kesimpulan. Gunakan gaya bahasa formal."
                    ),
                    lines=6,
                )

                with gr.Accordion("Pengaturan isi draft", open=False):
                    include_code = gr.Checkbox(label="Masukkan kode notebook ke draft", value=True)
                    include_output = gr.Checkbox(label="Masukkan output teks notebook ke draft", value=True)
                    include_images = gr.Checkbox(label="Masukkan gambar/grafik output notebook", value=True)

                review_btn = gr.Button("1. Buat Draft Review", variant="primary")
                export_btn = gr.Button("2. Export ke DOCX Setelah Edit", variant="secondary")

            with gr.Column(scale=2):
                status_box = gr.Textbox(label="Status", lines=3, interactive=False)
                draft_box = gr.Textbox(
                    label="Review & Edit Isi Laporan",
                    placeholder=(
                        "Draft laporan akan muncul di sini. Kamu bisa edit judul bagian, penjelasan, "
                        "kode, output, atau hapus bagian yang tidak diperlukan sebelum export DOCX."
                    ),
                    lines=30,
                    interactive=True,
                )
                with gr.Accordion("Minta AI Tambahkan / Ubah Draft", open=True):
                    edit_instruction = gr.Textbox(
                        label="Perintah Edit AI",
                        placeholder=(
                            "Contoh: Tambahkan bagian kesimpulan dan saran. "
                            "Atau: Ubah penjelasan bagian training agar lebih detail dan formal."
                        ),
                        lines=4,
                    )
                    edit_mode = gr.Radio(
                        label="Mode Edit AI",
                        choices=[
                            "Tambah bagian baru di akhir",
                            "Ubah draft penuh sesuai instruksi",
                        ],
                        value="Tambah bagian baru di akhir",
                    )
                    ai_edit_btn = gr.Button("Minta AI Tambahkan/Ubah Draft", variant="primary")

                output_file = gr.File(label="Download Laporan DOCX", file_types=[".docx"])
                gr.Markdown(
                    "**Catatan:** Notebook sebaiknya sudah dijalankan terlebih dahulu agar output, grafik, "
                    "akurasi, loss, dan hasil evaluasi tersimpan di file `.ipynb`. "
                    "Setelah AI mengubah draft, cek lagi isinya sebelum export ke DOCX."
                )

        review_btn.click(
            fn=generate_report_review,
            inputs=[
                ipynb_file,
                judul_laporan,
                nama,
                nim,
                dosen,
                kelas,
                anggota,
                prodi,
                kampus,
                tahun,
                assignment_instructions,
                include_code,
                include_output,
                include_images,
            ],
            outputs=[draft_box, report_state, status_box],
        )

        ai_edit_btn.click(
            fn=ai_update_review_draft,
            inputs=[
                draft_box,
                report_state,
                edit_instruction,
                edit_mode,
            ],
            outputs=[draft_box, status_box],
        )

        export_btn.click(
            fn=export_review_to_docx,
            inputs=[
                draft_box,
                report_state,
                judul_laporan,
                nama,
                nim,
                dosen,
                kelas,
                anggota,
                prodi,
                kampus,
                tahun,
            ],
            outputs=output_file,
        )

    with gr.Tab("Telegram Bot"):
        gr.Markdown("Hubungkan aplikasi ini ke Telegram supaya AI bisa dipakai dari chat Telegram, termasuk upload `.ipynb`, buat draft, edit draft, dan export DOCX lewat tombol.")
        gr.Markdown(
            "Buat bot melalui **@BotFather**, lalu simpan token bot sebagai Secret di Hugging Face dengan nama `TELEGRAM_BOT_TOKEN`. "
            "Setelah itu lakukan Restart Space."
        )
        gr.Markdown(
            "**Secret/Variable Telegram:**\n"
            "- `TELEGRAM_BOT_TOKEN` → Secret wajib untuk mengaktifkan bot.\n"
            "- `TELEGRAM_ENABLED=true` → Variable opsional.\n"
            "- `TELEGRAM_ALLOWED_USER_IDS=123456789,987654321` → Variable opsional untuk membatasi user.\n"
            "- `TELEGRAM_MAX_HISTORY=8` → Variable opsional untuk jumlah riwayat chat.\n- `TELEGRAM_REPORT_NAMA`, `TELEGRAM_REPORT_NIM`, dll. → Variable opsional untuk default cover."
        )
        telegram_check_btn = gr.Button("Cek Status Telegram Bot")
        telegram_output = gr.Textbox(label="Status Telegram", lines=10)
        telegram_check_btn.click(fn=telegram_status_text, inputs=None, outputs=telegram_output)

    with gr.Tab("Cek API"):
        gr.Markdown("Gunakan tab ini untuk mengecek apakah NVIDIA API key dan model sudah benar.")
        gr.Markdown(f"**Model saat ini:** `{NVIDIA_MODEL}`")
        gr.Markdown(f"**Base URL:** `{NVIDIA_BASE_URL}`")
        check_btn = gr.Button("Cek Koneksi NVIDIA API")
        check_output = gr.Textbox(label="Status", lines=4)
        check_btn.click(fn=test_api_connection, inputs=None, outputs=check_output)

if __name__ == "__main__":
    start_telegram_bot_if_configured()
    demo.queue(max_size=20).launch(server_name="0.0.0.0", server_port=7860, share=False, show_api=False)