import concurrent.futures
import traceback
import re
import os
import pandas as pd

import utils
import config
import preprocessing_funcs

COMPLEX_LIST = []


def count_reimbursements(page_text: str) -> int:
    """Count number of reimbursement items found on a page text, through use of a
    regex that catches either '$' or '%'.

    Args:
        page_text (str): Input page text, as a string

    Returns:
        int: Number of reimbursement items
    """
    pattern = r"[%$]"
    matches = re.findall(pattern, page_text)
    return len(matches)


def detect_complex(filename, contract_text):
    print(f"Processing {filename}")
    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)
    rate_dict = {}
    table_dict = {}
    for page_num, page_text in text_dict.items():
        rate_dict[page_num] = count_reimbursements(page_text)
        table_dict[page_num] = "-------Table Start--------" in page_text

    page_count, consecutive_page_count, consecutive_rate_count = 0, 0, 0

    for page_num in table_dict.keys():
        if table_dict[page_num] and rate_dict[page_num] > 0:
            consecutive_rate_count += rate_dict[page_num]
            consecutive_page_count += 1
        else:
            if consecutive_rate_count >= config.RATE_THRESHOLD:
                page_count += consecutive_page_count
                consecutive_page_count, consecutive_rate_count = 0, 0

    COMPLEX_LIST.append(
        {"Filename": filename, "COMPLEX_FLAG": page_count >= config.TABLE_THRESHOLD}
    )


def process_item(contract):
    data = s3_client.get_object(Bucket=config.BUCKET, Key=contract)
    contents = data["Body"].read()
    contract_text = contents.decode("utf-8")

    detect_complex(contract, contract_text)


s3_client = config.S3_CLIENT

file_list = []
paginator = s3_client.get_paginator("list_objects_v2")
pages = paginator.paginate(Bucket=config.BUCKET, Prefix=config.PREFIX)
for page in pages:
    for obj in page["Contents"]:
        if not obj["Key"].endswith("/"):
            file_list.append(obj["Key"])

contract_list = sorted(file_list)
files = {}

with concurrent.futures.ThreadPoolExecutor(
    max_workers=config.COMPLEX_MAX_WORKERS
) as executor:
    # Process each item individually
    futures = [executor.submit(process_item, contract) for contract in contract_list]

    # Wait for all futures to complete
    for future in concurrent.futures.as_completed(futures):
        # Retrieve any exceptions raised by the thread
        try:
            future.result()
        except Exception as e:
            print(f"Error: {e}")

complex_df = pd.DataFrame(COMPLEX_LIST)
# complex_df.to_csv(os.path.join(config.COMPLEX_OUTPUT_PATH, config.COMPLEX_OUTPUT_FILENAME))

complex_df.to_csv(
    os.path.join(
        config.COMPLEX_OUTPUT_PATH, "TX_priority_tin_complex_categorization.csv"
    )
)