import pandas as pd
from tqdm import tqdm
from src.logger import get_logger

logger = get_logger(__name__)


def extract_data(files: list[str]) -> pd.DataFrame:
    logger.info(f"Starting data extraction from {len(files)} files.")

    result_df = pd.DataFrame()
    for filename in tqdm(files, desc="Processing files"):
        chunk_df = pd.read_csv(filename, sep="\t")

        chunk_filter = (chunk_df["language"] == "pl") & ~chunk_df[
            "caption_reference_description"
        ].isna()
        chunk_df = chunk_df[chunk_filter]

        result_df = pd.concat([result_df, chunk_df], ignore_index=True)

    logger.info(f"Extracted {len(result_df)} rows after filtering.")
    return result_df


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

    return wit_df.merge(pages_with_categories_df, how="left", on="page_title")


if __name__ == "__main__":
    import os
    import argparse

    parser = argparse.ArgumentParser()

    parser.add_argument(
        "--input_dir",
        type=str,
        required=True,
        help="Path to the directory containing raw WIT files.",
    )
    parser.add_argument(
        "--categories_path",
        type=str,
        required=True,
        help="Path to the file containing list of categories",
    )
    parser.add_argument(
        "--output_path", type=str, required=True, help="Path to save the processed WIT data file"
    )

    args = parser.parse_args()

    input_files = [
        os.path.join(args.input_dir, f) for f in os.listdir(args.input_dir) if f.endswith(".tsv")
    ]
    logger.info(f"Found {len(input_files)} input files.")
    wit_df = extract_data(input_files)

    logger.info(f"Loading category tree from {args.categories_path}")
    category_tree_df = pd.read_json(args.categories_path, lines=args.categories_path.endswith(".jsonl"))

    wit_df = merge_category_tree(wit_df, category_tree_df)

    logger.info(f"Saving processed data to {args.output_path}")
    os.makedirs(os.path.dirname(args.output_path), exist_ok=True)
    use_lines = args.output_path.endswith(".jsonl")
    wit_df.to_json(
        args.output_path, orient="records", lines=use_lines, indent=2 if not use_lines else None
    )
