import gradio as gr
import re
import pandas as pd

def build_keywords_dict(primary_inputs, synonym_inputs):
    """Build keyword dictionary from separate primary and synonym inputs"""
    keywords_dict = {}
    
    for primary, synonyms in zip(primary_inputs, synonym_inputs):
        if primary and primary.strip():  # Only process if primary keyword exists
            primary_clean = primary.strip()
            if synonyms and synonyms.strip():
                synonym_list = [s.strip() for s in synonyms.split(';') if s.strip()]
            else:
                synonym_list = []
            keywords_dict[primary_clean] = synonym_list
    
    return keywords_dict

def find_keywords(story, keywords_dict):
    """Find keywords in the story text"""
    if not story or not isinstance(story, str):
        return ''
    
    found_keywords = set()
    
    # Search for each primary keyword and its synonyms
    for primary_keyword, synonyms in keywords_dict.items():
        keyword_group_found = False
        
        # Check primary keyword
        if primary_keyword.upper() == "US":
            if ' US ' in story or story.startswith('US ') or story.endswith(' US'):
                keyword_group_found = True
        else:
            pattern = r'\b' + re.escape(primary_keyword) + r'\b'
            if re.search(pattern, story, re.IGNORECASE):
                keyword_group_found = True
        
        # Check each synonym
        for synonym in synonyms:
            if synonym.upper() == "US":
                if ' US ' in story or story.startswith('US ') or story.endswith(' US'):
                    keyword_group_found = True
            else:
                if re.search(r'\b' + re.escape(synonym) + r'\b', story, re.IGNORECASE):
                    keyword_group_found = True
        
        # If any keyword from this group was found, add ALL keywords from the group
        if keyword_group_found:
            found_keywords.add(primary_keyword)  # Always include the primary
            found_keywords.update(synonyms)     # Add all synonyms
    
    return '; '.join(sorted(found_keywords))

def highlight_keywords_in_text(text, keywords_list):
    """Create HTML with highlighted keywords while preserving line breaks"""
    if not keywords_list:
        # Convert line breaks to HTML breaks for plain text
        formatted_text = text.replace('\n', '<br>')
        return formatted_text
    
    highlighted_text = text
    colors = ['#FF6B6B', '#4ECDC4', '#45B7D1', '#F9CA24', '#6C5CE7', '#A0E7E5', '#FD79A8', '#55A3FF', '#00B894', '#E17055']
    
    for i, keyword in enumerate(keywords_list):
        if keyword:
            color = colors[i % len(colors)]
            pattern = r'\b' + re.escape(keyword) + r'\b'
            replacement = f'<span style="background-color: {color}; padding: 2px 4px; border-radius: 3px; color: white; font-weight: bold;">{keyword}</span>'
            highlighted_text = re.sub(pattern, replacement, highlighted_text, flags=re.IGNORECASE)
    
    # Convert line breaks to HTML breaks after highlighting
    highlighted_text = highlighted_text.replace('\n', '<br>')
    return highlighted_text

def create_keyword_results_table(found_keywords_str, keywords_dict, input_text):
    """Create HTML table showing detailed keyword results"""
    if not found_keywords_str:
        return "<p style='text-align: center; padding: 20px;'>No keywords found.</p>"
    
    found_keywords = found_keywords_str.split('; ')
    
    # Group keywords by their primary category
    keyword_groups = {}
    for primary, synonyms in keywords_dict.items():
        found_in_group = []
        # Check if primary keyword was found
        if primary in found_keywords:
            found_in_group.append(primary)
        # Check if any synonyms were found
        for synonym in synonyms:
            if synonym in found_keywords:
                found_in_group.append(synonym)
        
        if found_in_group:
            keyword_groups[primary] = found_in_group
    
    if not keyword_groups:
        return "<p style='text-align: center; padding: 20px;'>No keyword groups matched.</p>"
    
    # Create the HTML table
    table_html = """
    <div style='max-height: 500px; overflow-y: auto; border: 2px solid #ddd; border-radius: 8px; padding: 20px; background-color: #fafafa; margin: 10px 0;'>
        <h4 style='margin: 0 0 15px 0; color: #333;'>📊 Detailed Keyword Results</h4>
        <table style="width: 100%; border-collapse: collapse; border: 1px solid #ddd; background-color: white;">
            <thead>
                <tr style="background-color: #6366f1; color: white;">
                    <th style="padding: 12px; text-align: left; border: 1px solid #ddd;">Primary Keyword</th>
                    <th style="padding: 12px; text-align: left; border: 1px solid #ddd;">Found Terms</th>
                    <th style="padding: 12px; text-align: left; border: 1px solid #ddd;">Count in Text</th>
                    <th style="padding: 12px; text-align: left; border: 1px solid #ddd;">Context Preview</th>
                </tr>
            </thead>
            <tbody>
    """
    
    colors = ['#FF6B6B', '#4ECDC4', '#45B7D1', '#F9CA24', '#6C5CE7', '#A0E7E5', '#FD79A8', '#55A3FF', '#00B894', '#E17055']
    
    for i, (primary, found_terms) in enumerate(keyword_groups.items()):
        color = colors[i % len(colors)]
        
        # Count total occurrences and get context
        total_count = 0
        contexts = []
        
        for term in found_terms:
            # Count occurrences (case insensitive)
            if term.upper() == "US":
                # Special handling for "US"
                count = len([m for m in re.finditer(r'\bUS\b', input_text)])
            else:
                pattern = r'\b' + re.escape(term) + r'\b'
                count = len(list(re.finditer(pattern, input_text, re.IGNORECASE)))
            
            total_count += count
            
            # Get context (first occurrence)
            if term.upper() == "US":
                match = re.search(r'\bUS\b', input_text)
            else:
                match = re.search(r'\b' + re.escape(term) + r'\b', input_text, re.IGNORECASE)
            
            if match:
                start = max(0, match.start() - 30)
                end = min(len(input_text), match.end() + 30)
                context = input_text[start:end].replace('\n', ' ')
                # Highlight the found term in context
                if term.upper() == "US":
                    highlighted_context = re.sub(
                        r'\bUS\b', 
                        f'<strong style="background-color: {color}; color: white; padding: 1px 2px; border-radius: 2px;">{term}</strong>', 
                        context
                    )
                else:
                    highlighted_context = re.sub(
                        r'\b' + re.escape(term) + r'\b', 
                        f'<strong style="background-color: {color}; color: white; padding: 1px 2px; border-radius: 2px;">{term}</strong>', 
                        context, 
                        flags=re.IGNORECASE
                    )
                contexts.append(highlighted_context)
        
        # Create found terms display
        found_terms_display = []
        for term in found_terms:
            found_terms_display.append(f'<span style="background-color: {color}; color: white; padding: 2px 6px; border-radius: 10px; font-size: 12px; margin: 1px;">{term}</span>')
        
        table_html += f"""
                <tr style="background-color: #fff;">
                    <td style="padding: 10px; border: 1px solid #ddd; font-weight: bold;">{primary}</td>
                    <td style="padding: 10px; border: 1px solid #ddd;">{' '.join(found_terms_display)}</td>
                    <td style="padding: 10px; border: 1px solid #ddd; text-align: center;">
                        <span style='background-color: #28a745; color: white; padding: 4px 8px; border-radius: 12px; font-weight: bold;'>
                            {total_count}
                        </span>
                    </td>
                    <td style="padding: 10px; border: 1px solid #ddd; font-style: italic; font-size: 14px;">
                        {contexts[0] if contexts else 'No context available'}...
                    </td>
                </tr>
        """
    
    table_html += """
            </tbody>
        </table>
    </div>
    """
    
    return table_html

def process_text(input_text, primary1, synonyms1, primary2, synonyms2, primary3, synonyms3, primary4, synonyms4, primary5, synonyms5):
    """Main processing function with added results table"""
    if not input_text.strip():
        return "Please enter some text to analyse", "", "", "No keywords found"
    
    # Build keywords dictionary from separate inputs
    primary_inputs = [primary1, primary2, primary3, primary4, primary5]
    synonym_inputs = [synonyms1, synonyms2, synonyms3, synonyms4, synonyms5]
    keywords_dict = build_keywords_dict(primary_inputs, synonym_inputs)
    
    if not keywords_dict:
        return "Please enter at least one primary keyword", "", "", "No keyword dictionary provided"
    
    # Find keywords in the text
    found_keywords_str = find_keywords(input_text, keywords_dict)
    
    if not found_keywords_str:
        return f"No keywords found in the text.\n\nKeyword dictionary loaded: {len(keywords_dict)} primary keywords", input_text, "", "No matches found"
    
    # Create highlighted version
    keywords_list = found_keywords_str.split('; ')
    highlighted_html = highlight_keywords_in_text(input_text, keywords_list)
    
    # Create results table
    results_table_html = create_keyword_results_table(found_keywords_str, keywords_dict, input_text)
    
    # Create results summary
    results_summary = f"""
## Results Summary

**Keywords Found:** {len(keywords_list)}

**Matched Keywords:** {found_keywords_str}

**Keyword Dictionary Stats:**
- Primary keywords loaded: {len(keywords_dict)}
- Total searchable terms: {sum(len(synonyms) + 1 for synonyms in keywords_dict.values())}

**Copy this result to your spreadsheet:**
{found_keywords_str}
"""
    
    return results_summary, highlighted_html, results_table_html, found_keywords_str

# Create the Gradio interface
def create_interface():
    with gr.Blocks(title="Keyword Tagging Tool", theme=gr.themes.Soft()) as demo:
        gr.HTML("""
        <h1>Controlled Vocabluary Keyword Tagging Tool</h1>
        
        <p>This tool demonstrates how a simple python script can be used to extract keywords from text using a controlled vocabulary of primary keywords and associated keywords/synonyms. 
       </p>
        
        <h2>How to use this tool:</h2>
        <ol>
        <li>📝 <strong>Enter your text</strong> in the left panel</li>
        <li>📚 <strong>Define your keyword dictionary</strong> in the right panel - enter primary keywords and their synonyms</li>
        <li>🔍 <strong>Click "Find Keywords"</strong> to see results</li>
        <li>📋 <strong>Copy the results</strong> to paste into your spreadsheet</li>
        </ol>
        """)
        
        with gr.Row():
            with gr.Column(scale=1):
                text_input = gr.Textbox(
                    label="Text to Analyse",
                    placeholder="Enter the text you want to tag with keywords...",
                    lines=22,
                    max_lines=25
                )
            
            with gr.Column(scale=1):
                gr.Markdown("**Keyword Dictionary** - Enter primary keywords and their synonyms:")
                
                # Row 1
                with gr.Row():
                    primary1 = gr.Textbox(label="Primary Keyword 1", placeholder="e.g., Prisoner of War", scale=1)
                    synonyms1 = gr.Textbox(label="Synonyms 1", placeholder="e.g., POW; POWs; prisoner of war", scale=2)
                
                # Row 2
                with gr.Row():
                    primary2 = gr.Textbox(label="Primary Keyword 2", placeholder="e.g., United States", scale=1)
                    synonyms2 = gr.Textbox(label="Synonyms 2", placeholder="e.g., USA; US; America", scale=2)
                
                # Row 3
                with gr.Row():
                    primary3 = gr.Textbox(label="Primary Keyword 3", placeholder="e.g., University", scale=1)
                    synonyms3 = gr.Textbox(label="Synonyms 3", placeholder="e.g., university; institution; college", scale=2)
                
                # Row 4
                with gr.Row():
                    primary4 = gr.Textbox(label="Primary Keyword 4", placeholder="Optional", scale=1)
                    synonyms4 = gr.Textbox(label="Synonyms 4", placeholder="Optional", scale=2)
                
                # Row 5
                with gr.Row():
                    primary5 = gr.Textbox(label="Primary Keyword 5", placeholder="Optional", scale=1)
                    synonyms5 = gr.Textbox(label="Synonyms 5", placeholder="Optional", scale=2)
        
        # Full width Find Keywords button
        with gr.Row():
            find_btn = gr.Button("Find Keywords", variant="primary", size="lg")
        
        # Clear buttons side by side under Find Keywords
        with gr.Row():
            clear_text_btn = gr.Button("Clear Text", size="lg", variant="secondary")
            clear_dict_btn = gr.Button("Clear Dictionary", size="lg", variant="secondary")
        
        # Horizontal line before results
        gr.HTML("<hr style='margin-top: 20px; margin-bottom: 15px;'>")
        
        with gr.Row():
            results_output = gr.Markdown(label="Results Summary")
        
        with gr.Row():
            highlighted_output = gr.HTML(label="Text with Highlighted Keywords")
        
        with gr.Row():
            results_table_output = gr.HTML(label="Detailed Results Table")
        
        with gr.Row():
            copy_output = gr.Textbox(
                label="Keywords for Spreadsheet (copy this text)",
                lines=3,
                max_lines=5
            )
        
        # Examples section with table format
        gr.Markdown("### Examples")
        
        example1 = [
            "During World War II, many prisoners of war were held in camps across Europe. The Geneva Convention established rules for POW treatment. American soldiers and British troops were among those captured.",
            "Prisoner of War", "POW; POWs; prisoner of war",
            "World War II", "WWII; Second World War",
            "United States", "USA; US; America; American",
            "", "", "", ""
        ]
        
        example2 = [
            "The University of Oxford is located in Oxford, England. Students from around the world study at this prestigious institution.",
            "University", "university; institution; college",
            "Oxford", "oxford",
            "England", "england; English",
            "Student", "student; students; pupils",
            "", ""
        ]
        
        gr.Examples(
            examples=[example1, example2],
            inputs=[text_input, primary1, synonyms1, primary2, synonyms2, primary3, synonyms3, primary4, synonyms4, primary5, synonyms5],
            label="Click an example to try it out"
        )
        
        # Clear functions
        def clear_text_only():
            """Clear only the text input field"""
            return ""
        
        def clear_dictionary_only():
            """Clear only the keyword dictionary fields"""
            return "", "", "", "", "", "", "", "", "", ""
        
        # Button functions
        find_btn.click(
            fn=process_text,
            inputs=[text_input, primary1, synonyms1, primary2, synonyms2, primary3, synonyms3, primary4, synonyms4, primary5, synonyms5],
            outputs=[results_output, highlighted_output, results_table_output, copy_output]
        )
        
        clear_text_btn.click(
            fn=clear_text_only,
            outputs=[text_input]
        )
        
        clear_dict_btn.click(
            fn=clear_dictionary_only,
            outputs=[primary1, synonyms1, primary2, synonyms2, primary3, synonyms3, primary4, synonyms4, primary5, synonyms5]
        )
        
        # Instructions
        gr.Markdown("""
        ## Format Guide
        
        **How to enter keywords:**
        - **Primary Keyword:** Enter the main/preferred term for a concept
        - **Synonyms:** Enter alternative terms separated by semicolons `;`
        - Leave rows blank if you don't need all 5 keyword groups
        - The tool will find ANY of these terms and return ALL related terms
        
        **Example:**
        - Primary: `Prisoner of War`
        - Synonyms: `POW; POWs; prisoner of war`
        
        **Special Handling:**
        - "US" is matched exactly to avoid confusion with the word "us"
        - Word boundaries are respected (prevents partial matches)
        - Results are alphabetised and deduplicated
        
        **How it works:**
        When ANY variant is found in your text (primary OR synonym), the tool returns the complete standardized set of terms for that concept.
        """)
        
        # Bottom horizontal line and footer
        gr.HTML("<hr style='margin-top: 40px; margin-bottom: 20px;'>")
        gr.HTML("""
        <div style="background-color: #f8f9fa; padding: 20px; border-radius: 8px; margin-top: 20px; text-align: center;">
                       <p style="font-size: 14px; line-height: 1.8; margin: 0;">
            The code for this tool was built with the aid of Claude Sonnet 4.
            </p>
        </div>
        """)
    
    return demo

if __name__ == "__main__":
    demo = create_interface()
    demo.launch()