import re
import logging

from src.prompts import investment_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):
    if not contract_text:
        return contract_text
    # First correction: Replace '$$' with '$'
    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 using LLM-based detection.  
    
    This function identifies pages that start new exhibits by checking each page (or
    subpage) for exhibit headers using an LLM prompt.  It handles both regular pages
    (e.g., "5") and subpages created by table splitting (e.g., "5.1", "5.2").

    Key behaviors:
    - The first page of the contract is always included in the exhibit pages.
    - For base pages with subpages, if ANY subpage contains an exhibit header,
      the FIRST subpage is added to the exhibit list
    - Each base page number is only processed once to avoid duplicates.

    Args:
        text_dict (dict[str, str]): Dictionary keyed by string-formatted page numbers and valued by page text
                                    Page numbers may include subpages (e.g., "5.1", "5.2")
        filename (str): Filename of the contract (for tracking purposes)

    Returns:
        list[str]: Ordered list of page numbers that start new exhibits.  May include
                   subpage numbers (e.g., ["1", "5.1", "5.2"]) depending on table splitting.
        
    Example:
        Given pages ["1", "2", "5.0", "5.1", "5.2", "7"] where "5.1" contains an exhibit header:
        Returns ["1", "5.1", "7"] (assuming pages 1 and 7 also have exhibit headers)
    """
    exhibit_pages = []
    first_page = True

    # Sort pages to handle subpages properly (5.0, 5.1, 5.2, etc.)
    # TODO: augment `string_utils.page_key_sort` to handle subpages
    sorted_pages = sorted(text_dict.keys(), key=lambda x: (
        int(x.split('.')[0]),
        float(x.split('.')[1]) if '.' in x else 0)
    )

    for page_num in sorted_pages:
        
        # Always include the first page as an exhibit
        if first_page:
            exhibit_pages.append(page_num)
            is_exhibit = True
            first_page = False
        else:
            if "." not in page_num or ".0" in page_num:
                page_content = text_dict[page_num]
                prompt = investment_prompts.EXHIBIT_CHECK(page_content[0:200])
                claude_answer_raw = llm_utils.invoke_claude(
                    prompt, config.MODEL_ID_CLAUDE3_HAIKU, filename, max_tokens=150
                )
                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)
                    is_exhibit = True
                else:
                    is_exhibit = False
            elif is_exhibit==True:
                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_chunk_mapping = {}
    current_exhibit = "0"
    for page_num in text_dict.keys():
        if page_num in exhibit_pages:
            current_exhibit = page_num
            exhibit_chunk_mapping[page_num] = current_exhibit
        else:
            exhibit_chunk_mapping[page_num] = current_exhibit
    return exhibit_chunk_mapping


def get_exhibit_dict(text_dict: dict,
                     exhibit_chunk_mapping: dict
                     ) -> dict:
    """
    Creates a dictionary of exhibit texts, concatenating pages with the same exhibit number,
    but separating decimal pages into their own chunks.
    
    Args:
        text_dict (dict): Dictionary mapping page numbers to page text.
        exhibit_chunk_mapping (dict): Dictionary mapping page numbers to exhibit numbers.

    Returns:
        dict: Dictionary mapping exhibit chunk keys to combined text. Format for keys: 
        "{exhibit}.{index}"
            NOTE: If exhibit identifiers contain decimals (e.g., "23.0") output keys will
            follow the same pattern (resulting in "23.0.0", "23.0.1", etc.)

    Guidelines: 
        - Every exhibit chunk contains at most 1 table page in it (i.e. one page with a decimal)
        - Tables are inserted in chronological order within the document flow.
        - Multiple decimal pages are separated into individual chunks
        - Each chunk gets a new key: "{exhibit}.{index}" (e.g. "2.0", "2.1", "2.2", etc.)
    
    Algorithm:
        For each exhibit:
        1. Identify base pages (no decimals) and table pages (with decimals).
        2. For each decimal page, create a chunk containing:
            - All base pages that came before the decimal page (in document order)
            - The specific decimal page
            - All base pages that came after the decimal page (in document order)

    Example 1:
        Input mapping {'1.0': '1.0', '2': '2', '3.0': '2', '3.1': '2', '3.2': '2', '4.0': '2'}
        - '1.0': text from '1.0' (single page exhibit)
        - '2.0': text from '2' + text from '3.0' (base page 2, then table 3.0)
        - '2.1': text from '2' + text from '3.1' (base page 2, then table 3.1)
        - '2.2': text from '2' + text from '3.2' (base page 2, then table 3.2)
        - '2.3': text from '2' + text from '4.0' (base page 2, then table 4.0)

    Example 2:
        Input mapping {'1.0': '1.0', '2': '2', '3.0': '2', '3.1': '2', '3.2': '2', '4': '2', '5' : '2'}
        - '1.0': text from '1.0'
        - '2.0': text from '2' + text from '3.0' + text from '4' + text from '5'
        - '2.1': text from '2' + text from '3.1' + text from '4' + text from '5'
        - '2.2': text from '2' + text from '3.2' + text from '4' + text from '5'
        - '2.3': text from '2' + text from '4' + text from '5'
        
        Note: Tables 3.0, 3.1, and 3.2 appear between pages 2 and 4 in document flow.
    
    Other cases:
        - No decimal pages: creates single chunk for the exhibit.
        - Single decimal page exhibit: uses original page as the chunk key.
    """
    exhibit_dict = {}
    
    # Group pages by exhibit
    exhibit_pages = {}
    for page, exhibit in exhibit_chunk_mapping.items():
        if exhibit not in exhibit_pages:
            exhibit_pages[exhibit] = []
        exhibit_pages[exhibit].append(page)
    
    # Process each exhibit
    for exhibit in exhibit_pages:
        pages = sorted(exhibit_pages[exhibit], key=string_utils.page_key_sort)

        base_pages = [p for p in pages if '.' not in p]
        decimal_pages = [p for p in pages if '.' in p]

        # Special case: single decimal page exhibit
        if len(decimal_pages) == 1 and not base_pages:
            exhibit_dict[decimal_pages[0]] = text_dict[decimal_pages[0]]
            continue

        # Case: No decimal pages - create single chunk with .0 suffix
        if not decimal_pages:
            exhibit_dict[f"{exhibit}.0"] = "\n".join(text_dict[p] for p in pages)
            continue
            
        # Case: With decimal pages, create chunks for each decimal page
        chunk_index = 0
        for decimal_page in decimal_pages:
            chunk_key = f"{exhibit}.{chunk_index}"

            # Get decimal page base number for chronological positioning
            decimal_base = int(decimal_page.split('.')[0])

            # Build chunk with pages in chronological order
            chunk_parts = []
            
            # Add base pages before decimal page
            for base_page in base_pages:
                if int(base_page) <= decimal_base:
                    chunk_parts.append(text_dict[base_page])
            
            # Add the decimal page itself
            chunk_parts.append(text_dict[decimal_page])

            # Add base pages after decimal page
            for base_page in base_pages:
                if int(base_page) > decimal_base:
                    chunk_parts.append(text_dict[base_page])

            exhibit_dict[chunk_key] = "\n".join(chunk_parts)
            chunk_index += 1

        # Add final chunk with just base pages (if there are base pages after all decimal pages)
        if base_pages and decimal_pages:
            # Check if there are base pages that come after all decimal pages
            max_decimal_base = max(int(dp.split('.')[0]) for dp in decimal_pages)
            remaining_base_pages = [bp for bp in base_pages if int(bp) > max_decimal_base]

            if remaining_base_pages:
                chunk_key = f"{exhibit}.{chunk_index}"
                # Include ALL base pages in the final chunk
                exhibit_dict[chunk_key] = "\n".join(text_dict[bp] for bp in base_pages)
    return exhibit_dict

def get_exhibit_headers(exhibit_dict, filename):
    """
    Extracts the exhibit header using an LLM prompt.

    Args:
        exhibit_chunk (str): A chunk of exhibit text to analyze.
        filename (str): The name of the file being processed.

    Returns:
        str: The extracted exhibit header.
    """
    exhibit_header_dict = {}
    base_pages = {}
    for page_num, exhibit_chunk in exhibit_dict.items():
        base_page_num = page_num.split('.')[0]
        if base_page_num in base_pages.keys():
            exhibit_header_dict[page_num] = base_pages[base_page_num]
        else:
            prompt = investment_prompts.EXHIBIT_HEADER(exhibit_chunk[0:400])
            llm_answer_raw = llm_utils.invoke_claude(
                prompt, config.MODEL_ID_CLAUDE35_SONNET, filename
            )
            llm_answer_final = string_utils.extract_text_from_delimiters(
                llm_answer_raw, Delimiter.PIPE
            )
            exhibit_header_dict[page_num] = llm_answer_final
            base_pages[base_page_num] = llm_answer_final
    
    return exhibit_header_dict