import os
import torch
import numpy as np
from torch.utils.data import Dataset
import glob
import json

class TextDataset(Dataset):
    """
    Pre-tokenize edilmiş binary dosyadan okur (wikipedia_tokens.bin).
    __getitem__ sadece numpy slice - çok hızlı.
    """
    def __init__(self, data_dir, tokenizer, max_length=512):
        self.tokenizer = tokenizer
        self.max_length = max_length

        bin_path = os.path.join(data_dir, "wikipedia_tokens.bin")
        if not os.path.exists(bin_path):
            raise FileNotFoundError(
                f"{bin_path} bulunamadı!\n"
                "Önce çalıştır: cd ~/pege && python3 pretokenize.py"
            )

        print(f"Token verisi yükleniyor: {bin_path}")
        self.tokens = np.memmap(bin_path, dtype=np.uint16, mode='r')
        print(f"Toplam token: {len(self.tokens):,}")

    def __len__(self):
        return max(0, len(self.tokens) - self.max_length - 1)

    def __getitem__(self, idx):
        chunk = self.tokens[idx : idx + self.max_length + 1].astype(np.int64)
        x = torch.from_numpy(chunk[:self.max_length])
        y = torch.from_numpy(chunk[1:self.max_length + 1])
        return x, y


class ConversationDataset(Dataset):
    """RLHF için konuşma verisi"""
    def __init__(self, feedback_file, tokenizer):
        self.tokenizer = tokenizer
        self.conversations = []

        if os.path.exists(feedback_file):
            with open(feedback_file, 'r', encoding='utf-8') as f:
                for line in f:
                    self.conversations.append(json.loads(line))

    def __len__(self):
        return len(self.conversations)

    def __getitem__(self, idx):
        conv = self.conversations[idx]
        text = f"{self.tokenizer.USER}{conv['input']}{self.tokenizer.ASSISTANT}{conv['output']}"
        tokens = self.tokenizer.encode(text)
        reward = 1.0 if conv['feedback'] == 'good' else -1.0
        return tokens, reward