#!/usr/bin/env python3
import os
import sys
import random
import json
from pathlib import Path

import torch
import numpy as np
from transformers import set_seed

from config import (
    DIRECTORIES, DATA_DIR, CONFIG_FILE,
    DEFAULT_MEMORY_LIMIT_GB, DEFAULT_MODEL
)
from utils import console, Theme, graceful_exit, error_logger
from models import ModelManager
from ui import enhanced_menu

def create_sample_datasets() -> None:
    sample_file = os.path.join(DATA_DIR, "dataset.txt")
    if not os.path.exists(sample_file):
        with open(sample_file, 'w', encoding='utf-8') as f:
            f.write("User: Halo, apa kabar?\nAI: Halo! Saya baik, terima kasih. Ada yang bisa saya bantu?\n\n")
            f.write("User: Siapa kamu?\nAI: Saya adalah chatbot AI yang sedang dikembangkan.\n\n")
            f.write("User: Bisakah kamu membantu saya?\nAI: Tentu! Saya siap membantu Anda.\n\n")
            f.write("User: Terima kasih\nAI: Sama-sama! Senang bisa membantu.\n")

    sample_jsonl = os.path.join(DATA_DIR, "dataset.jsonl")
    if not os.path.exists(sample_jsonl):
        with open(sample_jsonl, 'w', encoding='utf-8') as f:
            f.write('{"user": "Halo, apa kabar?", "assistant": "Halo! Saya baik, terima kasih."}\n')
            f.write('{"user": "Siapa kamu?", "assistant": "Saya adalah chatbot AI."}\n')
            f.write('{"user": "Bisakah kamu membantu?", "assistant": "Tentu! Saya siap membantu."}\n')

def initialize_config() -> None:
    model_manager = ModelManager()
    if 'memory_limit_gb' not in model_manager.config:
        model_manager.config['memory_limit_gb'] = DEFAULT_MEMORY_LIMIT_GB
        model_manager.save_config()

if __name__ == '__main__':
    for dir_path in DIRECTORIES.values():
        Path(dir_path).mkdir(exist_ok=True)

    create_sample_datasets()
    initialize_config()

    set_seed(42)
    np.random.seed(42)
    random.seed(42)

    console.print(Theme.success(" Ultimate AI Chatbot v2.5 FIXED initialized"))
    console.print(Theme.dim("Features: Auto-LoRA • Multi-GPU • 8-bit Quantization • Parallel Tokenization"))
    console.print(Theme.gold(" Rating: 9.8/10 - Production Grade"))

    with graceful_exit:
        while not graceful_exit.should_exit:
            try:
                enhanced_menu()
                if not graceful_exit.should_exit:
                    input("\n[yellow]Press Enter untuk melanjutkan...[/yellow]")
            except KeyboardInterrupt:
                console.print("\n" + Theme.error(" Aplikasi dihentikan"))
                break
            except Exception as e:
                console.print(Theme.error(f"Error: {rich_escape(str(e))}"))
                error_logger.error(f"Main loop error: {e}")
                input("\n[yellow]Press Enter...[/yellow]")

    console.print(Theme.success(" Sampai jumpa!"))