# File: src/explore.py (SỬA LỖI - CHỈ CHẠY hai2131)
import json
import os
import sys
from datasets import load_dataset
from itertools import islice
import traceback

DATASET_LIST = {
    # "trantac": "TraTacXiMuoi/Ielts_writing_task1_academic", # Tạm thời tắt do lỗi mạng
    "hai2131": "hai2131/IELTS-essays-task-1"
}

NUM_SAMPLES_TO_VIEW = 2
SPLIT_NAME = "train"

def safe_value_to_string(value):
    """Chuyển đổi value thành string an toàn"""
    if value is None:
        return None
    if isinstance(value, (str, int, float, bool)):
        return value
    if isinstance(value, dict):
        return value
    if isinstance(value, list):
        return value
    # Đối với các object khác (ảnh, audio, etc)
    return f"<{type(value).__name__}>"

def explore_dataset(name: str, path: str, split: str, n: int):
    """
    Tải N mẫu đầu tiên của một dataset từ Hugging Face và in cấu trúc của nó.
    """
    print("="*80)
    print(f"🕵️  Đang khám phá dataset: {name}")
    print(f"    Path: {path}")
    print(f"    Split: {split}")
    print("="*80)
    
    try:
        # Tải N mẫu đầu tiên (không dùng streaming nữa,
        # vì dataset hai2131 chỉ 8MB, tải luôn cho nhanh)
        dataset = load_dataset(path, split=f"{split}[:{n}]")
        
        print(f"\n✅ Tải thành công. Cấu trúc (Features):")
        # In ra các cột và kiểu dữ liệu
        print(dataset.features)
        
        print(f"\n--- Đang xem {n} mẫu đầu tiên ---")
        
        for i, item in enumerate(dataset):
            print(f"\n--- Mẫu {i+1} ---")
            
            printable_item = {}
            for key, value in item.items():
                printable_item[key] = safe_value_to_string(value)

            print(json.dumps(printable_item, ensure_ascii=False, indent=2))
            
    except Exception as e:
        print(f"\n❌ LỖI khi tải hoặc đọc dataset '{name}':")
        print(f"   {e}")
        traceback.print_exc()

def list_available_splits(path):
    # Hàm này không cần thiết nữa nếu chúng ta tải trực tiếp
    pass

def main():
    print("🚀 BẮT ĐẦU KHÁM PHÁ IELTS DATASETS (CHỈ hai2131)")
    print("="*80)
    
    for name, path in DATASET_LIST.items():
        try:
            explore_dataset(name, path, SPLIT_NAME, NUM_SAMPLES_TO_VIEW)
        except KeyboardInterrupt:
            print("\n⚠️  Bị gián đoạn bởi người dùng")
            break
        except Exception as e:
            print(f"\n❌ Lỗi không mong muốn: {e}")
            traceback.print_exc()
        
        print("\n" + "-"*80)
    
    print("\n" + "="*80)
    print("✅ KHÁM PHÁ HOÀN TẤT")
    print("="*80)

if __name__ == "__main__":
    main()