import os
import concurrent.futures
import pandas as pd
import metrics as m
from url_scrape import url_scrape
import zipfile
import argparse
from functools import partial
from clean_text import clean_text


def process_row(row_data, raw_data_folder):
    index, row = row_data  # Unpack the row_data tuple

    # Use the 'url' column from the row to fetch content
    url = row['url']
    text = url_scrape(url)

    if text:
        text = clean_text(text)

        # Create the raw_data folder if it doesn't exist
        if not os.path.exists(raw_data_folder):
            os.makedirs(raw_data_folder, exist_ok=True)

        # Save the text content in a file using fm_id, url_id, url_datetime
        file_name = f"{row['fm_id']}_{row['url_id']}_{row['url_datetime']}"
        file_name = file_name.replace('-', '').replace(':', '').replace(' ', '')
        text_file_path = os.path.join(raw_data_folder, file_name + '.txt')
        zip_file_path = os.path.join(raw_data_folder, file_name + '.zip')

        with open(text_file_path, "w") as file:
            file.write(text)
        
        with zipfile.ZipFile(zip_file_path, mode='w', compression=zipfile.ZIP_DEFLATED, compresslevel=9) as zipper:
            zipper.write(text_file_path, file_name + '.txt')
        
        os.remove(text_file_path)
        
        # Process the text using process_metrics
        metrics = m.process_metrics(text)

        # Convert the row to a dictionary, update it with the metrics, and convert it back to a Series
        row_dict = row.to_dict()
        row_dict.update(metrics)
        row = pd.Series(row_dict)

    return index, row  # Return the index and row as a tuple

script_dir = os.path.dirname(os.path.realpath(__file__))
processing_txt_path = os.path.join(script_dir, 'processing.txt')

def process_csvs(source_folder, target_folder, raw_data_folder):
    # Create the target_folder if it doesn't exist
    if not os.path.exists(target_folder):
        os.makedirs(target_folder)

    # Check if processing.txt exists, if not, create it.
    if not os.path.exists(processing_txt_path):
        open(processing_txt_path, 'w').close()

    # List all pickle files in the source_folder
    pkl_files = [file for file in os.listdir(source_folder) if file.endswith(".pkl")]
    
    for pkl_file in pkl_files:
        # Check if file already processed and exists in target_folder
        if pkl_file in os.listdir(target_folder):
            continue

        # Check if file is currently being processed by another instance
        if pkl_file in open(processing_txt_path).read():
            continue

        # Add the file to processing.txt
        with open(processing_txt_path, 'a') as f:
            f.write(pkl_file + '\n')

        source_pkl_path = os.path.join(source_folder, pkl_file)
        df = pd.read_pickle(source_pkl_path, compression='zip')

        # Process rows concurrently
        process_row_partial = partial(process_row, raw_data_folder=raw_data_folder)
        with concurrent.futures.ThreadPoolExecutor() as executor:
            data = list(executor.map(process_row_partial, df.iterrows()))

        # Reconstruct the DataFrame from the processed data
        df = pd.concat([row for index, row in data], axis=1).transpose()

        # Remove the 'url' column
        df = df.drop(columns=['url'])

        # Save the updated DataFrame to the target_folder
        target_pkl_path = os.path.join(target_folder, pkl_file)
        df.to_pickle(target_pkl_path, compression='zip')

        # Remove the file from processing.txt
        with open(processing_txt_path, 'r') as f:
            lines = f.readlines()
        with open(processing_txt_path, 'w') as f:
            for line in lines:
                if line.strip("\n") != pkl_file:
                    f.write(line)
                    
    if os.path.exists(processing_txt_path):
        os.remove(processing_txt_path)

if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="Process CSV files for scraping.")
    parser.add_argument("source_folder", help="Path to the source folder containing the pkl files for scraping.")
    parser.add_argument("result_folder", help="Path to the folder to store results.")
    parser.add_argument("raw_data_folder", help="Path to the folder to store raw data.")
    args = parser.parse_args()

    os.makedirs(args.raw_data_folder, exist_ok=True)
    os.makedirs(args.result_folder, exist_ok=True)

    try:
        process_csvs(args.source_folder, args.result_folder, args.raw_data_folder)
    except Exception as e:
        print(f"An error occurred: {e}")