import os
import boto3

from s3Ops import (
    list_files_in_s3_folder
)

from utils.prompts import  DOCUMENT_CLASSIFY_PROMPT

from backendAPIs import (
    get_onboarding_status,
    update_onboarding_status
)

bucket_name = os.getenv('S3_UPLOAD_BUCKET_NAME')

def document_classifier(record_details):

    FINAL_DOCUMENT_CLASSIFY_PROMPT = DOCUMENT_CLASSIFY_PROMPT.format(document=record_details)

    json_string = llm_bedrock.invoke(FINAL_DOCUMENT_CLASSIFY_PROMPT).content

    print(json_string)

    pattern = r'<output>(.*?)</output>'
    match = re.search(pattern, json_string, re.DOTALL)

    if match:
        return match.group(1).strip()
    else:
        return ""


def get_pdf_details(key):
    # Create a temporary file for the PDF
    with tempfile.NamedTemporaryFile(delete=False, suffix='.pdf', dir='/tmp') as temp_file:
        local_file_path = temp_file.name
    # Download the file from S3
    s3_client.download_file(bucket_name, key, local_file_path)

    with pdfplumber.open(local_file_path) as pdf:
        list_pages = []

        for page in pdf.pages:
            text = page.extract_text()
            if text:
                list_pages.append(text.strip())

        pdf_string = "".join(list_pages)

    if len(pdf_string.strip()) > 0:
        # Clean up the temporary PDF file
        os.remove(local_file_path)
        return pdf_string
    else:  # start image processing
        print("Starting image processing")
        # Create a temporary directory for storing images
        with tempfile.TemporaryDirectory(dir='/tmp') as temp_dir:
            pdf_images_list = pdf_to_images(local_file_path, temp_dir)
            result = read_images(pdf_images_list)

        # Clean up the temporary PDF file
        os.remove(local_file_path)

        return result

def process_single_pdf(pdf_file):
    pdf_content = get_pdf_details(key)
    classification = document_classifier(pdf_content)
    return pdf_file, classification


def pdf_classifier(files_list):

    categories = {
        "LPA(Limited Partner Agreement)": [],
        "Capital Statements": [],
        "Sub Docs(Subscription Document)": [],
        "Side Letter": [],
        "Capital Call Notices": [],
        "Distribution Notices": [],
        "Fund Reports": [],
        "K1s": [],
        "Other": []
    }

    with ThreadPoolExecutor(max_workers=10) as executor:
        future_to_pdf = {executor.submit(process_single_pdf, pdf_file): pdf_file for pdf_file in pdf_files}

        for future in as_completed(future_to_pdf):
            pdf_file = future_to_pdf[future]
            try:
                pdf_file, classification = future.result()

                # Determine which category the PDF belongs to
                found_category = False
                for category in categories.keys():
                    if category.lower() in classification.lower():
                        categories[category].append(pdf_file)
                        found_category = True
                        break

                if not found_category:
                    categories["Other"].append(pdf_file)

            except Exception as exc:
                print(f'{pdf_file} generated an exception: {exc}')
                categories["Other"].append(pdf_file)

        return categories

def validate_input_files(file_list):
    """
    Function to find specific files in a list based on naming patterns.

    Args:
        file_list (str): List of files to validate.

    Returns:
        dict: Dictionary containing file names that match the required patterns.
        str: Error message if any required file is missing.
    """
    # Patterns to look for
    required_files = {
        # "lpa": "_LP",
        "partner": "_partner",
        "financials": "_financials",
        "bankTransactions": "_bank_transactions",
        "journals": "_journals",
        "fund_performance": "_fund-performance"
    }

    # Dictionary to hold found file names
    found_files = {
        # "lpa": "",
        "partner": None,
        "financials": None,
        "bankTransactions": None,
        "journals": None,
        "fund_performance": None
    }

    pdf_files = []

    try:

        if file_list:
            # Loop through files in the list
            for file_name in file_list:
                if required_files['partner'] in file_name:
                    found_files['partner'] = file_name
                elif required_files['financials'] in file_name:
                    found_files['financials'] = file_name
                elif required_files['bankTransactions'] in file_name:
                    found_files['bankTransactions'] = file_name
                elif required_files['journals'] in file_name:
                    found_files['journals'] = file_name
                elif required_files['fund_performance'] in file_name:
                    found_files['fund_performance'] = file_name
                elif file_name.lower().endswith('.pdf'):
                    pdf_files.append(file_name)

        else:
            return(f"Error: No files found in folder.")

        # if len(pdf_files) > 0:
        #     classified_pdf_files = pdf_classifier(pdf_files)
        #
    except FileNotFoundError:
        return f"Error: Unable to find file.",""

    # Check if all required files are found
    missing_files = [key for key, value in found_files.items() if value is None]

    if missing_files:
        return f"Error: Missing required files - {', '.join(missing_files)}",""

    # return found_files,classified_pdf_files

    return found_files , pdf_files


def pretty_print_files(files_dict):
    """
    Pretty print the dictionary containing file names.
    
    Args:
        files_dict (dict): Dictionary where the keys are the file categories, and the values are the file names.
    """
    file_message = ""
    for key, value in files_dict.items():
        file_message += f"- {key.replace('_', ' ').capitalize()}: {value} \n"

    return file_message


def run_file_validation():

    bucket_name = os.getenv("S3_UPLOAD_BUCKET_NAME")
    step_number = 1
    # Array index starts from 0
    step_number -= 1


    print(f"Getting files from bucket {bucket_name}")
    success_message = ""
    error_message = ""
    status = ""

    try:

        onboarding_id = os.getenv('ONBOARDING_ID') 
        print(f"Get onBoardingStatus for {onboarding_id}")
        response = get_onboarding_status(onboarding_id)
        # Get current onboarding status values.
        onboarding_status = response['data']['data']
        success_message = onboarding_status['steps'][step_number]['summary']
        error_message = onboarding_status['steps'][step_number]['errorMessage']
        status = onboarding_status['steps'][step_number]['status']
        folder_path = onboarding_status['path']

        print(f"Pulling files list from {bucket_name}/{folder_path}")
        file_list = list_files_in_s3_folder(bucket_name, folder_path)
        print("S3 files:", file_list)
        found_files ,classified_files = validate_input_files(file_list)

        print("Files found:", found_files)

        if isinstance(found_files, dict):
            success_message = pretty_print_files(found_files)
            status =  "COMPLETE"
        #Validation failed
        else:
            error_message = found_files
            status = "FAILED"
            success_message = ""

        response = update_onboarding_status(step_number, status, error_message, success_message)
        if 'error' in response:
            print(f"Failed to fetch onboarding status: Error: {response['error']}")
            print(f"Status Code: {response['status_code']}")
            return False , False
        
        return found_files , classified_files

    except Exception as e:
        return False , False