import xml.etree.ElementTree as ET
import re

INFORMATION_EXTRACTION_PROMPT = """You will be given a text and a list of tags. Your task is to extract all important information from the text related to each tag separately and provide the information with each tag and their corresponding relevant information.

Here is the input text:
<text>
{TEXT}
</text>

Here is the list of tags:
<all_tags>
{TAGS}
</all_tags>

Follow these steps to complete the task:

1. Read through the entire text carefully.

2. For each tag in the list:
   a. Identify all relevant information in the text that relates to the tag.
   b. Extract and summarize the important points related to that tag.
   c. If no relevant information is found for a tag, note that as well.

3. Format your output as follows:
   - Use <tag_info> tags to enclose the information for each tag.
   - Within each <tag_info> section:
     - Start with the tag name in <tag> tags.
     - Follow with the extracted information in <info> tags.
   - If no relevant information is found for a tag, just leave section empty within the <info> tags but make sure the <info> tags are present.

4. Ensure that the extracted information is concise yet comprehensive, capturing all important points related to each tag.

Here's an example of how your output should be formatted:

<answer>
<tag_info>
<tag>Tag1</tag>
<info>Relevant information for Tag1 extracted from the text...</info>
</tag_info>

<tag_info>
<tag>Tag2</tag>
<info>Relevant information for Tag2 extracted from the text...</info>
</tag_info>

<tag_info>
<tag>Tag3</tag>
<info></info>
</tag_info>
</answer>

Remember to include all tags from the provided list, even if no relevant information is found for some of them.

Provide your complete answer within <answer> tags."""

INFORMATION_EXTRACTION_TAG_FORMAT = """<answer>
<tag_info>
<tag>Tag1</tag>
<info>Relevant information for Tag1 extracted from the text...</info>
</tag_info>

<tag_info>
<tag>Tag2</tag>
<info>Relevant information for Tag2 extracted from the text...</info>
</tag_info>

<tag_info>
<tag>Tag3</tag>
<info></info>
</tag_info>
</answer>"""

COMPARISON_PROMPT = """You are tasked with comparing information from two PDFs for a given set of tags. The input will be provided in the following format:

<tag_info>
{TAG_INFO}
</tag_info>

Follow these steps to complete the task:

1. Read through the entire input carefully.

2. For each tag data, you will see the tag and the information related to the tag from pdf1 and pdf2. Your job is to compare this information and explain the differences.

3. When comparing the information:
   - Look for differences in content, wording, or details.
   - Note if one PDF provides information that the other doesn't.
   - Identify any contradictions between the two PDFs.

4. Format your output as follows:
   - Use <comparison> tags to enclose your entire response.
   - For each tag, use <tag_difference> tags to enclose the explanation of differences.
   - Within each <tag_difference> section, use <tag> to enclose the tag name, and <difference> to enclose the explanation of differences.
   - Provide a clear, concise explanation of the differences between pdf1 and pdf2 for that tag.
   - Your output should be in korean language

5. Here's an example of how your output should look:

<comparison>
<tag_difference>
<tag>Tag 1</tag>
<difference>
PDF1 states that the product is available in three colors (red, blue, green), while PDF2 mentions four colors (red, blue, green, yellow).
</difference>
</tag_difference>
<tag_difference>
<tag>Tag 2</tag>
<difference>
The information in PDF1 and PDF2 is identical for this tag, both describing the product's weight as 500g.
</difference>
</tag_difference>
</comparison>

6. Additional guidelines:
   - Be thorough in your comparison, addressing all aspects of the information provided.
   - If there are no differences for a particular tag, state this clearly.
   - Use clear and concise language in your explanations.
   - If one PDF lacks information that the other provides, explicitly mention this.

Begin your analysis with the first tag in the input and proceed through all tags sequentially. Ensure that your output is comprehensive and accurately reflects the differences between the two PDFs for each tag and is in korean language."""

COMPARISON_INPUT_FORMAT = """<tag_data>
<tag>{tag}</tag>
<pdf1>{pdf1_information}</pdf1>
<pdf2>{pdf2_information}</pdf2>
</tag_data>"""

COMPARISON_TAG_FORMAT = """<comparison>
<tag_difference>
<tag>Tag 1</tag>
<difference>
PDF1 states that the product is available in three colors (red, blue, green), while PDF2 mentions four colors (red, blue, green, yellow).
</difference>
</tag_difference>
<tag_difference>
<tag>Tag 2</tag>
<difference>
The information in PDF1 and PDF2 is identical for this tag, both describing the product's weight as 500g.
</difference>
</tag_difference>
</comparison>"""

# def verify_INFORMATION_EXTRACTION_PROMPT(text):
#     try:
#         root = ET.fromstring(text)
        
#         # Check if the root element is 'answer'
#         if root.tag != 'answer':
#             print("root tag is wrong")
#             return False
        
#         # Check if all child elements are 'tag_info'
#         for child in root:
#             if child.tag != 'tag_info':
#                 print("tag_info tag is wrong")
#                 return False
        
#         # Verify each tag_info element
#         for tag_info in root.findall('.//tag_info'):
#             tag = tag_info.find('tag')
#             info = tag_info.find('info')
            
#             # Check if tag exists and is not empty
#             if tag is None:
#                 print("tag is missing")
#                 return False
            
#             if not tag.text.strip():
#                 print("tag is empty")
#                 return False
            
#             # Check if info exists and is optional
#             if info is None:
#                 print("info is missing")
#                 return False
                    
#         # If all checks pass, the format is valid
#         return True
    
#     except ET.ParseError:
#         # If parsing fails, the text doesn't follow the expected format
#         print(text)
#         return False


def verify_INFORMATION_EXTRACTION_PROMPT(text):
    # Regular expression pattern
    pattern = r'<answer>(.*?)<\/answer>'
    
    # Try to match the overall structure
    match = re.search(pattern, text, flags=re.DOTALL)
    
    # Extract the content between <answer> tags
    content = match.group(1)
    
    # Pattern for tag_info blocks
    tag_info_pattern = r'<tag_info>\s*<tag>(.*?)<\/tag>\s*<info>(.*?)<\/info>\s*<\/tag_info>'
    
    # Find all tag_info blocks
    matches = re.findall(tag_info_pattern, content,flags=re.DOTALL)
    # Check each match
    for match in matches:
        tag, info = match
        
        # Check if tag exists and is not empty
        if not tag.strip():
            print("Tag is empty")
            return False
        
    
    # If all checks pass, the format is valid
    return True

# def verify_COMPARISON_PROMPT(text):
#     try:
#         root = ET.fromstring(text)
        
#         # Check if the root element is 'answer'
#         if root.tag != 'comparison':
#             return False
        
#         # Check if all child elements are 'tag_difference'
#         for child in root:
#             if child.tag != 'tag_difference':
#                 return False
        
#         # Verify each tag_difference element
#         for tag_difference in root.findall('.//tag_difference'):
#             tag = tag_difference.find('tag')
#             difference = tag_difference.find('difference')
            
#             # Check if tag exists and is not empty
#             if tag is None:
#                 return False
            
#             if not tag.text.strip():
#                 return False
            
#             # Check if info exists and is optional
#             if difference is None:
#                 return False
                    
#         # If all checks pass, the format is valid
#         return True
    
#     except ET.ParseError:
#         # If parsing fails, the text doesn't follow the expected format
#         return False

def verify_COMPARISON_PROMPT(text):
    # Regular expression pattern
    pattern = r'<comparison>(.*?)<\/comparison>'
    
    # Try to match the overall structure
    match = re.search(pattern, text, flags=re.DOTALL)
    if not match:
        print("no comparison match")
        return False
    
    # Extract the content between <comparison> tags
    content = match.group(1)
    
    # Pattern for tag_difference blocks
    tag_difference_pattern = r'<tag_difference>\s*<tag>(.*?)<\/tag>\s*<difference>(.*?)<\/difference>\s*<\/tag_difference>'
    
    # Find all tag_info blocks
    matches = re.findall(tag_difference_pattern, content,flags=re.DOTALL)
    # Check each match
    for match in matches:
        tag, difference = match
        
        # Check if tag exists and is not empty
        if not tag.strip():
            print("Tag is empty")
            return False
        
    
    # If all checks pass, the format is valid
    return True

def verify_all_tags_present(text, tags):
    tag_pattern = r'<tag>(.*?)<\/tag>'
    tags_llm = re.findall(tag_pattern, text,flags=re.DOTALL)
    tags = [t.strip() for t in tags]
    tags_llm = [t.strip() for t in tags_llm]
    extra_tags = []
    missing_tags = []
    for tag in tags:
        if tag not in tags_llm:
            missing_tags.append(tag)
    for tag in tags_llm:
        if tag not in tags:
            extra_tags.append(tag)
    return missing_tags, extra_tags

def extract_INFORMATION_EXTRACTION_PROMPT(text):
    # Regular expression pattern
    pattern = r'<answer>(.*?)<\/answer>'
    
    # Try to match the overall structure
    match = re.search(pattern, text, flags=re.DOTALL)
    
    # Extract the content between <answer> tags
    content = match.group(1)
    
    # Pattern for tag_info blocks
    tag_info_pattern = r'<tag_info>\s*<tag>(.*?)<\/tag>\s*<info>(.*?)<\/info>\s*<\/tag_info>'
    
    # Find all tag_info blocks
    matches = re.findall(tag_info_pattern, content,flags=re.DOTALL)
    # Check each match
    data = {}
    for match in matches:
        tag, info = match
        
        data.update({tag:info})

    return data

def extract_COMPARISON_PROMPT(text):
    # Regular expression pattern
    pattern = r'<comparison>(.*?)<\/comparison>'
    
    # Try to match the overall structure
    match = re.search(pattern, text, flags=re.DOTALL)
    
    # Extract the content between <comparison> tags
    content = match.group(1)
    
    # Pattern for tag_difference blocks
    tag_difference_pattern = r'<tag_difference>\s*<tag>(.*?)<\/tag>\s*<difference>(.*?)<\/difference>\s*<\/tag_difference>'
    
    # Find all tag_info blocks
    matches = re.findall(tag_difference_pattern, content,flags=re.DOTALL)
    # Check each match
    data = {}
    for match in matches:
        tag, difference = match
        data.update({tag:difference})

    return data