ui.py

import os
import sys
import platform
import json
import shutil
from pathlib import Path
from typing import Optional

import torch

from rich.panel import Panel
from rich.prompt import Prompt, Confirm, FloatPrompt
from rich.table import Table
from rich import box
from rich.markup import escape as rich_escape

from config import (
    MODEL_DIR, PLOT_DIR, EXPORT_DIR, DATA_DIR,
    DEFAULT_MEMORY_LIMIT_GB, DEFAULT_MODEL
)
from utils import console, Theme, set_memory_hard_limit
from models import ModelManager
from history import TrainingHistory
from data_loader import EnhancedDatasetLoader, DatasetStats
from trainer import EnhancedTrainingModule, get_gpu_info, get_device
from chat import EnhancedChatModule
from backup_module import backup_menu

try:
    HAS_PSUTIL = True
    import psutil
except ImportError:
    HAS_PSUTIL = False

try:
    HAS_MATPLOTLIB = True
    import matplotlib
    matplotlib.use('Agg')
    import matplotlib.pyplot as plt
except ImportError:
    HAS_MATPLOTLIB = False

try:
    HAS_PANDAS = True
    import pandas as pd
except ImportError:
    HAS_PANDAS = False

try:
    HAS_EVALUATE = True
    import evaluate
except ImportError:
    HAS_EVALUATE = False

try:
    HAS_PEFT = True
    from peft import PeftModel
except ImportError:
    HAS_PEFT = False

try:
    HAS_BNB = True
    import bitsandbytes as bnb
except ImportError:
    HAS_BNB = False

try:
    HAS_SACREBLEU = True
    import sacrebleu
except ImportError:
    HAS_SACREBLEU = False


def logo() -> None:
    console.print(Panel(
        """[bold red]
  ████████  ███████   █████   ██ █████   ██ ██   ██
  ██░░░░░░  ██░░░░██ ██░░░██ ██ ██░░░██ ██ ░██ ░██
  ██        ██   ░██░██  ░░ ██ ██   ░██ ██  ░████ 
  ███████   ███████ ░██   ████ ███████ ██   ░██░  
  ██░░░░    ██░░░░  ░░██ ████ ██░░░░  ██    ██░   
  ██        ██       ░░████  ██      ██     ██░   
  ████████  ██        ░░██   ██      ██     ██░   
  ░░░░░░░░  ░░         ░░    ░░      ░░     ░░    
""",
        width=55, title="[bold yellow]ULTIMATE AI CHATBOT v2.5 FIXED", style="bold yellow"
    ))


def manage_models() -> None:
    model_manager = ModelManager()

    console.print(Panel(Theme.header(" MANAJEMEN MODEL"), title="MODEL", style="yellow"))

    models = model_manager.list_models()

    if not models:
        console.print(Theme.warning("Belum ada model"))
        return

    table = Table(title="Daftar Model", box=box.ROUNDED)
    table.add_column("No", style="cyan")
    table.add_column("ID", style="green")
    table.add_column("Base Model", style="white")
    table.add_column("Dataset", style="yellow")
    table.add_column("Samples", style="blue")
    table.add_column("Status", style="magenta")
    table.add_column("Created", style="dim")
    table.add_column("Metrics", style="cyan")

    current = model_manager.get_current_model()

    for idx, (key, val) in enumerate(models.items(), 1):
        status = " Active" if key == current else "Available"

        metrics = val.get("metrics", {})
        if metrics.get('bleu'):
            metric_str = f"BLEU: {metrics['bleu']:.3f}"
        elif metrics.get('perplexity'):
            metric_str = f"PPL: {metrics['perplexity']:.2f}"
        elif metrics.get('bleu_sacrebleu'):
            metric_str = f"BLEU: {metrics['bleu_sacrebleu']:.3f}"
        else:
            metric_str = "-"

        samples = val.get('training_samples', '?')
        if isinstance(samples, int):
            samples = f"{samples:,}"

        table.add_row(
            str(idx),
            key,
            val.get("base_model", "?"),
            val.get("dataset_format", "?"),
            samples,
            status,
            val.get("created_date", "-")[:10],
            metric_str
        )

    console.print(table)

    console.print("\n" + Theme.info("Pilihan:"))
    console.print("  [green]1. Set model aktif[/green]")
    console.print("  [green]2. Lihat detail model[/green]")
    console.print("  [green]3. Hapus model[/green]")
    console.print("  [green]4. Export model[/green]")
    console.print("  [green]5. Kembali[/green]")

    choice = Prompt.ask("[yellow]Pilih", choices=['1', '2', '3', '4', '5'])

    if choice == '1':
        model_id = Prompt.ask("[cyan]ID model", choices=list(models.keys()))
        model_manager.set_current_model(model_id)
        console.print(Theme.success(f" Model aktif: {model_id}"))

    elif choice == '2':
        model_id = Prompt.ask("[cyan]ID model", choices=list(models.keys()))
        info = models.get(model_id)

        if info:
            table = Table(title=f"Detail: {model_id}", box=box.ROUNDED)
            table.add_column("Key", style="cyan")
            table.add_column("Value", style="green")

            for k, v in info.items():
                if k not in ['metrics', 'config']:
                    table.add_row(k, str(v)[:200])

            console.print(table)

            if 'metrics' in info and info['metrics']:
                console.print("\n" + Theme.header(" Metrics:"))
                for k, v in info['metrics'].items():
                    if isinstance(v, float):
                        console.print(f"  {k}: {v:.4f}")
                    else:
                        console.print(f"  {k}: {v}")

    elif choice == '3':
        model_id = Prompt.ask("[cyan]ID model yang akan dihapus", choices=list(models.keys()))
        if Confirm.ask(Theme.error(f"Hapus {model_id}?"), default=False):
            if model_manager.delete_model(model_id):
                console.print(Theme.success(" Model dihapus"))
            else:
                console.print(Theme.error("Gagal hapus"))

    elif choice == '4':
        model_id = Prompt.ask("[cyan]ID model yang akan diexport", choices=list(models.keys()))
        model_path = os.path.join(MODEL_DIR, model_id)

        if os.path.exists(model_path):
            try:
                export_dir = Path(EXPORT_DIR) / model_id
                shutil.copytree(model_path, export_dir)
                console.print(Theme.success(f" Model exported to {export_dir}"))
            except Exception as e:
                console.print(Theme.error(f"Export failed: {e}"))


def dataset_inspector() -> None:
    console.print(Panel(Theme.header(" DATASET INSPECTOR"), title="INSPECTOR", style="bold yellow"))

    filepath = select_dataset_file()
    if filepath is None:
        return

    if not os.path.exists(filepath):
        console.print(Theme.error("File tidak ditemukan!"))
        return

    console.print(Theme.info("Analyzing dataset..."))
    try:
        loader = EnhancedDatasetLoader()
        samples, stats = loader.load(filepath)
    except Exception as e:
        console.print(Theme.error(f"Error: {rich_escape(str(e))}"))
        return

    console.print("\n" + Theme.header(" Dataset Statistics:"))
    table = Table(title="Dataset Info", box=box.ROUNDED)
    table.add_column("Metric", style="cyan")
    table.add_column("Value", style="green")

    for key, value in stats.__dict__.items():
        if key not in ['warnings', 'conversation_pairs', 'word_frequency', 'char_frequency']:
            if isinstance(value, float):
                table.add_row(key, f"{value:.2f}")
            else:
                table.add_row(key, str(value))

    console.print(table)

    if stats.warnings:
        console.print("\n" + Theme.warning("Warnings:"))
        for warn in stats.warnings[:5]:
            console.print(f"  {warn}")

    if samples:
        console.print("\n" + Theme.header(" Sample Texts:"))
        for idx, text in enumerate(samples[:3], 1):
            preview = text[:200] + "..." if len(text) > 200 else text
            console.print(f"[green][{idx}][/green]\n{rich_escape(preview)}\n")

    if stats.word_frequency:
        console.print("\n" + Theme.info(" Top Words:"))
        word_table = Table(title="Word Frequency", box=box.SIMPLE)
        word_table.add_column("Word", style="cyan")
        word_table.add_column("Count", style="green")
        for word, count in list(stats.word_frequency.items())[:20]:
            word_table.add_row(word, str(count))
        console.print(word_table)


def visualize_history_menu() -> None:
    model_manager = ModelManager()
    models = model_manager.list_models()

    if not models:
        console.print(Theme.warning("Tidak ada model untuk divisualisasi"))
        return

    console.print(Theme.info("Pilih model:"))
    for idx, (mid, _) in enumerate(models.items(), 1):
        console.print(f"  {idx}. {mid}")

    choice = Prompt.ask("[yellow]Pilih", default="1")
    try:
        model_id = list(models.keys())[int(choice)-1]
    except (ValueError, IndexError):
        console.print(Theme.error("Pilihan invalid"))
        return

    model_path = os.path.join(MODEL_DIR, model_id)
    history_path = os.path.join(model_path, "training_history.json")

    if not os.path.exists(history_path):
        console.print(Theme.warning("Tidak ada history untuk model ini"))
        return

    try:
        with open(history_path, 'r') as f:
            history_data = json.load(f)

        history = TrainingHistory.from_dict(history_data)
        history.print_summary()

        if HAS_MATPLOTLIB:
            if Confirm.ask("[yellow]Generate plot?", default=True):
                plot_path = os.path.join(PLOT_DIR, f"{model_id}_history.png")
                history.plot(plot_path)

        overfit = history.detect_overfitting()
        if overfit['overfitting_detected']:
            console.print(Theme.error(overfit['warning']))

    except Exception as e:
        console.print(Theme.error(f"Error reading history: {e}"))


def enhanced_system_stats() -> None:
    console.print(Panel(Theme.header(" SYSTEM STATISTICS"), title="STATS", style="yellow"))

    table = Table(title="System Information", box=box.ROUNDED)
    table.add_column("Info", style="cyan")
    table.add_column("Value", style="green")

    table.add_row("Python Version", platform.python_version())
    table.add_row("PyTorch Version", torch.__version__)

    try:
        import transformers
        table.add_row("Transformers Version", transformers.__version__)
    except:
        pass

    try:
        import datasets
        table.add_row("Datasets Version", datasets.__version__)
    except:
        pass

    cuda_available = torch.cuda.is_available()
    table.add_row("CUDA Available", "Yes" if cuda_available else "No")

    if cuda_available:
        table.add_row("GPU Name", torch.cuda.get_device_name(0))
        props = torch.cuda.get_device_properties(0)
        vram_gb = props.total_memory / (1024**3)
        table.add_row("GPU VRAM", f"{vram_gb:.2f} GB")
        table.add_row("GPU Compute", f"{props.major}.{props.minor}")
        table.add_row("GPU Count", str(torch.cuda.device_count()))

        allocated = torch.cuda.memory_allocated() / (1024**3)
        table.add_row("GPU Used", f"{allocated:.2f} GB")

    table.add_row("CPU Count", str(os.cpu_count()))

    if HAS_PSUTIL:
        memory = psutil.virtual_memory()
        table.add_row("RAM Used", f"{memory.used / (1024**3):.2f} GB")
        table.add_row("RAM Total", f"{memory.total / (1024**3):.2f} GB")
        table.add_row("RAM Percent", f"{memory.percent}%")

        cpu_freq = psutil.cpu_freq()
        if cpu_freq:
            table.add_row("CPU Frequency", f"{cpu_freq.current:.0f} MHz")
        table.add_row("CPU Usage", f"{psutil.cpu_percent(interval=0.5)}%")

    model_manager = ModelManager()
    models = model_manager.list_models()
    table.add_row("Models Saved", str(len(models)))
    table.add_row("Active Model", model_manager.get_current_model() or "None")
    table.add_row("Multi-GPU", "Yes" if model_manager.is_multi_gpu() else "No")

    mem_limit = model_manager.config.get('memory_limit_gb', 0)
    table.add_row("Memory Protection", f"{mem_limit} GB" if mem_limit > 0 else "DISABLED")

    table.add_row("PEFT Available", "Yes" if HAS_PEFT else "No")
    table.add_row("8-bit Available", "Yes" if HAS_BNB else "No")
    table.add_row("Evaluate Available", "Yes" if HAS_EVALUATE else "No")
    table.add_row("SacreBLEU Available", "Yes" if HAS_SACREBLEU else "No")
    table.add_row("Matplotlib Available", "Yes" if HAS_MATPLOTLIB else "No")
    table.add_row("Pandas Available", "Yes" if HAS_PANDAS else "No")

    console.print(table)

    input("\n[yellow]Press Enter...[/yellow]")


def enhanced_config_menu() -> None:
    console.print(Panel(Theme.header(" KONFIGURASI"), title="CONFIG", style="yellow"))

    console.print(Theme.info("Pilihan:"))
    console.print("  [green]1. Training Config[/green]")
    console.print("  [green]2. Generation Config[/green]")
    console.print("  [green]3. Memory Protection Settings[/green]")
    console.print("  [green]4. Kembali[/green]")

    cfg_choice = Prompt.ask("Pilih", choices=['1', '2', '3', '4'])

    if cfg_choice == '1':
        _config_training()
    elif cfg_choice == '2':
        _config_generation()
    elif cfg_choice == '3':
        _config_memory_protection()
    else:
        return


def _config_training() -> None:
    model_manager = ModelManager()
    train_cfg = model_manager.get_training_config()

    console.print(Theme.header("Training Configuration:"))

    table = Table(title="Current Settings", box=box.ROUNDED)
    table.add_column("Parameter", style="cyan")
    table.add_column("Value", style="green")

    for k, v in train_cfg.items():
        if k not in ['peft_config', 'data_augmentation']:
            if isinstance(v, float):
                v_str = f"{v:.6f}" if v < 0.001 else f"{v:.4f}"
            else:
                v_str = str(v)
            table.add_row(k, v_str)

    console.print(table)

    if Confirm.ask("[yellow]Ubah konfigurasi?", default=False):
        for k in train_cfg.keys():
            if k in ['peft_config', 'data_augmentation']:
                continue
            new_val = Prompt.ask(f"{k}", default=str(train_cfg[k]))
            try:
                if k in ['learning_rate', 'validation_split', 'warmup_ratio', 'weight_decay', 'gradient_clip_value']:
                    train_cfg[k] = float(new_val)
                elif k in ['batch_size', 'num_epochs', 'max_length', 'warmup_steps',
                          'save_steps', 'logging_steps', 'early_stopping_patience',
                          'seed', 'gradient_accumulation_steps', 'max_samples_limit', 'num_proc']:
                    train_cfg[k] = int(new_val)
                elif k in ['fp16', 'bf16', 'use_peft', 'freeze_embeddings',
                          'dynamic_grad_accumulation', 'use_8bit', 'load_best_model_at_end']:
                    train_cfg[k] = new_val.lower() in ['true', '1', 'yes']
                else:
                    train_cfg[k] = new_val
            except ValueError:
                console.print(Theme.warning(f"Skipped {k}"))

        model_manager.update_training_config(**train_cfg)
        console.print(Theme.success(" Konfigurasi diperbarui"))


def _config_generation() -> None:
    model_manager = ModelManager()
    gen_cfg = model_manager.get_generation_config()

    console.print(Theme.header("Generation Configuration:"))

    table = Table(title="Current Settings", box=box.ROUNDED)
    table.add_column("Parameter", style="cyan")
    table.add_column("Value", style="green")

    for k, v in gen_cfg.items():
        if isinstance(v, float):
            v_str = f"{v:.4f}"
        else:
            v_str = str(v)
        table.add_row(k, v_str)

    console.print(table)

    if Confirm.ask("[yellow]Ubah konfigurasi?", default=False):
        for k in gen_cfg.keys():
            new_val = Prompt.ask(f"{k}", default=str(gen_cfg[k]))
            try:
                if k in ['temperature', 'top_p', 'repetition_penalty']:
                    gen_cfg[k] = float(new_val)
                elif k in ['max_new_tokens', 'top_k', 'no_repeat_ngram_size', 'max_context_length']:
                    gen_cfg[k] = int(new_val)
                else:
                    gen_cfg[k] = new_val
            except ValueError:
                console.print(Theme.warning(f"Skipped {k}"))

        model_manager.update_generation_config(**gen_cfg)
        console.print(Theme.success(" Konfigurasi diperbarui"))


def _config_memory_protection() -> None:
    model_manager = ModelManager()

    console.print("\n" + Theme.header("Memory Protection Settings:"))
    console.print(Theme.warning(" Memory protection helps prevent OOM killer"))

    current_limit = model_manager.config.get('memory_limit_gb', 0)
    console.print(f"  Current memory limit: {current_limit if current_limit > 0 else 'DISABLED'} GB")

    if Confirm.ask("[yellow]Enable memory protection?", default=current_limit > 0):
        new_limit = FloatPrompt.ask("Memory limit (GB)", default=current_limit or DEFAULT_MEMORY_LIMIT_GB)
        model_manager.config['memory_limit_gb'] = new_limit
        console.print(Theme.success(f" Memory limit set to {new_limit} GB"))

        if platform.system().lower() == "linux":
            if set_memory_hard_limit(new_limit):
                console.print(Theme.success(" Protection active"))
            else:
                console.print(Theme.warning(" Protection may not be active"))
    else:
        model_manager.config['memory_limit_gb'] = 0
        console.print(Theme.warning(" Memory protection disabled"))

    model_manager.save_config()


def select_dataset_file(prompt: str = "Pilih file dataset") -> Optional[str]:
    from pathlib import Path
    from rich.prompt import Prompt
    from config import DATA_DIR

    supported_ext = ('.txt', '.json', '.jsonl', '.csv', '.tsv', '.parquet', '.arrow', '.json.gz', '.jsonl.gz')
    data_dir = Path(DATA_DIR)

    files = []
    for ext in supported_ext:
        files.extend(data_dir.glob(f"*{ext}"))

    files = sorted(files)

    if not files:
        console.print(Theme.warning("Tidak ada file dataset di folder 'data/'. Silakan masukkan path manual."))
        path = Prompt.ask("[cyan]Masukkan path file dataset")
        if os.path.exists(path):
            return path
        console.print(Theme.error("File tidak ditemukan!"))
        return None

    console.print(Theme.info("File dataset tersedia:"))

    for idx, f in enumerate(files, 1):
        try:
            size = f.stat().st_size
            if size > 100 * 1024 * 1024:
                size_str = f"{size / (1024*1024*1024):.2f} GB"
            elif size > 1024 * 1024:
                size_str = f"{size / (1024*1024):.2f} MB"
            elif size > 1024:
                size_str = f"{size / 1024:.2f} KB"
            else:
                size_str = f"{size} B"
            console.print(f"  {idx}. {f.name} ({size_str})")
        except Exception:
            console.print(f"  {idx}. {f.name}")

    console.print("  0. Masukkan path manual")

    choice = Prompt.ask("[yellow]Pilih nomor", default="1")

    if choice == "0":
        path = Prompt.ask("[cyan]Masukkan path file dataset")
        if os.path.exists(path):
            return path
        console.print(Theme.error("File tidak ditemukan!"))
        return None

    try:
        idx = int(choice) - 1
        if 0 <= idx < len(files):
            return str(files[idx])
        console.print(Theme.error("Nomor tidak valid!"))
        return None
    except ValueError:
        console.print(Theme.error("Input harus berupa angka!"))
        return None


# ============================================================
# HELPER: load model untuk RAG pipeline
# ============================================================
def _load_rag_model(model_path: str):
    """Load tokenizer + model untuk RAG pipeline. Return (tokenizer, model)."""
    from transformers import AutoTokenizer, AutoModelForCausalLM

    device = get_device()
    dtype = torch.float16 if device == "cuda" else torch.float32

    is_peft = False
    base_model_name = None

    if os.path.exists(os.path.join(model_path, "adapter_config.json")):
        is_peft = True

    base_info_path = os.path.join(model_path, "base_model_info.json")
    if os.path.exists(base_info_path):
        try:
            with open(base_info_path) as f:
                info = json.load(f)
            if info.get("use_peft", False):
                is_peft = True
                base_model_name = info.get("base_model")
        except Exception:
            pass

    tokenizer = AutoTokenizer.from_pretrained(model_path)
    if tokenizer.pad_token is None:
        tokenizer.pad_token = tokenizer.eos_token

    if is_peft and base_model_name and HAS_PEFT:
        try:
            from peft import PeftModel
            base = AutoModelForCausalLM.from_pretrained(
                base_model_name, torch_dtype=dtype, low_cpu_mem_usage=True
            )
            model = PeftModel.from_pretrained(base, model_path)
        except Exception as e:
            console.print(Theme.warning(f"PEFT load failed: {e}, fallback full model"))
            model = AutoModelForCausalLM.from_pretrained(
                model_path, torch_dtype=dtype, low_cpu_mem_usage=True
            )
    else:
        model = AutoModelForCausalLM.from_pretrained(
            model_path, torch_dtype=dtype, low_cpu_mem_usage=True
        )

    model = model.to(device)
    model.eval()
    return tokenizer, model


# ============================================================
# SUBMENU: QUIZ MODULE
# ============================================================
def quiz_module_menu() -> None:
    try:
        from quiz_module import QuizManager
    except Exception as e:
        console.print(Theme.error(f"Quiz module tidak tersedia: {rich_escape(str(e))}"))
        input("\n[yellow]Press Enter...[/yellow]")
        return

    try:
        quiz = QuizManager()
    except Exception as e:
        console.print(Theme.error(f"Gagal inisialisasi QuizManager: {rich_escape(str(e))}"))
        input("\n[yellow]Press Enter...[/yellow]")
        return

    # --- Auto-seed pertanyaan dari dataset kalau masih kosong ---
    if not quiz.questions:
        console.print(Theme.warning("Belum ada pertanyaan tersimpan."))
        if Confirm.ask("[yellow]Generate pertanyaan dari dataset sekarang?", default=True):
            filepath = select_dataset_file()
            if filepath:
                console.print(Theme.info("Loading dataset..."))
                try:
                    loader = EnhancedDatasetLoader()
                    samples, stats = loader.load(filepath)
                    qa_pairs = stats.conversation_pairs or []

                    if not qa_pairs:
                        console.print(Theme.warning("Tidak ada pasangan Q&A di dataset"))
                        console.print(Theme.dim("Quiz hanya mendukung format user/assistant, prompt/response, instruction/response"))
                    else:
                        max_q = min(len(qa_pairs), 200)
                        selected_pairs = qa_pairs[:max_q]
                        new_questions = quiz.generate_questions_from_qa_pairs(selected_pairs)
                        quiz.add_questions(new_questions)
                        console.print(Theme.success(f" ✓ {len(new_questions)} pertanyaan dibuat dari {max_q} Q&A pairs"))
                except Exception as e:
                    console.print(Theme.error(f"Gagal generate: {rich_escape(str(e))}"))
                    input("\n[yellow]Press Enter...[/yellow]")
                    return

    console.print(Panel(Theme.header(" QUIZ MODULE"), title="QUIZ", style="bold yellow"))
    console.print(f"  [dim]Total pertanyaan tersedia: {len(quiz.questions)}[/dim]")
    console.print("  [green]1. Start Quiz[/green]")
    console.print("  [green]2. View Results[/green]")
    console.print("  [green]3. Generate lebih banyak pertanyaan dari dataset[/green]")
    console.print("  [green]4. Kembali[/green]")

    sub = Prompt.ask("[yellow]Pilih", choices=['1', '2', '3', '4'])

    try:
        if sub == '1':
            if not quiz.questions:
                console.print(Theme.error("Tidak ada pertanyaan. Pilih opsi 3 dulu."))
            else:
                num_str = Prompt.ask("[cyan]Jumlah soal", default="5")
                try:
                    num = int(num_str)
                except ValueError:
                    num = 5
                num = max(1, min(num, len(quiz.questions)))
                quiz_id = quiz.create_auto_quiz(num)
                if quiz_id:
                    quiz.start_quiz(quiz_id)

        elif sub == '2':
            quiz.display_results()

        elif sub == '3':
            filepath = select_dataset_file()
            if filepath:
                try:
                    loader = EnhancedDatasetLoader()
                    samples, stats = loader.load(filepath)
                    qa_pairs = stats.conversation_pairs or []
                    if not qa_pairs:
                        console.print(Theme.error("Dataset tidak punya pasangan Q&A"))
                    else:
                        new_questions = quiz.generate_questions_from_qa_pairs(qa_pairs)
                        quiz.add_questions(new_questions)
                        console.print(Theme.success(f" ✓ {len(new_questions)} pertanyaan ditambahkan"))
                except Exception as e:
                    console.print(Theme.error(f"Gagal generate: {rich_escape(str(e))}"))

    except Exception as e:
        console.print(Theme.error(f"Error: {rich_escape(str(e))}"))

    input("\n[yellow]Press Enter...[/yellow]")


# ============================================================
# SUBMENU: KNOWLEDGE BASE (SEMANTIC SEARCH)
# ============================================================
def knowledge_base_menu() -> None:
    try:
        from semantic_search import KnowledgeBase
    except Exception as e:
        console.print(Theme.error(f"Knowledge Base tidak tersedia: {rich_escape(str(e))}"))
        input("\n[yellow]Press Enter...[/yellow]")
        return

    try:
        kb = KnowledgeBase()
    except Exception as e:
        console.print(Theme.error(f"Gagal inisialisasi KnowledgeBase: {rich_escape(str(e))}"))
        input("\n[yellow]Press Enter...[/yellow]")
        return

    console.print(Panel(Theme.header(" KNOWLEDGE BASE"), title="KB", style="bold yellow"))
    console.print(f"  [dim]Dokumen tersimpan: {len(kb.documents)}[/dim]")
    console.print("  [green]1. Search[/green]")
    console.print("  [green]2. Add Document[/green]")
    console.print("  [green]3. Import dari Dataset[/green]")
    console.print("  [green]4. Kembali[/green]")

    sub = Prompt.ask("[yellow]Pilih", choices=['1', '2', '3', '4'])

    try:
        if sub == '1':
            query = Prompt.ask("[cyan]Search query")
            if not query.strip():
                return
            top_k_str = Prompt.ask("[cyan]Top K", default="5")
            try:
                top_k = int(top_k_str)
            except ValueError:
                top_k = 5

            results = kb.search(query, top_k=top_k)
            if not results:
                console.print(Theme.warning("Tidak ada hasil"))
            else:
                for i, r in enumerate(results, 1):
                    score = r.get('score', 0.0)
                    text = r.get('text', '')
                    console.print(f"[green]#{i} Score: {score:.3f}[/green]")
                    console.print(rich_escape(text[:300]))
                    console.print()

        elif sub == '2':
            text = Prompt.ask("[cyan]Masukkan teks dokumen")
            if text.strip():
                if hasattr(kb, 'add_document'):
                    try:
                        import time as _t
                        doc_id = f"doc_{int(_t.time())}"
                        kb.add_document(doc_id, text)
                        console.print(Theme.success(f" Dokumen ditambahkan ({doc_id})"))
                    except TypeError:
                        kb.add_document(text)
                        console.print(Theme.success(" Dokumen ditambahkan"))

        elif sub == '3':
            filepath = select_dataset_file()
            if filepath:
                try:
                    loader = EnhancedDatasetLoader()
                    samples, stats = loader.load(filepath)
                    added = 0
                    import time as _t
                    for i, text in enumerate(samples[:500]):
                        if text.strip():
                            try:
                                kb.add_document(f"ds_{i}_{int(_t.time() * 1000)}", text)
                                added += 1
                            except Exception:
                                pass
                    console.print(Theme.success(f" ✓ {added} dokumen di-import"))
                except Exception as e:
                    console.print(Theme.error(f"Gagal import: {rich_escape(str(e))}"))

    except Exception as e:
        console.print(Theme.error(f"Error: {rich_escape(str(e))}"))

    input("\n[yellow]Press Enter...[/yellow]")


# ============================================================
# SUBMENU: RAG PIPELINE
# ============================================================
def rag_pipeline_menu() -> None:
    try:
        from rag_module import RAGPipeline
    except Exception as e:
        console.print(Theme.error(f"RAG Pipeline tidak tersedia: {rich_escape(str(e))}"))
        input("\n[yellow]Press Enter...[/yellow]")
        return

    try:
        rag = RAGPipeline()
    except Exception as e:
        console.print(Theme.error(f"Gagal inisialisasi RAGPipeline: {rich_escape(str(e))}"))
        input("\n[yellow]Press Enter...[/yellow]")
        return

    console.print(Panel(Theme.header(" RAG PIPELINE"), title="RAG", style="bold yellow"))

    # --- Load model aktif untuk generation ---
    model_manager = ModelManager()
    current_model = model_manager.get_current_model()

    if not current_model:
        console.print(Theme.warning("Tidak ada model aktif. RAG akan bekerja tanpa generation (retrieve only)."))
        console.print(Theme.dim("Set model aktif via menu [03] Manage Model terlebih dahulu."))
    else:
        model_path = os.path.join(MODEL_DIR, current_model)
        if not os.path.exists(model_path):
            console.print(Theme.error(f"Model path tidak ditemukan: {model_path}"))
        else:
            console.print(Theme.info(f"Loading model untuk RAG: {current_model}..."))
            try:
                tokenizer, model = _load_rag_model(model_path)
                rag.set_model(tokenizer, model)
                console.print(Theme.success(" Model RAG siap!"))
            except Exception as e:
                console.print(Theme.error(f"Gagal load model: {rich_escape(str(e))}"))
                console.print(Theme.warning("RAG akan berjalan tanpa generation"))

    try:
        if hasattr(rag, 'query_interactive'):
            rag.query_interactive()
        else:
            query = Prompt.ask("[cyan]Query")
            if query.strip():
                result = rag.query(query)
                console.print(Panel(rich_escape(str(result)), title="Answer", style="green"))
    except Exception as e:
        console.print(Theme.error(f"Error: {rich_escape(str(e))}"))

    input("\n[yellow]Press Enter...[/yellow]")


# ============================================================
# MAIN MENU
# ============================================================
def enhanced_menu() -> None:
    while True:
        os.system('clear' if os.name == 'posix' else 'cls')
        logo()

        model_manager = ModelManager()
        current_model = model_manager.get_current_model()
        model_display = f"[green]{current_model}" if current_model else "[red]None"
        mem_limit = model_manager.config.get('memory_limit_gb', 0)
        mem_status = "[green]ACTIVE" if mem_limit > 0 else "[red]INACTIVE"

        # Status Google Drive
        drive_mounted = Path("/content/drive/MyDrive").exists()
        drive_status = "[green]CONNECTED" if drive_mounted else "[yellow]NOT MOUNTED"

        console.print(Panel(
            f"""[bold white][[bold green]01[/bold white]] [bold white]Training Ultimate - Auto-LoRA • Multi-GPU
[bold white][[bold green]02[/bold white]] [bold white]Chat Enhanced     - Chat dengan context management
[bold white][[bold green]03[/bold white]] [bold white]Manage Model      - List, set aktif, hapus
[bold white][[bold green]04[/bold white]] [bold white]Dataset Inspector - Analisa dataset
[bold white][[bold green]05[/bold white]] [bold white]Visualize History - Plot training metrics
[bold white][[bold green]06[/bold white]] [bold white]Konfigurasi       - Training & Generation settings
[bold white][[bold green]07[/bold white]] [bold white]System Stats      - Informasi sistem
[bold white][[bold green]08[/bold white]] [bold white]Quiz Module       - Latihan soal otomatis
[bold white][[bold green]09[/bold white]] [bold white]Knowledge Base    - Semantic search dokumen
[bold white][[bold green]10[/bold white]] [bold white]RAG Pipeline      - Retrieval-Augmented Generation
[bold white][[bold green]11[/bold white]] [bold white]Backup Project    - Export ke Google Drive
[bold white][[bold green]12[/bold white]] [bold white]Exit""",
            width=70, title="[bold yellow]ULTIMATE MENU v2.5 FIXED", style="bold yellow"
        ))

        console.print(Panel(
            f"[cyan]Status: [white]ACTIVE\n"
            f"[cyan]Model Aktif: {model_display}\n"
            f"[cyan]Memory Protection: {mem_status}\n"
            f"[cyan]Google Drive: {drive_status}\n"
            f"[cyan]Multi-GPU: {'[green]Enabled[/green]' if model_manager.is_multi_gpu() else '[dim]Disabled[/dim]'}",
            width=40, title="[yellow]INFO", style="yellow"
        ))

        pilihan = Prompt.ask("[bold yellow] Pilih (1-12)")

        if pilihan in ['1', '01']:
            EnhancedTrainingModule().run()
        elif pilihan in ['2', '02']:
            EnhancedChatModule().run()
        elif pilihan in ['3', '03']:
            manage_models()
        elif pilihan in ['4', '04']:
            dataset_inspector()
        elif pilihan in ['5', '05']:
            visualize_history_menu()
        elif pilihan in ['6', '06']:
            enhanced_config_menu()
        elif pilihan in ['7', '07']:
            enhanced_system_stats()
        elif pilihan in ['8', '08']:
            quiz_module_menu()
        elif pilihan in ['9', '09']:
            knowledge_base_menu()
        elif pilihan in ['10']:
            rag_pipeline_menu()
        elif pilihan in ['11']:
            backup_menu()
        elif pilihan in ['12', 'exit', 'q']:
            console.print(Theme.success(" Keluar..."))
            sys.exit(0)
        else:
            console.print(Theme.error("Pilihan tidak valid!"))
            input("\n[yellow]Press Enter...[/yellow]")