# import os
# import json
# import datasets

import os
import json
import datasets

# Hugging Face Dataset Info
DS_NAME = "anna-tch/generation-results"
HF_TOKEN = os.getenv("HF_TOKEN")
PROGRESS_FILE = "progress.json"  # Save progress here

# Load dataset
def load_dataset():
    """Loads the dataset from Hugging Face."""
    try:
        dataset = datasets.load_dataset(DS_NAME, token=HF_TOKEN)["train"]
        return dataset
    except Exception as e:
        print(f"Error loading dataset: {e}")
        return None

# Get generation columns
def get_generation_columns(dataset):
    """Returns a list of generation columns, excluding metadata columns."""
    return [col for col in dataset.column_names if col not in ["comment_id", "manual_annotation"]]

# Load progress from the JSON file
def load_progress():
    """Load the saved annotations progress from a JSON file."""
    if os.path.exists(PROGRESS_FILE):
        with open(PROGRESS_FILE, "r") as f:
            return json.load(f)
    return {}

# Save progress to the JSON file
def save_progress(progress):
    """Save annotations progress to a JSON file."""
    with open(PROGRESS_FILE, "w") as f:
        json.dump(progress, f, indent=4)

# Fetch the next unannotated sample
def get_next_sample(dataset, progress):
    """Finds the next unannotated comment, or None if no more unannotated samples."""
    for sample in dataset:
        if sample["comment_id"] not in progress:
            return sample
    return None  # No more unannotated samples

# Annotate text and update dataset
def annotate_text(dataset, comment_id, grammar_scores, coherence_scores, preferred_text, generation_columns, progress):
    """Annotate the sample and update it in the dataset."""
    
    # Store annotation in the progress
    progress[comment_id] = {
        "grammar": dict(zip(generation_columns, grammar_scores)),
        "coherence": dict(zip(generation_columns, coherence_scores)),
        "preferred_text": preferred_text
    }

    # Save progress after annotating
    save_progress(progress)

    # Find the sample with the given comment_id and update it
    df = dataset.to_pandas()
    df.loc[df["comment_id"] == comment_id, "manual_annotation"] = progress[comment_id]

    # Push the updated dataset to Hugging Face
    dataset.push_to_hub(DS_NAME)

    return dataset


# # Hugging Face Dataset Info
# DS_NAME = "anna-tch/generation-results"
# HF_TOKEN = os.getenv("HF_TOKEN")
# PROGRESS_FILE = "progress.json"

# # Load dataset
# def load_dataset():
#     """Loads the dataset from Hugging Face."""
#     try:
#         dataset = datasets.load_dataset(DS_NAME, token=HF_TOKEN)["train"]
#         return dataset
#     except Exception as e:
#         print(f"Error loading dataset: {e}")
#         return None

# def load_progress():
#     """Loads the progress file."""
#     try:
#         with open(PROGRESS_FILE, "r") as f:
#             return json.load(f)
#     except FileNotFoundError:
#         # Create progress file if it doesn’t exist
#         with open(PROGRESS_FILE, "w") as f:
#             json.dump({}, f)
#             return {}

# # Get generation columns
# def get_generation_columns(dataset):
#     """Returns a list of generation columns, excluding metadata columns."""
#     return [col for col in dataset.column_names if col not in ["comment_id", "manual_annotation"]]

# # Fetch the next unannotated sample
# def get_next_sample(dataset):
#     """Finds the next unannotated comment."""
#     for sample in dataset:
#         if sample["manual_annotation"] is None:
#             return sample
#     return None  # No more samples

# # Update dataset using comment_id
# def annotate_text(dataset, comment_id, grammar_scores, coherence_scores, preferred_text, generation_columns):
#     """Finds the correct sample by comment_id and updates it."""
    
#     def update_sample(example):
#         """Updates only the sample with the given comment_id."""
#         if example["comment_id"] == comment_id:
#             return {
#                 **example,
#                 "manual_annotation": {
#                     "grammar": dict(zip(generation_columns, grammar_scores)),
#                     "coherence": dict(zip(generation_columns, coherence_scores)),
#                     "preferred_text": preferred_text
#                 }
#             }
#         return example  # Return unchanged sample if it doesn’t match

#     # Apply update to dataset
#     dataset = dataset.map(update_sample)

#     # Push updated dataset to Hugging Face
#     dataset.push_to_hub(DS_NAME)

#     return dataset