import pandas as pd
import argparse
from src.custom_utils import get_logger
from pathlib import Path

logger = get_logger(__name__)

if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument('--input_path', type=str, required=True, help='Path to the input TSV file')
    parser.add_argument('--dataset_passages_path', type=str)
    parser.add_argument('--passages_nrows', type=int, default=None, help='Number of rows to read from the passages file for debugging')
    parser.add_argument('--top_k', type=int)
    
    args = parser.parse_args()

    tsv = pd.read_csv(args.input_path, sep='\t', header=0)
    if args.passages_nrows is not None:
        logger.info(f"Limiting to {args.passages_nrows} rows for debugging")
        tsv = tsv[tsv.pidx < args.passages_nrows].reset_index()
    logger.info(f"Read {len(tsv)} rows from {args.input_path}")

    if args.top_k is not None:
        tsv = pd.concat([v.head(args.top_k) for k, v in tsv.groupby('qid', sort=False)], ignore_index=True, axis=0)
        logger.info(f"Filtered to top {args.top_k} rows per query, total rows: {len(tsv)}")

    data = pd.read_json(args.dataset_passages_path, lines=True, nrows=args.passages_nrows)
    logger.info(f"Read {len(data)} rows from {args.dataset_passages_path}")
    data = data.iloc[tsv.pidx].passage_text.reset_index(drop=True)
    
    tsv['passage_text'] = data
    assert not tsv.passage_text.isna().any()

    output_path = Path(args.input_path).with_suffix('.jsonl')
    tsv.to_json(output_path, orient='records', lines=True)
    logger.info(f"Saved {len(tsv)} rows to {output_path}")