import re
import hashlib
from functools import lru_cache
from config import CACHE_SIZE
from datasets import load_dataset # Chúng ta chỉ cần load_dataset, không cần load_from_disk cho JSONL
import os
from pyvi import ViTokenizer

def load_iwslt_dataset(data_folder="data"): # Đổi tên tham số 'path' thành 'data_folder' để rõ nghĩa hơn
    """Load IWSLT15 dataset from local JSONL files in the specified folder."""
    
    # Lấy đường dẫn đến thư mục chứa script hiện tại (utils.py)
    script_dir = os.path.dirname(os.path.abspath(__file__))
    
    # Kết hợp đường dẫn script với tên thư mục data để có đường dẫn tuyệt đối
    full_data_path = os.path.join(script_dir, data_folder)

    train_file = os.path.join(full_data_path, "train.jsonl")
    validation_file = os.path.join(full_data_path, "validation.jsonl")
    test_file = os.path.join(full_data_path, "test.jsonl")

    # Kiểm tra xem các tệp có tồn tại không
    if not (os.path.exists(train_file) and os.path.exists(validation_file) and os.path.exists(test_file)):
        print(f"Error: Missing one or more data files in '{full_data_path}'.")
        print(f"Expected: {train_file}, {validation_file}, {test_file}")
        print("Please ensure you have run the download script (như hướng dẫn trước) và các file được đặt đúng chỗ.")
        return None, None, None

    # Khai báo các tệp dữ liệu cho Hugging Face datasets loader
    data_files = {
        "train": train_file,
        "validation": validation_file,
        "test": test_file,
    }

    try:
        print(f"Đang tải dataset từ các tệp cục bộ tại: {full_data_path}...")
        # Sử dụng load_dataset với builder "json" để tải các tệp .jsonl
        # Thư viện sẽ tự động phân tích cấu trúc {"translation": {"en": ..., "vi": ...}}
        dataset = load_dataset("json", data_files=data_files)
        print("Dataset đã được tải thành công từ các tệp cục bộ.")

        train_data = dataset["train"]
        val_data = dataset["validation"]
        test_data = dataset["test"]
        
        print(f"Loaded IWSLT15 dataset: {len(train_data)} train, {len(val_data)} val, {len(test_data)} test samples")
        return train_data, val_data, test_data
    except Exception as e:
        print(f"Error loading IWSLT15 dataset from local files: {e}")
        return None, None, None

@lru_cache(maxsize=CACHE_SIZE)
def _preprocess_text_internal(text):
    """Cached version of text preprocessing (only whitespace normalization and lowercasing)"""
    if not text:
        return ""
    text = text.strip()
    text = re.sub(r"\s+", " ", text) # Chuẩn hóa nhiều khoảng trắng thành một
    text = text.lower() # Chuyển về chữ thường
    # Đã loại bỏ: text = re.sub(r'[^\w\s]', '', text) để giữ lại dấu câu
    return text

def preprocess_text(text):
    """Preprocess text with caching (main entry for general text)"""
    return _preprocess_text_internal(text)

def preprocess_pair(en_text, vi_text): # Đã bỏ tham số max_length, sẽ xử lý ở tokenizer
    """Preprocess a pair of English-Vietnamese sentences"""
    if not en_text or not vi_text:
        return "", ""
    
    # Áp dụng ViTokenizer trước khi tiền xử lý chung cho tiếng Việt
    # PyVi thường hoạt động tốt nhất trên văn bản gốc hoặc đã chuẩn hóa cơ bản
    vi_text_tokenized = ViTokenizer.tokenize(vi_text)
    
    # Sau đó áp dụng tiền xử lý chung (chuyển chữ thường, chuẩn hóa khoảng trắng)
    en_text_final = _preprocess_text_internal(en_text)
    vi_text_final = _preprocess_text_internal(vi_text_tokenized)
    
    return en_text_final, vi_text_final

def create_cache_key(text, model_choice, max_length, num_beams):
    """Create hash key for cache"""
    content = f"{text}_{model_choice}_{max_length}_{num_beams}"
    return hashlib.md5(content.encode()).hexdigest()

def clear_preprocessing_cache():
    """Clear the preprocessing cache"""
    _preprocess_text_internal.cache_clear() # Đổi tên hàm