import re

from src.prompts import preprocessing_prompts
import src.utils.llm_utils as llm_utils
import src.utils.string_utils as string_utils
from src import config, keywords
from src.regex.regex_patterns import PIPE_PATTERN


from src import config, keywords
from src.enums.delimiters import Delimiter

def remove_page_indicators(contract_text: str) -> str:
    """Clean textract output by removing page number indicators in the form of "Page X of Y"

    This function processes input text to remove lines that indicate page numbers (e.g. 'Page 1 of 10')

    Args:
        contract_text (str): Raw text output from Textract to be cleaned

    Returns:
        str: cleaned text with newlines and page number indicators removed
    """
    
    if contract_text:
        cleaned_text = re.sub(r"Page [0-9]+ of [0-9]+\n\n", " ", contract_text)
    else:
        cleaned_text = contract_text
    return cleaned_text


# TODO: write unit tests
def split_text(text: str) -> dict[str, str]:
    """Split text on pages by the string `Start of Page No. = '

    Args:
        text (str): Raw text output from Textract to be split

    Returns:
        dict[str, str]: A dictionary, keyed by the string page number and valued by the page text.
    """
    temp_list = text.split("Start of Page No. = ")
    text_list = re.split(r"Start of Page No. = [0-9]+\n", text)

    text_dict = {}
    for i in range(len(text_list)):
        if temp_list[i]:
            text_dict[temp_list[i].split()[0]] = text_list[i]   # splits on whitespace characters, which includes spaces, tabs, and newline characters

    return {k: v for k, v in text_dict.items() if k != "Document"}


def clean_law_symbols(contract_text):
    contract_text = contract_text.replace("$$", "$")

    contract_text = re.sub(r"(U\.?S\.?C\.?) \$", r"\1§", contract_text) # replaces $ with § when it follows abbreviations like "U.S.C." (United States Code).
    contract_text = re.sub(r"(C\.?F\.?R\.?) \$", r"\1§", contract_text) # replaces $ with § when it follows abbreviations like "C.F.R." (Code of Federal Regulations).

    # Second correction: Replace '$' with '§' when followed by a number with three decimal places
    contract_text = re.sub(r"\$(?=\d+\.\d{3})", "§", contract_text)

    return contract_text


# ORIGINAL
def chunk_consecutive_og(text_dict, exhibit_pages):
    # If needed - extend page 1 to page 1 and 2, then cut chunking off after (edge case: compensation terms not found on first page of exhibit)

    reimbursement_pages = [page_num for page_num in text_dict.keys() if string_utils.contains_reimbursement(text_dict, page_num)]

    page_dict = {}
    current_exhibit = None
    for page_num in text_dict.keys():

        # Page is Reimbursement AND Exhibit
        if page_num in reimbursement_pages and page_num in exhibit_pages:
            current_exhibit = page_num
            page_dict[page_num] = [page_num]

        # Page is Reimbursement NOT Exhibit
        elif page_num in reimbursement_pages and page_num not in exhibit_pages:
            if current_exhibit:
                page_dict[current_exhibit].append(page_num)
            else:
                page_dict[page_num] = [page_num]

        # Page is Exhibit NOT Reimbursement
        elif page_num in exhibit_pages and page_num not in reimbursement_pages:
            current_exhibit = page_num
            page_dict[page_num] = [page_num]

        # Page is NOT Exhibit NOT Reimbursment
        elif page_num not in exhibit_pages and page_num not in reimbursement_pages:
            current_exhibit = None
            page_dict[page_num] = [page_num]

    final_dict = {page_num : '' for page_num in page_dict.keys()}
    for page_num in page_dict.keys():
        for p in page_dict[page_num]:
            final_dict[page_num] += text_dict[p]
    return final_dict


def chunk_consecutive(text_dict, exhibit_pages):
    exhibit_pages = set(str(page) for page in exhibit_pages)
    reimbursement_pages = {
        page_num
        for page_num in text_dict.keys()
        if string_utils.contains_reimbursement(text_dict, page_num)
    }

    page_dict = {}
    current_chunk_start = None
    last_exhibit = None
    in_exhibit = False

    def word_count(text):
        return len(text.split())

    for page_str in sorted(text_dict.keys(), key=int):
        page_num = int(page_str)

        if page_str in exhibit_pages:
            current_chunk_start = page_str
            last_exhibit = page_str
            in_exhibit = True
            page_dict[current_chunk_start] = [page_str]
            # print(f"found page {page_num} in exhibit, starting new chunk")

        if page_str in reimbursement_pages:
            if not in_exhibit or current_chunk_start is None:
                # Check if this page has less than 200 words and should be added to the previous chunk
                if current_chunk_start and word_count(text_dict[page_str]) < 200:
                    page_dict[current_chunk_start].append(page_str)
                    # print(f"found page {page_num} in reimbursement with less than 200 words, adding to previous chunk")
                else:
                    current_chunk_start = page_str
                    page_dict[current_chunk_start] = [page_str]
                    # print(f"found page {page_num} in reimbursement, starting new chunk")
            else:
                page_dict[current_chunk_start].append(page_str)
                # print(f"found page {page_num} in reimbursement, adding to current chunk")
        else:
            if in_exhibit and current_chunk_start is not None:
                page_dict[current_chunk_start].append(page_str)
                # print(f"found non-reimbursement page {page_num} in exhibit, adding to current chunk")
            else:
                in_exhibit = False

        # Check next page if it's non-reimbursement and not in exhibit
        next_page_str = str(page_num + 1)
        if (
            next_page_str in text_dict
            and next_page_str not in reimbursement_pages
            and next_page_str not in exhibit_pages
        ):
            if current_chunk_start is not None:
                page_dict[current_chunk_start].append(next_page_str)
                # print(f'next page {next_page_str} was found to be a non-reimbursement and added')

        # If we've moved past the last exhibit page, reset in_exhibit
        if in_exhibit and int(page_str) > int(last_exhibit):
            in_exhibit = False

    # Remove duplicates and sort page numbers in each chunk
    for key in page_dict:
        page_dict[key] = sorted(list(set(page_dict[key])), key=int)

    final_dict = {page_num: "" for page_num in page_dict.keys()}
    for page_num in page_dict.keys():
        for p in page_dict[page_num]:
            final_dict[page_num] += text_dict[p]

    # print("Final page_dict:", {k: v for k, v in page_dict.items()})
    return final_dict


def filter_quick_review(text_dict):
    """
    Cover Sheets (aka Top Sheets or Quick Review pages) are pages stapled to the front of the contract that contain manually written summaries of the contract's contents.
    They are NOT legally binding documents, and because they are often manually filled out and hand-written, are more prone to have erroneous information than the rest of the contract.
    For Doczy.AI, we should not pull any information from top sheets, except as a very last resort.

    filter_quick_review() splits the input dictionary into two dictionaries, one containing the contract pages, the other containing the top sheet pages.

    The function checks each page's text for the keywords "QUICK REVIEW", "TOP SHEET", and "COVER SHEET". It categorizes the text into two separate dictionaries: one for texts that do not contain any of these keywords, and another for texts that do.

    Parameters:
    text_dict (dict): A dictionary where the key is the page number and the value is the text of that page.

    Returns:
    tuple of two dicts:
        - The first dictionary contains pages that do not have the specified keywords.
        - The second dictionary includes pages that contain any of the specified keywords.
    """
    return {
        page_num: page_text
        for page_num, page_text in text_dict.items()
        if "QUICK REVIEW" not in page_text.upper()
        and "TOP SHEET" not in page_text.upper()
        and "COVER SHEET" not in page_text.upper()
    }, {
        page_num: page_text
        for page_num, page_text in text_dict.items()
        if "QUICK REVIEW" in page_text.upper()
        or "TOP SHEET" in page_text.upper()
        or "COVER SHEET" in page_text.upper()
    }

def get_exhibit_pages(text_dict: dict[str, str], filename: str) -> list[str]:
    """Extract beginning-of-exhibit pages from a contract.  The first page is also 
    always included in the exhibit pages.

    Args:
        text_dict (dict[str, str]): Dictionary valued by string-formatted page numbers and valued by page text
        filename (str): Filename of the contract (for tracking purposes)

    Returns:
        list[str]: A list of exhibit page numbers extracted from the contract.
    """
    exhibit_pages = []
    first_page = True
    for page_num, page in text_dict.items():
        if first_page: # Include the first page in the exhibit pages
            exhibit_pages.append(page_num)
            first_page = False
        else:
            prompt = preprocessing_prompts.EXHIBIT_CHECK(page[0:100])
            claude_answer_raw = llm_utils.invoke_claude(
                prompt, config.MODEL_ID_CLAUDE3_HAIKU, filename, max_tokens=10 # TODO: low priority, try increasing max_tokens and maybe pass multiple pages in to reduce overall calls
            )
            claude_answer_extracted = string_utils.extract_text_from_delimiters(
                                claude_answer_raw, Delimiter.PIPE
                            )
            if "Y" in claude_answer_extracted:
                exhibit_pages.append(page_num)
    return exhibit_pages

def chunk_by_exhibit(text_dict: dict, 
                     exhibit_pages: list
                     ) -> dict:
    """
    Organizes pages into groups based on their association with specific exhibits.

    This function assigns each page number from the `text_dict` dictionary to an exhibit
    based on the `exhibit_pages` list. Pages are grouped under the nearest preceding 
    page number in `exhibit_pages`. If a page number in `text_dict` is itself in 
    `exhibit_pages`, it starts a new exhibit group.

    Parameters:
        text_dict (dict): A dictionary where keys are page numbers and values are page text
        exhibit_pages (list): A list of page numbers that mark the beginning of a new
                              exhibit.

    Returns:
        dict: A dictionary mapping each page number in `text_dict` to its corresponding
              exhibit identifier. The exhibit identifier is the page number of the first
              page in that exhibit as listed in `exhibit_pages`. If there are pages before
              the first `exhibit_page`, they are grouped under the exhibit identifier "0".
    """
    if len(exhibit_pages) == 0:
        return {key : key for key in text_dict.keys()}
    
    exhibit_dict = {}
    current_exhibit = "0"
    for page_num in text_dict.keys():
        if page_num in exhibit_pages:
            current_exhibit = page_num
            exhibit_dict[page_num] = current_exhibit
        else:
            exhibit_dict[page_num] = current_exhibit
    return exhibit_dict