import os
import pandas as pd


def merge_category_tree(
    wit_df: pd.DataFrame, category_tree_df: pd.DataFrame
) -> pd.DataFrame:
    pages_categories_df = (
        category_tree_df[['title', 'subpages']]
            .explode('subpages')
            .rename(columns={'title': 'categories', 'subpages': 'page_title'})
            .groupby('page_title')['categories']
            .agg(set)
            .reset_index()
    )

    merged_df = wit_df.merge(
        pages_categories_df, how='left', on='page_title'
    )

    return merged_df




if __name__ == "__main__":
    import argparse

    parser = argparse.ArgumentParser()
    parser.add_argument(
        "--input_file",
        type=str,
        required=True,
        help="Path to the WIT dataset directory",
    )
    parser.add_argument(
        "--category_tree_file",
        type=str,
        required=True,
        help="Path to the category tree JSON file",
    )
    parser.add_argument(
        "--output_file",
        type=str,
        required=True,
        help="Path to save the merged category tree JSON file",
    )
    args = parser.parse_args()

    wit_df = pd.read_json(args.input_file, lines=args.input_file.endswith(".jsonl"))
    category_tree_df = pd.read_json(
        args.category_tree_file, lines=args.category_tree_file.endswith(".jsonl")
    )

    merged_df = merge_category_tree(wit_df, category_tree_df)

    os.makedirs(os.path.dirname(args.output_file), exist_ok=True)
    use_lines = args.output_file.endswith(".jsonl")
    merged_df.to_json(
        args.output_file,
        orient="records",
        lines=use_lines,
        indent=2 if not use_lines else None,
    )
