import json
import gzip
import random
import requests
from transformers import DistilBertTokenizerFast

genre_url_dict = {
    'poetry': 'https://mcauleylab.ucsd.edu/public_datasets/gdrive/goodreads/byGenre/goodreads_reviews_poetry.json.gz',
    'children': 'https://mcauleylab.ucsd.edu/public_datasets/gdrive/goodreads/byGenre/goodreads_reviews_children.json.gz',
    'comics_graphic': 'https://mcauleylab.ucsd.edu/public_datasets/gdrive/goodreads/byGenre/goodreads_reviews_comics_graphic.json.gz',
    'fantasy_paranormal': 'https://mcauleylab.ucsd.edu/public_datasets/gdrive/goodreads/byGenre/goodreads_reviews_fantasy_paranormal.json.gz',
    'history_biography': 'https://mcauleylab.ucsd.edu/public_datasets/gdrive/goodreads/byGenre/goodreads_reviews_history_biography.json.gz',
    'mystery_thriller_crime': 'https://mcauleylab.ucsd.edu/public_datasets/gdrive/goodreads/byGenre/goodreads_reviews_mystery_thriller_crime.json.gz',
    'romance': 'https://mcauleylab.ucsd.edu/public_datasets/gdrive/goodreads/byGenre/goodreads_reviews_romance.json.gz',
    'young_adult': 'https://mcauleylab.ucsd.edu/public_datasets/gdrive/goodreads/byGenre/goodreads_reviews_young_adult.json.gz'
}

def load_reviews(url, head=10000, sample_size=2000):
    reviews = []
    response = requests.get(url, stream=True)
    with gzip.open(response.raw, 'rt', encoding='utf-8') as file:
        for i, line in enumerate(file):
            if i >= head: break
            reviews.append(json.loads(line)['review_text'])
    return random.sample(reviews, min(sample_size, len(reviews)))

def prepare_data(model_name='distilbert-base-cased'):
    tokenizer = DistilBertTokenizerFast.from_pretrained(model_name)
    train_texts, train_labels, test_texts, test_labels = [], [], [], []
    
    for genre, url in genre_url_dict.items():
        reviews = load_reviews(url, head=1000, sample_size=200)
        split = int(len(reviews) * 0.8)
        train_texts.extend(reviews[:split])
        train_labels.extend([genre] * split)
        test_texts.extend(reviews[split:])
        test_labels.extend([genre] * (len(reviews) - split))

    unique_labels = sorted(list(set(train_labels)))
    label2id = {label: i for i, label in enumerate(unique_labels)}
    
    train_encodings = tokenizer(train_texts, truncation=True, padding=True, max_length=512)
    test_encodings = tokenizer(test_texts, truncation=True, padding=True, max_length=512)
    
    return train_encodings, [label2id[l] for l in train_labels], test_encodings, [label2id[l] for l in test_labels], label2id