import os
import pandas as pd
import time
import csv
import json

import config
import keywords
import preprocess
import preprocessing_funcs
import table_funcs
import bottom_up_funcs
import top_down_funcs
import conditional_funcs
import postprocess
import merge_funcs
import utils
import ac_funcs
import prompts
import valid
import claude_funcs


def run_ac_prompts(file_object):

    ac_answers_dict = {}
    log_data = []
    chunk_data = []

    ################## INITIATE PROCESSING ##################
    filename, contract_text = file_object
    print(f"Processing AC for {filename}...")
    print(f"Total contract word count: {len(contract_text.split())}")

    ################## PREPROCESS - SMART CHUNK ##################
    text_dict, exhibit_pages, num_pages, ac_chunks = preprocess.preprocess(
        contract_text, filename, fields="ac"
    )
    print(f"AC Preprocessing Complete - {filename}")

    ################## RUN REGEX FUNCTIONALITY ##################
    if config.IRS_REGEX_CHECK == True:
        irs_answers = ac_funcs.tin_regex(filename, text_dict)
        ac_answers_dict.update(irs_answers)

    ################## RUN CHUNKED PROMPTS ##################

    ## Check `ac_funcs.create_prompt()` vs. `prompts.AC_*_template()`
    all_fields = set(prompts.AC_DICT.keys())
    field_groups = [
        field_group for field_group in keywords.GROUPED_KEYWORD_MAPPINGS.keys()
    ]

    # list to append the fields with no matching keywords
    no_keyword_matches = []

    for field_group in field_groups:
        keyword_dict = keywords.GROUPED_KEYWORD_MAPPINGS[field_group]["keywords"]
        fields = keywords.GROUPED_KEYWORD_MAPPINGS[field_group]["fields"]
        fields.sort()
        questions = {}
        for field in fields:
            questions[field] = prompts.AC_DICT[field]

        # Use chunk if available and different from full context, otherwise use full context
        ############# Important Note #############
        # the above logic is changed
        # New Logic - chunk if available and different from full context, otherwise add the field to full context list
        if (
            field_group in ac_chunks
            and ac_chunks[field_group].strip()
            and ac_chunks[field_group] != contract_text
        ):

            context = ac_chunks[field_group]
            context = context[0 : min(100000, len(context)) - 1]
            context_type = "Smart Chunking"

            if len(fields) == 1:
                question = questions[
                    fields[0]
                ]  # Take the one question for the single field included
                prompt = prompts.AC_SINGLE_FIELD_TEMPLATE(context, question)
            elif len(fields) > 1:
                prompt = prompts.AC_MULTI_FIELD_TEMPLATE(context, questions)
            else:
                raise ValueError("Field group has no defined fields")

            try:
                field_group_answer_raw = claude_funcs.invoke_claude(
                    prompt, config.MODEL_ID_CLAUDE35_SONNET, filename, 8192
                )

                if len(fields) == 1:
                    field_group_answer_dict = {field: field_group_answer_raw}
                elif len(fields) > 1:
                    field_group_answer_raw = field_group_answer_raw.replace("_DATE", "_DT")
                    field_group_answer_dict = ac_funcs.json_parsing_search(field_group_answer_raw, fields)

                # TODO: Validate field names coming out of Claude, check against fields in group, make sure they're all there, make sure there aren't any extras
                for field in fields:
                    field_answer = field_group_answer_dict[field]
                    ac_answers_dict[field] = field_answer

                log_data.append(
                    {
                        "Field": field,
                        "Prompt": f"Question: {questions}\n\nContext: {context}",  # Include full context
                        "Context Type": context_type,
                        "Response": field_answer,
                    }
                )

                chunk_data.append(
                    {
                        "Contract name": filename,
                        "Field group": field_group,
                        "Field list": fields,
                        "Methodology": ac_chunks[field_group + "_methodology"],
                        "Case sensitivity": ac_chunks[field_group + "_case"],
                        "Page list": ac_chunks[field_group + "_pages"],
                        "Chunk size": len(context),
                        "% Reduction": f"{100*(1-len(context)/len(contract_text)):0.2f}%",
                    }
                )
            except Exception as e:
                print(f"Error processing field {field}: {type(e)}, {str(e)}")
                ac_answers_dict[field] = f"Error: {str(e)}"
                log_data.append(
                    {
                        "Field": field,
                        "Prompt": f"Question: {questions}\n\nContext: {context}",  # Include full context
                        "Context Type": context_type,
                        "Response": f"Error: {str(e)}",
                    }
                )
                chunk_data.append(
                    {
                        "Contract name": filename,
                        "Field group": field_group,
                        "Field list": fields,
                        "Methodology": ac_chunks[field_group + "_methodology"],
                        "Case sensitivity": ac_chunks[field_group + "_case"],
                        "Page list": ac_chunks[field_group + "_pages"],
                        "Chunk size": len(context),
                        "% Reduction": f"{100*(1-len(context)/len(contract_text)):0.2f}%",
                    }
                )
        else:
            no_keyword_matches.append(field)

    ################## RUN FULL CONTEXT PROMPTS ##################
    # Process 'full_context' fields together

    smartly_chunked = []
    for field_groups in keywords.GROUPED_KEYWORD_MAPPINGS.keys():
        for fields in keywords.GROUPED_KEYWORD_MAPPINGS[field_groups]["fields"]:
            smartly_chunked.append(fields)
    
    # Instantiate full context fields from those that weren't smart chunked
    full_context_fields = [
        field for field in prompts.AC_DICT if field not in smartly_chunked
    ]

    # Add fields with no keyword matches to full context list
    if len(no_keyword_matches) > 0:
        full_context_fields.extend(no_keyword_matches)
        
    # Filter out fields that already have answers
    # e.g., some have already been filled out by regex chunking
    full_context_fields = [
        field for field in full_context_fields if field not in ac_answers_dict.keys()
    ]

    full_context_questions = {
        field: prompts.AC_DICT.get(field) for field in full_context_fields
    }
    if full_context_questions:
        full_context_prompt = ac_funcs.create_prompt(
            contract_text[0 : min(100000, len(contract_text) - 1)],
            question=full_context_questions,
        )
        try:
            full_context_answers = ac_funcs.get_ac_answer(
                full_context_prompt,
                filename,
                fields=list(full_context_questions.keys()),
            )
            ac_answers_dict.update(full_context_answers)

            for field, answer in full_context_answers.items():
                log_data.append(
                    {
                        "Field": field,
                        "Prompt": f"Question: {full_context_questions[field]}\n\nContext: {contract_text}",  # Include full context
                        "Context Type": "Full Context (High Accuracy)",
                        "Response": answer,
                    }
                )
        except Exception as e:
            print(f"Error processing high accuracy fields: {str(e)}")
            for field in full_context_questions.keys():
                ac_answers_dict[field] = f"Error: {str(e)}"
                log_data.append(
                    {
                        "Field": field,
                        "Prompt": f"Question: {full_context_questions[field]}\n\nContext: {contract_text}",  # Include full context
                        "Context Type": "Full Context (High Accuracy)",
                        "Response": f"Error: {str(e)}",
                    }
                )

    ################## AC CONDITIONAL PROMPTS ##################
    # Non-Renewal - Days
    nrd_prompt = prompts.AC_SINGLE_FIELD_TEMPLATE(
        ac_answers_dict["NON_RENEWAL_LANGUAGE"], prompts.AC_DICT["NON_RENEWAL_DAYS"]
    )
    nrd_answer = claude_funcs.invoke_claude(
        nrd_prompt, config.MODEL_ID_CLAUDE35_SONNET, filename, 8192
    )
    ac_answers_dict["NON_RENEWAL_DAYS"] = nrd_answer
    ## TODO - add Notice Provider Name/Address here
    # Clean Up Contract Effective Date
    if (
        utils.is_empty(ac_answers_dict["CONTRACT_EFFECTIVE_DT"])
        and "meridian" in ac_answers_dict["PAYER_NAME"].lower()
    ):
        date_prompt = prompts.AC_EFFECTIVE_DATE_CLEANUP(
            contract_text[0 : min(100000, len(contract_text)) - 1]
        )
        date_answer = claude_funcs.invoke_claude(
            date_prompt, config.MODEL_ID_CLAUDE35_SONNET, filename, 124
        )
        ac_answers_dict["CONTRACT_EFFECTIVE_DT"] = date_answer

    ################## ENSURE ALL FIELDS ARE PRESENT ##################
    for field in all_fields:
        if field not in ac_answers_dict:
            print(f"Field not found in results. {field}")

    ################## CREATE OUTPUT DIRECTORIES ##################
    base_filename = os.path.splitext(filename)[0].strip()
    output_dir = os.path.join(config.OUTPUT_DIRECTORY, base_filename)
    os.makedirs(output_dir, exist_ok=True)

    ################## POSTPROCESS ##################
    for key in ac_answers_dict.keys():
        print(key, ac_answers_dict[key])

    ac_df = pd.DataFrame([ac_answers_dict])
    ac_df = postprocess.ac_postprocess(ac_df, filename, num_pages)

    ################## WRITE TO OUTPUT ##################
    ac_df.to_csv(os.path.join(output_dir, config.AC_RESULTS_NAME), index=False)

    chunk_file_path = "chunk_log.csv"
    with open(chunk_file_path, "a", newline="", encoding="utf-8") as chunk_file:
        fieldnames = [
            "Contract name",
            "Field group",
            "Field list",
            "Methodology",
            "Case sensitivity",
            "Page list",
            "Chunk size",
            "% Reduction",
        ]
        writer = csv.DictWriter(chunk_file, fieldnames=fieldnames)
        writer.writeheader()
        writer.writerows(chunk_data)

    print(f"Chunk log written to {chunk_file_path}")

    return ac_answers_dict


def run_b_prompts(file_object):
    ################## INITIATE PROCESSING ##################
    filename, contract_text = file_object
    print(f"Processing B for {filename}...")

    ################## PREPROCESS ##################
    text_dict, exhibit_pages, num_pages, ac_chunks = preprocess.preprocess(
        contract_text, filename, fields="b"
    )
    print(f"B Preprocessing Complete - {filename}")

    ################## RUN BOTTOM UP PROMPTS ##################
    bu_results = bottom_up_funcs.run_bottom_up(
        filename, text_dict
    )  # Returns list of dictionaries
    print(f"B Bottom Up Complete - {filename}")

    ################## RUN TOP DOWN PROMPTS ##################
    combined_results = top_down_funcs.run_top_down(
        filename, text_dict, bu_results, exhibit_pages
    )  # Returns list of dictionaries
    print(f"B Top Down Complete - {filename}")

    ################## RUN CONDITIONAL PROMPTS ##################
    conditional_results = conditional_funcs.run_conditional(
        combined_results, text_dict, filename
    )  # List of dictionaries
    print(f"B Conditional Prompts Complete - {filename}")

    ################## CREATE OUTPUT DIRECTORIES ##################
    base_filename = os.path.splitext(filename)[0].strip()
    output_dir = os.path.join(config.OUTPUT_DIRECTORY, base_filename)
    os.makedirs(output_dir, exist_ok=True)

    ################## WRITE UNPROCESSED OUTPUT ##################
    combined_df = pd.DataFrame(conditional_results)
    combined_df.to_csv(
        os.path.join(output_dir, config.UNPROCESSED_RESULTS_NAME), index=False
    )

    ################## RUN POSTPROCESSING ##################
    final_df = postprocess.b_postprocess(filename, combined_df, num_pages)
    print(f"B Postprocessing Complete - {filename} ")

    ################## WRITE FINAL ##################
    final_df.to_csv(os.path.join(output_dir, config.B_RESULTS_NAME), index=False)
    print(f"B Output Complete - {filename} ")