import torch
import os
import glob
import json
from tokenizer import Byte_Tokenizer

def load_all_text_data(data_dir):
    text_data = ""
    # Load .txt files
    for filepath in glob.glob(os.path.join(data_dir, "*.txt")):
        with open(filepath, 'r', encoding='utf-8', errors='ignore') as f:
            text_data += f.read() + "\n"
    
    # Load .jsonl files
    for filepath in glob.glob(os.path.join(data_dir, "*.jsonl")):
        with open(filepath, 'r', encoding='utf-8', errors='ignore') as f:
            for line in f:
                if not line.strip():
                    continue
                try:
                    obj = json.loads(line.strip())
                    if "turns" in obj:
                        for turn in obj["turns"]:
                            text_data += f"{turn.get('speaker', '')}: {turn.get('text', '')}\n"
                    elif "text" in obj:
                        text_data += obj["text"] + "\n"
                except:
                    pass
    return text_data

def get_real_data_loader(config, data_dir="data"):
    print(f"Loading datasets from {data_dir}...")
    text_data = load_all_text_data(data_dir)
    if not text_data:
        raise ValueError(f"No text data found in {data_dir}!")
        
    tokenizer = Byte_Tokenizer()
    print("Tokenizing the entire dataset using Character-level bytes...")
    tokens = tokenizer.encode(text_data)
    
    tensor_data = torch.tensor(tokens, dtype=torch.long)
    total_len = tensor_data.size(0)
    
    chunk_size = config.batch_size * config.seq_len
    num_chunks = total_len // chunk_size
    
    if num_chunks == 0:
        raise ValueError("Dataset is too small for the configured batch_size and seq_len.")
        
    print(f"Dataset loaded: {total_len} tokens, {num_chunks} batches available.")
    
    # Infinite generator
    while True:
        # Create random offset to yield different batches
        for _ in range(num_chunks):
            idx = torch.randint(0, total_len - chunk_size - 1, (1,)).item()
            chunk = tensor_data[idx:idx+chunk_size+1]
            x = chunk[:-1].view(config.batch_size, config.seq_len)
            y = chunk[1:].view(config.batch_size, config.seq_len)
            yield x, y

def save_checkpoint(model, optimizer, step, path="checkpoint.pt"):
    torch.save({
        'step': step,
        'model_state_dict': model.state_dict(),
        'optimizer_state_dict': optimizer.state_dict(),
    }, path)
    print(f"Checkpoint saved to {path}")

def load_checkpoint(model, optimizer, path="checkpoint.pt"):
    if os.path.exists(path):
        checkpoint = torch.load(path)
        model.load_state_dict(checkpoint['model_state_dict'])
        optimizer.load_state_dict(checkpoint['optimizer_state_dict'])
        print(f"Loaded checkpoint from {path} (Step {checkpoint['step']})")
        return checkpoint['step']
    return 0