"""This creates a dataset from the AndroidControl dataset in GCP and saves it to a local directory.
Before running it, you need to install tensorflow and the google-cloud-storage library:

pip install tensorflow google-cloud-storage
"""

import argparse
import base64
import gc
import json
import multiprocessing
import os
import time
from concurrent.futures import ThreadPoolExecutor, as_completed
from io import BytesIO

import psutil
import tensorflow as tf
from datasets import Dataset, DatasetDict, concatenate_datasets
from google.cloud import storage
from PIL import Image
from tqdm import tqdm


def parse_arguments():
    parser = argparse.ArgumentParser(description="Process AndroidControl dataset")
    parser.add_argument("--bucket", type=str, default="gresearch", help="Google Cloud Storage bucket name")
    parser.add_argument("--prefix", type=str, default="android_control", help="Prefix for files in the bucket")
    parser.add_argument(
        "--output_dir", type=str, default="android_control_dataset", help="Directory to save the processed dataset"
    )
    parser.add_argument("--data_dir", type=str, default="android_data", help="Directory to store downloaded files")
    parser.add_argument("--max_files", type=int, default=None, help="Maximum number of files to process (None for all)")
    parser.add_argument("--download", action="store_true", help="Download files from GCS")
    parser.add_argument("--workers", type=int, default=multiprocessing.cpu_count(), help="Number of worker processes")
    parser.add_argument("--chunk_size", type=int, default=10, help="Number of files to process per worker")
    parser.add_argument("--splits_file", type=str, default=None, help="JSON file with custom train/test splits")
    parser.add_argument("--test_size", type=float, default=0.2, help="Size of test split if no splits file provided")
    parser.add_argument("--seed", type=int, default=42, help="Random seed for reproducibility")
    parser.add_argument(
        "--repo_id", type=str, default="smolagents/android-control", help="Hugging Face repository ID for streaming"
    )
    parser.add_argument("--token", type=str, default=None, help="Hugging Face token for streaming")
    parser.add_argument("--batch_size", type=int, default=100, help="Number of examples to process at once")
    parser.add_argument(
        "--max_memory_gb", type=float, default=8.0, help="Maximum memory usage in GB before processing batch"
    )
    parser.add_argument(
        "--skip_train_split",
        action="store_true",
        help="Skip the training split (default: do not skip)",
    )
    return parser.parse_args()


def download_file(blob, destination_path):
    os.makedirs(os.path.dirname(destination_path), exist_ok=True)
    blob.download_to_filename(destination_path)
    return True


def download_files_parallel(bucket_name, prefix, destination_folder, max_workers=16):
    print("Starting download process...")
    print(f"Downloading files from gs://{bucket_name}/{prefix} to {destination_folder}")

    os.makedirs(destination_folder, exist_ok=True)

    # Create an anonymous client for public bucket access
    print("Initializing Google Cloud Storage client...")
    client = storage.Client.create_anonymous_client()
    bucket = client.bucket(bucket_name)

    print("Listing files in bucket...")
    blobs = list(bucket.list_blobs(prefix=prefix))
    print(f"Found {len(blobs)} files to download")

    blobs = [blob for blob in blobs if not blob.name.endswith("/")]
    print(f"Filtered to {len(blobs)} non-directory files")

    download_tasks = []
    for blob in blobs:
        destination_path = os.path.join(destination_folder, blob.name.replace(prefix + "/", ""))
        download_tasks.append((blob, destination_path))

    print(f"Starting parallel download with {max_workers} workers...")
    with ThreadPoolExecutor(max_workers=max_workers) as executor:
        futures = [executor.submit(download_file, blob, path) for blob, path in download_tasks]

        with tqdm(total=len(futures), desc="Downloading files") as pbar:
            for future in as_completed(futures):
                pbar.update(1)

    print(f"Download completed: {len(blobs)} files downloaded successfully")
    return destination_folder


def get_memory_usage():
    """Get current memory usage in GB."""
    process = psutil.Process(os.getpid())
    return process.memory_info().rss / 1024 / 1024 / 1024  # Convert to GB


def process_tfrecord_batch(file_paths, max_memory_gb=8.0, output_dir="temp_processed"):
    print(f"\nProcessing batch of {len(file_paths)} files...")
    os.makedirs(output_dir, exist_ok=True)
    batch_num = 0
    examples_count = 0
    current_batch = []
    batch_size = 100

    for file_path in file_paths:
        print(f"\nProcessing file: {os.path.basename(file_path)}")
        for record in tf.data.TFRecordDataset([file_path], compression_type="GZIP"):
            example = tf.train.Example()
            example.ParseFromString(record.numpy())

            features = example.features.feature
            episode_id = features["episode_id"].int64_list.value[0]
            goal = features["goal"].bytes_list.value[0].decode("utf-8")

            screenshots_raw = [ss for ss in features["screenshots"].bytes_list.value]
            screenshots_b64 = [base64.b64encode(img).decode("utf-8") for img in screenshots_raw]

            width_list = [w for w in features["screenshot_widths"].int64_list.value]
            height_list = [h for h in features["screenshot_heights"].int64_list.value]
            assert len(width_list) == len(height_list) == len(screenshots_b64)

            # Verify screenshot dimensions match the width/height lists
            for i, screenshot_b64 in enumerate(screenshots_b64):
                img = Image.open(BytesIO(base64.b64decode(screenshot_b64)))
                if img.width != width_list[i] or img.height != height_list[i]:
                    if img.width == height_list[i] and img.height == width_list[i]:
                        continue
                        # Here image seems rotated: ignore for the moment
                        # print("Rotating image!")
                        # img_transposed = img.transpose(Image.Transpose.ROTATE_90)
                        # img_transposed.save("TEMP_transposed_image.png")
                        # screenshots_b64[i] = base64.b64encode(img_transposed.tobytes()).decode("utf-8")
                    else:
                        raise ValueError(
                            f"Screenshot {i} dimensions mismatch: expected {width_list[i]}x{height_list[i]}, got {img.width}x{img.height}"
                        )

            actions = []
            for action_bytes in features["actions"].bytes_list.value:
                action = json.loads(action_bytes.decode("utf-8"))
                actions.append(action)

            step_instructions = [step.decode("utf-8") for step in features["step_instructions"].bytes_list.value]

            processed_example = {
                "episode_id": episode_id,
                "goal": goal,
                "screenshots": screenshots_b64,
                # "screenshot_widths": width_list,
                # "screenshot_heights": height_list,
                # "accessibility_trees": [],  # NOTE: we'll need to add this if we want to implement agents that use the accessibility trees
                "actions": actions,
                "step_instructions": step_instructions,
            }

            current_batch.append(processed_example)
            examples_count += 1

            # Check memory usage and save batch if needed
            current_memory = get_memory_usage()
            if current_memory > max_memory_gb:
                print(f"\nMemory usage high ({current_memory:.2f} GB), saving batch...")
                batch_path = os.path.join(output_dir, f"batch_{batch_num}.json")
                with open(batch_path, "w") as f:
                    json.dump(current_batch, f)
                batch_num += 1
                current_batch = []
                gc.collect()  # Force garbage collection
                print(f"Memory after cleanup: {get_memory_usage():.2f} GB")

        # Save remaining examples in the current batch
        if current_batch:
            batch_path = os.path.join(output_dir, f"batch_{batch_num}.json")
            with open(batch_path, "w") as f:
                json.dump(current_batch, f)
            batch_num += 1
            current_batch = []
            print(f"Processed {len(current_batch)} examples from {os.path.basename(file_path)}")

    print(f"\nBatch processing complete. Total examples processed: {examples_count}")
    return output_dir


def find_tfrecord_files(data_dir):
    tf_files = []
    for root, _, files in os.walk(data_dir):
        for file in files:
            if file.endswith(".tfrecord") or (
                file.startswith("android_control") and "-of-" in file and file.endswith("00020")
            ):
                tf_files.append(os.path.join(root, file))
    return tf_files


def process_files(data_dir, max_files=None, chunk_size=10, max_memory_gb=4.0):
    tf_files = find_tfrecord_files(data_dir)
    print(f"Found {len(tf_files)} TFRecord files")

    if max_files:
        tf_files = tf_files[:max_files]
        print(f"Limited to processing {len(tf_files)} files")

    batches = [tf_files[i : i + chunk_size] for i in range(0, len(tf_files), chunk_size)]
    print(f"Created {len(batches)} batches with up to {chunk_size} files each")

    temp_dir = "temp_processed"
    os.makedirs(temp_dir, exist_ok=True)

    for batch_num, batch in enumerate(tqdm(batches, desc="Processing TFRecord batches")):
        print(f"\nCurrent memory usage: {get_memory_usage():.2f} GB")
        batch_dir = os.path.join(temp_dir, f"batch_{batch_num}")
        process_tfrecord_batch(batch, max_memory_gb=max_memory_gb, output_dir=batch_dir)
        gc.collect()  # Force garbage collection after each batch
        print(f"Memory after batch: {get_memory_usage():.2f} GB")

    return temp_dir


def process_batch(batch, is_train=True, batch_num=0):
    if not batch:
        return

    print(f"\nProcessing {'train' if is_train else 'test'} batch {batch_num + 1}")
    processed_batch = []

    # Process in smaller chunks to manage memory
    chunk_size = 50
    for i in range(0, len(batch), chunk_size):
        chunk = batch[i : i + chunk_size]
        chunk_processed = []

        for example in chunk:
            processed_example = {
                "episode_id": example["episode_id"],
                "goal": example["goal"],
                "screenshots_b64": example["screenshots"],
                # "screenshot_widths": example["screenshot_widths"],
                # "screenshot_heights": example["screenshot_heights"],
                # "accessibility_trees": json.dumps(example["accessibility_trees"]),
                "actions": example["actions"],
                "step_instructions": example["step_instructions"],
            }
            chunk_processed.append(processed_example)

        # Create dataset from chunk and append
        chunk_dataset = Dataset.from_list(chunk_processed)
        processed_batch.append(chunk_dataset)

        # Clear memory
        del chunk_processed
        gc.collect()

        # Print memory usage
        current_memory = get_memory_usage()
        print(f"Memory usage after processing chunk {i // chunk_size + 1}: {current_memory:.2f} GB")

    print(f"Processed {len(batch)} examples in batch {batch_num + 1}")

    # Concatenate all chunks
    if processed_batch:
        final_dataset = processed_batch[0]
        for chunk_dataset in processed_batch[1:]:
            final_dataset = concatenate_datasets([final_dataset, chunk_dataset])
        return final_dataset
    return None


def convert_to_hf_dataset(
    temp_dir,
    splits_file=None,
    test_size=0.2,
    seed=42,
    repo_id=None,
    token=None,
    batch_size=100,
    output_dir="android_control_dataset",
    export_train_split=False,
):
    print("\nStarting conversion to HuggingFace dataset format")

    # Create output directories
    train_dir = os.path.join(output_dir, "train")
    test_dir = os.path.join(output_dir, "test")
    os.makedirs(test_dir, exist_ok=True)
    if export_train_split:
        os.makedirs(train_dir, exist_ok=True)

    print("\nStarting batch processing...")

    # Get all batch files
    batch_files = []
    for root, _, files in os.walk(temp_dir):
        for file in files:
            if file.endswith(".json"):
                batch_files.append(os.path.join(root, file))

    # Load splits if provided
    if splits_file and os.path.exists(splits_file):
        print(f"Using custom splits from file: {splits_file}")
        with open(splits_file) as f:
            splits = json.load(f)
        train_ids = set(splits.get("train", []))
        test_ids = set(splits.get("test", []))
    else:
        print("Using random split with test_size:", test_size)
        import random

        random.seed(seed)
        train_ids = set()
        test_ids = set()

    # Process each batch file
    train_batch_count = 0
    test_batch_count = 0

    for batch_file in tqdm(batch_files, desc="Processing batch files"):
        # Clear memory before loading new batch
        gc.collect()
        current_memory = get_memory_usage()
        print(f"\nMemory usage before loading batch: {current_memory:.2f} GB")

        with open(batch_file, "r") as f:
            batch = json.load(f)

        if splits_file and os.path.exists(splits_file):
            train_batch = [ex for ex in batch if ex["episode_id"] in train_ids] if export_train_split else []
            test_batch = [ex for ex in batch if ex["episode_id"] in test_ids]
        else:
            train_batch = []
            test_batch = []
            for ex in batch:
                if ex["episode_id"] not in train_ids and ex["episode_id"] not in test_ids:
                    if random.random() < (1 - test_size):
                        if export_train_split:
                            train_ids.add(ex["episode_id"])
                            train_batch.append(ex)
                    else:
                        test_ids.add(ex["episode_id"])
                        test_batch.append(ex)

        # Clear original batch to free memory
        del batch
        gc.collect()

        # Use smaller chunk size to prevent large parquet files
        chunk_size = 300
        # Process and save train batch
        if train_batch:
            # Process in smaller chunks
            for i in range(0, len(train_batch), chunk_size):
                chunk = train_batch[i : i + chunk_size]
                chunk_processed = []

                for example in chunk:
                    processed_example = {
                        "episode_id": example["episode_id"],
                        "goal": example["goal"],
                        "screenshots_b64": example["screenshots"],
                        "actions": example["actions"],
                        "step_instructions": example["step_instructions"],
                    }
                    chunk_processed.append(processed_example)

                # Save chunk to disk with max_shard_size parameter
                chunk_dataset = Dataset.from_list(chunk_processed)
                chunk_path = os.path.join(train_dir, f"batch_{train_batch_count}")
                chunk_dataset.save_to_disk(chunk_path, max_shard_size="200MB")
                train_batch_count += 1

                # Clear memory
                del chunk_processed
                del chunk_dataset
                gc.collect()

            del train_batch
            gc.collect()

        # Process and save test batch
        if test_batch:
            # Process in smaller chunks
            for i in range(0, len(test_batch), chunk_size):
                chunk = test_batch[i : i + chunk_size]
                chunk_processed = []

                for example in chunk:
                    processed_example = {
                        "episode_id": example["episode_id"],
                        "goal": example["goal"],
                        "screenshots_b64": example["screenshots"],
                        "actions": example["actions"],
                        "step_instructions": example["step_instructions"],
                    }
                    chunk_processed.append(processed_example)

                # Save chunk to disk with max_shard_size parameter
                chunk_dataset = Dataset.from_list(chunk_processed)
                chunk_path = os.path.join(test_dir, f"batch_{test_batch_count}")
                chunk_dataset.save_to_disk(chunk_path, max_shard_size="200MB")
                test_batch_count += 1

                # Clear memory
                del chunk_processed
                del chunk_dataset
                gc.collect()

            del test_batch
            gc.collect()

        # Print memory usage after processing batch
        current_memory = get_memory_usage()
        print(f"Memory usage after processing batch: {current_memory:.2f} GB")

    # If streaming to hub, upload the saved datasets
    if repo_id:
        print(f"\nPreparing to stream dataset to Hugging Face Hub: {repo_id}")
        from huggingface_hub import HfApi

        api = HfApi(token=token)

        # Upload train directory if needed
        if export_train_split:
            print("\nUploading train split to Hugging Face Hub...")
            train_datasets = [
                Dataset.load_from_disk(os.path.join(train_dir, f"batch_{i}")) for i in range(train_batch_count)
            ]
        test_datasets = [Dataset.load_from_disk(os.path.join(test_dir, f"batch_{i}")) for i in range(test_batch_count)]

        dataset_dict = {"test": concatenate_datasets(test_datasets)}
        if export_train_split:
            dataset_dict["train"] = concatenate_datasets(train_datasets)

        dataset = DatasetDict(dataset_dict)

        # Push to hub with max_shard_size parameter
        dataset.push_to_hub(repo_id, token=token, max_shard_size="200MB")

        # Clean up temporary directories
        print("\nCleaning up temporary files...")
        import shutil

        if export_train_split:
            shutil.rmtree(train_dir)
        shutil.rmtree(test_dir)
        shutil.rmtree(temp_dir)

        return None

    # If not streaming to hub, load and return the dataset
    print("\nLoading final dataset from disk...")
    dataset_dict = {}

    if export_train_split:
        train_datasets = []
        for i in range(train_batch_count):
            train_datasets.append(Dataset.load_from_disk(os.path.join(train_dir, f"batch_{i}")))
        dataset_dict["train"] = concatenate_datasets(train_datasets) if train_datasets else None

    test_datasets = []
    for i in range(test_batch_count):
        test_datasets.append(Dataset.load_from_disk(os.path.join(test_dir, f"batch_{i}")))
    dataset_dict["test"] = concatenate_datasets(test_datasets) if test_datasets else None

    # Clean up temporary directories
    import shutil

    shutil.rmtree(temp_dir)

    # Return the dataset
    return DatasetDict(dataset_dict)


def main():
    args = parse_arguments()
    start_time = time.time()

    print("\n=== Starting AndroidControl Dataset Processing ===")
    print("Configuration:")
    print(f"- Bucket: {args.bucket}")
    print(f"- Prefix: {args.prefix}")
    print(f"- Output directory: {args.output_dir}")
    print(f"- Data directory: {args.data_dir}")
    print(f"- Max files: {args.max_files}")
    print(f"- Workers: {args.workers}")
    print(f"- Chunk size: {args.chunk_size}")
    print(f"- Test size: {args.test_size}")
    print(f"- Batch size: {args.batch_size}")
    print(f"- Max memory: {args.max_memory_gb} GB")
    print(f"- Skip train split: {args.skip_train_split}")

    if args.download:
        print("\n=== Download Phase ===")
        download_files_parallel(args.bucket, args.prefix, args.data_dir)

    print("\n=== Processing Phase ===")
    temp_dir = process_files(
        args.data_dir, max_files=args.max_files, chunk_size=args.chunk_size, max_memory_gb=args.max_memory_gb
    )

    print("\n=== Conversion Phase ===")
    dataset = convert_to_hf_dataset(
        temp_dir,
        splits_file=args.splits_file,
        test_size=args.test_size,
        seed=args.seed,
        repo_id=args.repo_id if hasattr(args, "repo_id") else None,
        token=args.token if hasattr(args, "token") else os.environ.get("HF_TOKEN"),
        batch_size=args.batch_size,
        output_dir=args.output_dir,
        export_train_split=not args.skip_train_split,
    )

    if dataset:  # Only save to disk if not streaming to hub
        print(f"\nSaving dataset to {args.output_dir}")
        dataset.save_to_disk(args.output_dir)

    elapsed_time = time.time() - start_time
    print("\n=== Processing Complete ===")
    print(f"Total processing time: {elapsed_time:.2f} seconds")
    if dataset:
        print(f"Dataset saved to {args.output_dir}")
        if not args.skip_train_split:
            print(
                f"Dataset contains {len(dataset['train'])} training examples and {len(dataset['test'])} test examples"
            )
        else:
            print(f"Dataset contains {len(dataset['test'])} test examples")
    else:
        print("Dataset streamed to Hugging Face Hub")


if __name__ == "__main__":
    main()
