import re
from collections import defaultdict
from itertools import chain, count, groupby

import claude_funcs
import config
import keywords
import prompts
import smart_chunking_funcs
import string_funcs
import valid
from valid import DERIVED_INDICATOR_FIELDS
from regex_patterns import PIPE_PATTERN


def clean_newlines(contract_text: str) -> str:
    """Clean textract output by removing newlines and page number indicators.

    This function processes input text to:
    1. Remove all newline characters.
    2. 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
    """
    cleaned_text = re.sub(r"(?<!\n)\n(?!\n)", " ", contract_text)
    cleaned_text = re.sub(r"Page [0-9]+ of [0-9]+\n\n", " ", contract_text)
    return cleaned_text


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)):
        text_dict[temp_list[i].split()[0]] = text_list[i]

    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)
    contract_text = re.sub(r"(C\.?F\.?R\.?) \$", r"\1§", contract_text)

    # 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_funcs.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_funcs.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 smart_chunk_ac(
    text_dict: dict, 
    contract_text: str,
    keyword_mappings: dict = keywords.GROUPED_KEYWORD_MAPPINGS,
) -> dict:
    """Performs a keyword search on textract outputs to retrieve relevant chunks for a given field.

    Args:
        text_dict (dict): Dictionary keyed by string page-number and valued by actual textract output page
        contract_text(str): Contract Text

    Raises:
        ValueError: When the `key_dict` from keyword_mappings has a methodology that is not implemented.
                    Current options are `or` and `hierarchy`.

    Returns:
        dict: A dictionary keyed by field (defined by the keys in keyword_mappings) valued by
              corresponding chunks of text (input strings delimited by newlines).
    """
    field_pages = defaultdict(set)  # TODO: remove this, it's unused
    chunks = defaultdict(str)

    full_context = "\n".join(
        [f"Page {page}:\n{content}" for page, content in text_dict.items()]
    )

    for field_group, key_dict in keyword_mappings.items():
        if key_dict["methodology"] == "or":
            page_list = smart_chunking_funcs.chunk_or(
                text_dict, key_dict["keywords"], key_dict["case_sensitive"]
            )
            keyword_page_dict = smart_chunking_funcs.keyword_search(
                text_dict, key_dict["keywords"], key_dict["case_sensitive"]
            )
        elif key_dict["methodology"] == "hierarchy":
            page_list = smart_chunking_funcs.chunk_hierarchical(
                text_dict, key_dict["keywords"], key_dict["case_sensitive"]
            )
            keyword_page_dict = smart_chunking_funcs.keyword_search(
                text_dict, key_dict["keywords"], key_dict["case_sensitive"]
            )
        elif key_dict["methodology"] == "and":
            page_list = smart_chunking_funcs.chunk_and(
                text_dict, key_dict["keywords"], key_dict["case_sensitive"]
            )
            keyword_page_dict = smart_chunking_funcs.keyword_search(
                text_dict, key_dict["keywords"], key_dict["case_sensitive"]
            )
        elif key_dict["methodology"] == "regex":
            page_list = smart_chunking_funcs.chunk_regex(text_dict, key_dict["regex"])
            keyword_page_dict = smart_chunking_funcs.keyword_search(
                text_dict, key_dict["keywords"], key_dict["case_sensitive"], True
            )
        elif key_dict["methodology"] == "include_exclude":
            # Chunk pages based on included and excluded keywords
            page_list = smart_chunking_funcs.chunk_with_include_exclude_keywords(
                text_dict, key_dict["included_keywords"], key_dict["excluded_keywords"]
            )
            
            # Create a copy of keywords list by combining included and excluded keywords - this is required for keyword search
            included_keywords = key_dict["included_keywords"]
            excluded_keywords = key_dict["excluded_keywords"]
            
            keywords = list(chain(included_keywords,excluded_keywords))

            # Perform keyword search on the text dictionary - this returns a dictionary keyed by keyword and values are list of pages containing that keyword
            # This is not currently used but is available for debugging purposes (or future use cases)
            keyword_page_dict = smart_chunking_funcs.keyword_search(
                text_dict, keywords, key_dict["case_sensitive"]
            )
        else:
            page_list = list(text_dict.keys())
            keyword_page_dict = {"All pages": list(text_dict.keys())}
            raise ValueError(
                "methodology must be 'hierarchy' or 'or' - using entire contract as chunk"
            )
        if field_group == "initial_page_group":
            keyword_page_dict["initial_pages"] = [1, 2]
        
        # chunks[field_group] = concatenated page texts
        if len(page_list) == len(text_dict.keys()):
            chunks[field_group] = contract_text
        else:
            chunks[field_group] = "\n".join(
                [text_dict[str(page_num)] for page_num in page_list]
            )

        chunks[field_group + "_pages"] = keyword_page_dict
        chunks[field_group + "_methodology"] = key_dict["methodology"]
        chunks[field_group + "_case"] = key_dict["case_sensitive"]
    return chunks


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, filename):
    exhibit_pages = []
    for page_num, page in text_dict.items():
        prompt = prompts.EXHIBIT_CHECK(page[0:100])
        claude_answer_raw = claude_funcs.invoke_claude(
            prompt, config.MODEL_ID_CLAUDE3_HAIKU, filename, max_tokens=10
        )
        claude_answer_extracted = re.findall(pattern=PIPE_PATTERN, string=claude_answer_raw)[-1]
        if "Y" in claude_answer_extracted:
            exhibit_pages.append(page_num)
    return exhibit_pages