import os
import json
from datasets import load_dataset
from tqdm import tqdm


def process_chartqa_dataset(base_dataset_dir, output_image_dir, output_jsonl_path):
    """
    Downloads ChartQA, saves images, and writes JSONL with paths relative
    to base_dataset_dir.
    """

    # 1. Create the output directory if it doesn't exist
    if not os.path.exists(output_image_dir):
        os.makedirs(output_image_dir)
        print(f"Created directory: {output_image_dir}")

    # 2. Load the dataset
    print("Loading HuggingFaceM4/ChartQA dataset...")
    try:
        dataset = load_dataset("HuggingFaceM4/ChartQA", split="train")
    except Exception as e:
        print(f"Error loading dataset: {e}")
        return

    print(f"Processing {len(dataset)} items...")

    # 3. Iterate and Process
    with open(output_jsonl_path, 'w', encoding='utf-8') as f:
        for idx, item in enumerate(tqdm(dataset)):
            try:
                # --- A. Save Image ---
                image = item['image']

                # Create filename
                image_filename = f"chartqa_train_{idx}.png"
                full_image_path = os.path.join(output_image_dir, image_filename)

                # Save the image physically to the full path
                image.save(full_image_path)

                # --- B. Calculate Relative Path ---
                # This converts /home/yuheng/datasets/ChartQA/images/img.png
                # to ChartQA/images/img.png
                relative_image_path = os.path.relpath(full_image_path, base_dataset_dir)

                # --- C. Process Data Fields ---
                question_text = item.get('query', item.get('question', ''))
                raw_label = item.get('label', [])
                answers_list = raw_label if isinstance(raw_label, list) else [str(raw_label)]
                primary_answer = answers_list[0] if answers_list else ""

                # --- D. Construct JSON Object ---
                entry = {
                    "question_id": str(idx),
                    "image": relative_image_path,  # <--- UPDATED: Uses relative path now
                    "prompt": question_text,
                    "category": "default",
                    "answers": answers_list,
                    "answer": primary_answer,
                    "dataset": "ChartQA"
                }

                f.write(json.dumps(entry) + "\n")

            except Exception as e:
                print(f"Skipping index {idx} due to error: {e}")

    print(f"\nProcessing complete.")
    print(f"Images saved to: {output_image_dir}")
    print(f"Annotations saved to: {output_jsonl_path}")


# --- Configuration ---
if __name__ == "__main__":
    # The root directory you want the paths to be relative to
    BASE_ROOT = "/home/yuheng/datasets"

    # Where the images and jsonl will actually be stored
    # Ensure this path is inside BASE_ROOT if you want a clean relative path
    IMAGE_SAVE_PATH = os.path.join(BASE_ROOT, "ChartQA/images")
    JSONL_SAVE_PATH = os.path.join(BASE_ROOT, "chartqa_28k_qa.jsonl")

    process_chartqa_dataset(BASE_ROOT, IMAGE_SAVE_PATH, JSONL_SAVE_PATH)