import json
import os

import click
import jsonlines
from rich import print


@click.command()
@click.option("--data_dir", type=str, required=True, help="Directory containing data files.")
@click.option("--input_file", type=str, required=True, help="Path to the raw JSON file.")
@click.option("--output_file", type=str, required=True, help="Path to the processed JSON file.")
def main(data_dir: str, input_file: str, output_file: str) -> None:
    """
    Process the JSON file containing Sotopia episodes.
    """
    with open(os.path.join(data_dir, input_file), "r") as f:
        episodes = [json.loads(line) for line in f]
    print("[bold green]Successfully loaded episodes:[/bold green]")

    behavior_cloning_episodes = []
    for episode in episodes:
        if episode["experiment_model_name_pairs"][0] == "gpt-4" and episode["experiment_model_name_pairs"][1] == "gpt-4":
            behavior_cloning_episodes.append(episode)
    print(f"[bold green]Successfully filtered {len(behavior_cloning_episodes)} episodes [/bold green]")
    with open(os.path.join(data_dir, output_file), "w") as f:
        jsonlines.Writer(f).write_all(behavior_cloning_episodes)

if __name__ == "__main__":
    main()