import argparse
import os
from utils import read_local
import preprocessing_funcs
import ac_funcs


def parse_arguments():
    parser = argparse.ArgumentParser(description="Smart Chunking Tester")
    parser.add_argument(
        "--input_dir", help="Input directory (local path or S3 URI)", default="src/new/"
    )
    return parser.parse_args()


def main():
    args = parse_arguments()
    input_dir = args.input_dir

    for file in os.listdir(input_dir):
        full_path = os.path.join(input_dir, file)
        print("--" * 20)
        print(file)
        contract_text = read_local(full_path)
        contract_text = preprocessing_funcs.clean_newlines(contract_text)
        contract_text = preprocessing_funcs.clean_law_symbols(contract_text)
        text_dict = preprocessing_funcs.split_text(
            contract_text
        )  # return a dictionary with keys - page_num (str), values as the page_text

        irs_answers = ac_funcs.tin_regex(filename=str(file), text_dict=text_dict)
        if irs_answers is None:
            irs_answers = {}
            irs_answers["PROV_GROUP_TIN"] = "N/A"
            irs_answers["PROV_TIN_OTHER"] = "N/A"

        print(irs_answers)


if __name__ == "__main__":
    main()