app.py

import streamlit as st
import os
import json
import re
from pathlib import Path
from datetime import datetime
from typing import Dict, List, Any, Optional, Tuple
import pymupdf as fitz  # PyMuPDF
from docx import Document
from docx.shared import Inches
import tempfile
from groq import Groq

# Configure Streamlit page
st.set_page_config(
    page_title="GLR Pipeline Automation - Enhanced",
    page_icon="📋",
    layout="wide",
    initial_sidebar_state="expanded"
)

class EnhancedGLRProcessor:
    """Enhanced GLR processor that handles multiple template formats and data types"""
    
    def __init__(self):
        self.output_dir = Path("output")
        self.output_dir.mkdir(exist_ok=True)
        self.groq_client = None
        self.initialize_groq()
    
    def initialize_groq(self):
        """Initialize Groq client with API key from Hugging Face secrets"""
        try:
            # Try to get API key from Hugging Face secrets
            api_key = st.secrets["GROQ_API_KEY"]
            
            if api_key:
                self.groq_client = Groq(api_key=api_key)
                st.success("✅ Groq API initialized successfully!")
            else:
                st.warning("⚠️ Groq API key not found. Using enhanced pattern matching only.")
                st.info("To use AI-powered extraction, add GROQ_API_KEY to your Hugging Face Space secrets.")
            
        except Exception as e:
            st.warning("⚠️ Using enhanced pattern matching only (Groq API not configured)")
            self.groq_client = None

    def extract_text_from_pdf(self, pdf_file) -> str:
        """Extract text from uploaded PDF file"""
        try:
            with tempfile.NamedTemporaryFile(delete=False, suffix='.pdf') as tmp_file:
                tmp_file.write(pdf_file.read())
                tmp_file_path = tmp_file.name
            
            doc = fitz.open(tmp_file_path)
            text = ""
            for page in doc:
                text += page.get_text()
            doc.close()
            
            os.unlink(tmp_file_path)
            return text.strip()
        except Exception as e:
            st.error(f"Error extracting text from PDF: {str(e)}")
            return ""
    
    def extract_placeholders_from_docx(self, docx_file) -> Tuple[List[str], str]:
        """Extract placeholders from DOCX template and detect template type"""
        try:
            with tempfile.NamedTemporaryFile(delete=False, suffix='.docx') as tmp_file:
                tmp_file.write(docx_file.read())
                tmp_file_path = tmp_file.name
            
            doc = Document(tmp_file_path)
            placeholders = set()
            full_text = ""
            
            # Extract from paragraphs
            for paragraph in doc.paragraphs:
                full_text += paragraph.text + "\n"
                # Enhanced regex to capture different placeholder formats
                matches = re.findall(r'\[([A-Z0-9_]+)\]', paragraph.text)
                placeholders.update(matches)
            
            # Extract from tables
            for table in doc.tables:
                for row in table.rows:
                    for cell in row.cells:
                        full_text += cell.text + "\n"
                        matches = re.findall(r'\[([A-Z0-9_]+)\]', cell.text)
                        placeholders.update(matches)
            
            os.unlink(tmp_file_path)
            
            # Detect template type based on placeholder patterns
            template_type = self._detect_template_type(list(placeholders), full_text)
            
            return list(placeholders), template_type
        except Exception as e:
            st.error(f"Error extracting placeholders from DOCX: {str(e)}")
            return [], "unknown"
    
    def _detect_template_type(self, placeholders: List[str], full_text: str) -> str:
        """Detect template type based on placeholder patterns and content"""
        xm8_count = sum(1 for p in placeholders if p.startswith('XM8_'))
        usaa_indicators = ['USAA', 'member', 'mortgage company']
        elevate_indicators = ['Elevate', 'Wayne', 'COV_NAME']
        guideone_indicators = ['GuideOne', 'Eberl Claims']
        
        if xm8_count > len(placeholders) * 0.5:  # More than 50% are XM8_ prefixed
            if any(indicator in full_text for indicator in elevate_indicators):
                return "elevate_wayne"
            elif any(indicator in full_text for indicator in guideone_indicators):
                return "guideone"
            else:
                return "xm8_generic"
        elif any(indicator in full_text for indicator in usaa_indicators):
            return "usaa"
        else:
            return "generic"
    
    def get_field_descriptions(self, template_type: str) -> Dict[str, str]:
        """Get field descriptions based on template type"""
        
        # Common field descriptions
        common_fields = {
            "DATE_LOSS": "Date when the loss/damage occurred (format: MM/DD/YYYY)",
            "DATE_RECEIVED": "Date when the claim was received (format: MM/DD/YYYY)", 
            "DATE_INSPECTED": "Date when the property was inspected (format: MM/DD/YYYY)",
            "DATE_CURRENT": "Current date (format: MM/DD/YYYY)",
            "INSURED_NAME": "Full name of the insured person/entity",
            "CLAIM_NUMBER": "Insurance claim number",
            "CLAIM_NUM": "Insurance claim number",
            "POLICY_NUMBER": "Insurance policy number",
            "PHONE": "Phone number",
            "EMAIL": "Email address",
            "TOL_CODE": "Type of Loss code",
            "TOL_DESC": "Type of Loss description (wind, hail, vehicle, etc.)"
        }
        
        # USAA specific fields
        usaa_fields = {
            **common_fields,
            "INSURED_H_STREET": "Insured home street address",
            "INSURED_H_CITY": "Insured home city",
            "INSURED_H_STATE": "Insured home state (2 letters)",
            "INSURED_H_ZIP": "Insured home ZIP code",
            "MORTGAGE_CO": "Mortgage company name",
            "MORTGAGEE": "Mortgagee information"
        }
        
        # XM8 (Elevate/GuideOne) specific fields
        xm8_fields = {
            **common_fields,
            "XM8_DATE_LOSS": "Date when the loss/damage occurred",
            "XM8_DATE_RECEIVED": "Date when the claim was received",
            "XM8_DATE_INSPECTED": "Date when the property was inspected",
            "XM8_DATE_CURRENT": "Current date",
            "XM8_INSURED_NAME": "Full name of the insured person/entity",
            "XM8_INSURED_P_STREET": "Property street address",
            "XM8_INSURED_P_CITY": "Property city",
            "XM8_INSURED_P_STATE": "Property state (2 letters)",
            "XM8_INSURED_P_ZIP": "Property ZIP code",
            "XM8_CLAIM_NUM": "Insurance claim number",
            "XM8_POLICY_NUM": "Insurance policy number",
            "XM8_ESTIMATOR_NAME": "Name of the estimator/adjuster",
            "XM8_ESTIMATOR_E_MAIL": "Estimator email address",
            "XM8_ESTIMATOR_C_PHONE": "Estimator cell phone",
            "XM8_ESTIMATOR_B_PHONE": "Estimator business phone",
            "XM8_TOL_DESC": "Type of Loss description",
            "XM8_REFERENCE_COMPANY": "Reference company name",
            "XM8_CLAIM_REP_NAME": "Claim representative name",
            "XM8_FILE_NO": "File number",
            "XM8_COV_NAME_1": "Coverage name 1",
            "XM8_COV_NAME_2": "Coverage name 2", 
            "XM8_COV_NAME_3": "Coverage name 3",
            "XM8_COV_RCV_1": "Coverage RCV amount 1",
            "XM8_COV_RCV_2": "Coverage RCV amount 2",
            "XM8_COV_RCV_3": "Coverage RCV amount 3",
            "XM8_LR_RC_LOSS": "Loss reserve RC amount",
            "XM8_SUM_RECOVERABLE_DEPRECIATION": "Sum of recoverable depreciation",
            "XM8_SUM_NONRECOVERABLE_DEPRECIATION": "Sum of non-recoverable depreciation",
            "XM8_SUM_ACV": "Sum of ACV",
            "XM8_SUM_DEDUCTIBLE_APPLIED": "Sum of deductible applied",
            "XM8_LR_ACV_CLAIM": "Loss reserve ACV claim amount"
        }
        
        if template_type == "usaa":
            return usaa_fields
        elif template_type in ["elevate_wayne", "guideone", "xm8_generic"]:
            return xm8_fields
        else:
            return common_fields
    
    def extract_data_with_groq(self, pdf_texts: List[str], placeholders: List[str], template_type: str) -> Dict[str, str]:
        """Use Groq API to extract data from PDF texts with template-specific prompting"""
        if self.groq_client is None:
            st.warning("⚠️ Groq API not available. Using enhanced fallback extraction...")
            return self._enhanced_fallback_extraction("\n\n".join(pdf_texts), placeholders, template_type)
        
        try:
            combined_text = "\n\n".join(pdf_texts)
            
            # Truncate if too long but keep more context
            if len(combined_text) > 15000:
                combined_text = combined_text[:15000] + "..."
            
            # Get template-specific field descriptions
            field_descriptions = self.get_field_descriptions(template_type)
            
            # Build detailed extraction instructions
            extraction_instructions = []
            for field in placeholders:
                description = field_descriptions.get(field, f"Extract the {field.lower().replace('_', ' ')}")
                extraction_instructions.append(f"- {field}: {description}")
            
            # Template-specific context
            template_context = self._get_template_context(template_type)
            
            prompt = f"""You are an expert insurance document analyst specializing in {template_type.upper()} templates. Extract the following specific information from this insurance report text.

TEMPLATE TYPE: {template_type.upper()}
{template_context}

REQUIRED FIELDS TO EXTRACT:
{chr(10).join(extraction_instructions)}

EXTRACTION GUIDELINES:
1. Search the ENTIRE document text carefully for each field
2. Look for variations in field names and formats
3. For dates, convert to MM/DD/YYYY format if in different format
4. For names, extract complete full names as they appear
5. For addresses, look for complete street addresses with numbers
6. For states, use 2-letter codes (California = CA, New York = NY, etc.)
7. Look in headers, footers, tables, and body text
8. If a field has multiple possible values, choose the most complete/relevant one
9. For claim/policy numbers, look for alphanumeric codes
10. For estimator information, look for adjuster or inspector names
11. Only use "N/A" if the information is absolutely not present anywhere

DOCUMENT TEXT TO ANALYZE:
{combined_text}

Return ONLY a valid JSON object with the extracted data:
{{
  "FIELD_NAME": "extracted_value"
}}

JSON:"""

            # Make API call to Groq
            completion = self.groq_client.chat.completions.create(
                model="deepseek-r1-distill-llama-70b",
                messages=[
                    {
                        "role": "system",
                        "content": f"You are an expert at extracting structured data from {template_type} insurance documents. Always return valid JSON and be thorough in your extraction."
                    },
                    {
                        "role": "user",
                        "content": prompt
                    }
                ],
                temperature=0.0,
                max_tokens=4000,
                top_p=0.9,
                stream=False,
                stop=None,
            )
            
            response = completion.choices[0].message.content
            
            # Parse JSON response
            json_text = response.strip()
            
            # Remove markdown formatting
            if json_text.startswith("```json"):
                json_text = json_text[7:]
            if json_text.startswith("```"):
                json_text = json_text[3:]
            if json_text.endswith("```"):
                json_text = json_text[:-3]
            json_text = json_text.strip()
            
            # Extract JSON object
            json_match = re.search(r'\{.*\}', json_text, re.DOTALL)
            if json_match:
                json_text = json_match.group()
            
            try:
                extracted_data = json.loads(json_text)
            except json.JSONDecodeError:
                # Try to fix common JSON issues
                json_text = json_text.replace("'", '"')
                json_text = re.sub(r',\s*}', '}', json_text)
                extracted_data = json.loads(json_text)
            
            # Ensure all placeholders are present and post-process
            final_data = {}
            for placeholder in placeholders:
                if placeholder in extracted_data:
                    value = str(extracted_data[placeholder]).strip()
                    if value and value.lower() not in ['n/a', 'null', 'none', '', 'not found']:
                        final_data[placeholder] = value
                    else:
                        fallback_value = self._extract_single_field(combined_text, placeholder, template_type)
                        final_data[placeholder] = fallback_value
                else:
                    fallback_value = self._extract_single_field(combined_text, placeholder, template_type)
                    final_data[placeholder] = fallback_value
            
            st.success("✅ Groq AI extraction completed successfully!")
            return final_data
            
        except Exception as e:
            st.warning(f"Groq API extraction failed: {str(e)}")
            return self._enhanced_fallback_extraction(combined_text, placeholders, template_type)
    
    def _get_template_context(self, template_type: str) -> str:
        """Get template-specific context for better extraction"""
        contexts = {
            "usaa": """
USAA CONTEXT:
- Look for member information and mortgage company details
- Claims often involve wind, hail, or vehicle damage
- Adjuster company is typically Alacrity Solutions
- Look for member names like "Richard Daly"
            """,
            "elevate_wayne": """
ELEVATE/WAYNE CONTEXT:
- Look for Elevate Claims Solutions company information
- Claims often involve vehicle collisions or property damage
- Look for coverage amounts and depreciation details
- Estimator is typically from Elevate Claims Solutions
            """,
            "guideone": """
GUIDEONE CONTEXT:
- Look for GuideOne Insurance Company information
- Claims often involve church properties and wind/hail damage
- Look for Eberl Claims Service as adjuster
- Focus on property damage assessments
            """,
            "generic": """
GENERIC CONTEXT:
- Extract standard insurance claim information
- Look for basic claim details and property information
            """
        }
        return contexts.get(template_type, contexts["generic"])
    
    def _enhanced_fallback_extraction(self, text: str, placeholders: List[str], template_type: str) -> Dict[str, str]:
        """Enhanced fallback extraction with template awareness"""
        st.info(f"🔍 Using enhanced pattern matching extraction for {template_type} template...")
        
        fallback_data = {}
        
        for placeholder in placeholders:
            value = self._extract_single_field(text, placeholder, template_type)
            fallback_data[placeholder] = value
        
        return fallback_data
    
    def _extract_single_field(self, text: str, field: str, template_type: str) -> str:
        """Extract a single field using enhanced pattern matching with template awareness"""
        text_lower = text.lower()
        
        # Remove XM8_ prefix for pattern matching
        clean_field = field.replace("XM8_", "").replace("INSURED_H_", "INSURED_").replace("INSURED_P_", "INSURED_")
        
        # Date patterns
        if "DATE" in clean_field:
            date_patterns = [
                r'(?:date[^:]*taken|inspected|loss|received)[\s:]*(\d{1,2}[/\-]\d{1,2}[/\-]\d{2,4})',
                r'\b(\d{1,2}[/\-]\d{1,2}[/\-]\d{2,4})\b',
                r'\b([A-Za-z]+ \d{1,2}, \d{4})\b'
            ]
            for pattern in date_patterns:
                match = re.search(pattern, text, re.IGNORECASE)
                if match:
                    return match.group(1)
        
        # Name patterns (enhanced for different templates)
        elif "NAME" in clean_field:
            if template_type == "usaa":
                # Look for member names in USAA format
                patterns = [
                    r'(?:insured|member)[\s:]*([A-Z][a-z]+ [A-Z][a-z]+)',
                    r'Insured:\s*([A-Z][a-z]+ [A-Z][a-z]+)',
                ]
            elif "ESTIMATOR" in field:
                # Look for estimator/adjuster names
                patterns = [
                    r'(?:taken by|estimator|adjuster)[\s:]*([A-Z][a-z]+ [A-Z][a-z]+)',
                    r'Taken By:\s*([A-Z][a-z]+ [A-Z][a-z]+)',
                ]
            else:
                patterns = [
                    r'(?:insured|name)[\s:]*([A-Z][a-z]+ [A-Z][a-z]+(?:\s+[A-Z][a-z]+)?)',
                ]
            
            for pattern in patterns:
                match = re.search(pattern, text, re.IGNORECASE)
                if match:
                    return match.group(1).strip()
        
        # Address patterns
        elif "STREET" in clean_field:
            street_patterns = [
                r'\b(\d+\s+[A-Z][A-Za-z\s]+(?:ST|STREET|AVE|AVENUE|RD|ROAD|DR|DRIVE|LN|LANE|CT|COURT|BLVD|BOULEVARD|PL|PLACE|WAY)\.?)\b',
                r'(?:address|street)[\s:]+(\d+\s+[A-Za-z\s]+)',
                r'Risk address\s*(\d+\s+[A-Za-z\s]+)',
                r'(\d{3,5}\s+[A-Za-z][A-Za-z\s]+(?:Dr|Drive|St|Street|Ave|Avenue)\.?)'
            ]
            for pattern in street_patterns:
                match = re.search(pattern, text, re.IGNORECASE)
                if match:
                    return match.group(1).strip()
        
        # City patterns
        elif "CITY" in clean_field:
            city_patterns = [
                r'(?:city)[\s:]+([A-Z][a-z\s]+?)(?:,|\s+[A-Z]{2}\s|\n)',
                r',\s*([A-Z][a-z\s]+?)\s+[A-Z]{2}\s+\d{5}',
                r'\b([A-Z][a-z]+(?: [A-Z][a-z]+)*),\s*[A-Z]{2}\b'
            ]
            for pattern in city_patterns:
                match = re.search(pattern, text)
                if match:
                    city = match.group(1).strip()
                    if len(city) > 2:
                        return city
        
        # State patterns
        elif "STATE" in clean_field:
            state_patterns = [
                r'\b([A-Z]{2})\s+\d{5}',
                r'(?:state)[\s:]+([A-Z]{2})\b',
                r',\s*[A-Za-z\s]+,?\s*([A-Z]{2})\s+\d{5}'
            ]
            for pattern in state_patterns:
                match = re.search(pattern, text)
                if match:
                    return match.group(1)
        
        # ZIP patterns
        elif "ZIP" in clean_field:
            zip_patterns = [
                r'\b(\d{5}-\d{4})\b',
                r'\b(\d{5})\b(?!\d)'
            ]
            for pattern in zip_patterns:
                match = re.search(pattern, text)
                if match:
                    return match.group(1)
        
        # Phone patterns
        elif "PHONE" in clean_field:
            phone_patterns = [
                r'\b(\d{3}[-.\s]?\d{3}[-.\s]?\d{4})\b',
                r'\((\d{3})\)\s*(\d{3})[-.\s]?(\d{4})'
            ]
            for pattern in phone_patterns:
                match = re.search(pattern, text)
                if match:
                    if len(match.groups()) == 1:
                        return match.group(1)
                    else:
                        return f"({match.group(1)}) {match.group(2)}-{match.group(3)}"
        
        # Email patterns
        elif "EMAIL" in clean_field:
            email_pattern = r'\b([a-zA-Z0-9._%+-]+@[a-zA-Z0-9.-]+\.[a-zA-Z]{2,})\b'
            match = re.search(email_pattern, text)
            if match:
                return match.group(1)
        
        # Claim/Policy number patterns
        elif "CLAIM" in clean_field or "POLICY" in clean_field:
            number_patterns = [
                r'(?:claim|policy)[\s#:]*([A-Z0-9\-]+)',
                r'Claim #:\s*([A-Z0-9\-]+)',
                r'Policy #:\s*([A-Z0-9\-]+)',
                r'\b([A-Z]{2,4}\d{6,})\b',
                r'\b(\d{8,})\b'
            ]
            for pattern in number_patterns:
                match = re.search(pattern, text, re.IGNORECASE)
                if match:
                    return match.group(1)
        
        # Coverage and financial patterns
        elif "COV_" in field or "RCV" in field or "ACV" in field:
            # Look for coverage names and amounts
            if "NAME" in field:
                cov_patterns = [
                    r'(?:dwelling|coverage|building)[\s:]*([A-Za-z\s]+)',
                    r'([A-Z][a-z]+(?: [A-Z][a-z]+)*)\s*\$'
                ]
            else:
                cov_patterns = [
                    r'\$([0-9,]+\.?\d*)',
                    r'([0-9,]+\.?\d*)'
                ]
            
            for pattern in cov_patterns:
                match = re.search(pattern, text)
                if match:
                    return match.group(1).strip()
        
        # File number patterns
        elif "FILE" in field:
            file_patterns = [
                r'(?:file|elevate file)[\s#:]*([A-Z0-9\-]+)',
                r'\b([A-Z]+-[A-Z]+)\b'
            ]
            for pattern in file_patterns:
                match = re.search(pattern, text, re.IGNORECASE)
                if match:
                    return match.group(1)
        
        # Company/Reference patterns
        elif "COMPANY" in field or "REFERENCE" in field:
            company_patterns = [
                r'(Wayne Mutual Insurance Company)',
                r'(GuideOne Insurance Company)',
                r'(USAA)',
                r'([A-Z][a-z]+ Insurance Company)'
            ]
            for pattern in company_patterns:
                match = re.search(pattern, text)
                if match:
                    return match.group(1)
        
        # TOL (Type of Loss) patterns
        elif "TOL" in field:
            tol_patterns = [
                r'(?:wind|hail|vehicle|collision|storm|water|fire)',
                r'(?:type of loss)[\s:]*([A-Za-z]+)'
            ]
            for pattern in tol_patterns:
                match = re.search(pattern, text, re.IGNORECASE)
                if match:
                    return match.group(0) if len(match.groups()) == 0 else match.group(1)
        
        return "N/A"
    
    def populate_docx_template(self, docx_file, extracted_data: Dict[str, str]) -> str:
        """Populate DOCX template with extracted data"""
        try:
            with tempfile.NamedTemporaryFile(delete=False, suffix='.docx') as tmp_file:
                tmp_file.write(docx_file.read())
                tmp_file_path = tmp_file.name
            
            doc = Document(tmp_file_path)
            replacements_made = 0
            
            # Replace in paragraphs
            for paragraph in doc.paragraphs:
                for placeholder, value in extracted_data.items():
                    if f"[{placeholder}]" in paragraph.text:
                        paragraph.text = paragraph.text.replace(f"[{placeholder}]", str(value))
                        replacements_made += 1
            
            # Replace in tables
            for table in doc.tables:
                for row in table.rows:
                    for cell in row.cells:
                        for paragraph in cell.paragraphs:
                            for placeholder, value in extracted_data.items():
                                if f"[{placeholder}]" in paragraph.text:
                                    paragraph.text = paragraph.text.replace(f"[{placeholder}]", str(value))
                                    replacements_made += 1
            
            timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
            output_filename = f"GLR_Report_{timestamp}.docx"
            output_path = self.output_dir / output_filename
            
            doc.save(str(output_path))
            os.unlink(tmp_file_path)
            
            st.info(f"Made {replacements_made} field replacements in the template")
            return str(output_path)
            
        except Exception as e:
            st.error(f"Error populating template: {str(e)}")
            return None

def main():
    """Main Streamlit application"""
    st.title("📋 GLR Pipeline Automation - Enhanced Multi-Template Support")
    st.markdown("**Automate General Loss Report generation with support for USAA, Elevate/Wayne, and GuideOne templates**")
    
    # Initialize processor
    processor = EnhancedGLRProcessor()
    
    # Show template support information
    with st.expander("📋 Supported Template Types"):
        st.markdown("""
        **USAA Templates:**
        - Fields: [DATE_LOSS], [INSURED_NAME], [MORTGAGE_CO], etc.
        - Handles member information and mortgage details
        
        **Elevate/Wayne Templates:**
        - Fields: [XM8_DATE_LOSS], [XM8_INSURED_NAME], [XM8_COV_NAME_1], etc.
        - Handles coverage amounts and depreciation details
        
        **GuideOne Templates:**
        - Fields: [XM8_DATE_INSPECTED], [XM8_ESTIMATOR_NAME], [XM8_TOL_DESC], etc.
        - Handles church properties and detailed damage assessments
        """)
    
    col1, col2 = st.columns([1, 1])
    
    with col1:
        st.header("📄 Upload Template")
        template_file = st.file_uploader("Upload GLR Template (.docx)", type=['docx'])
    
    with col2:
        st.header("📋 Upload Reports")
        photo_reports = st.file_uploader("Upload PDF Reports", type=['pdf'], accept_multiple_files=True)
    
    if template_file and photo_reports:
        st.header("🔄 Processing")
        
        with st.spinner("Processing..."):
            placeholders, template_type = processor.extract_placeholders_from_docx(template_file)
            
            if placeholders:
                st.success(f"✅ Detected **{template_type.upper()}** template with {len(placeholders)} placeholders")
                
                # Show detected placeholders
                with st.expander("🔍 Detected Placeholders"):
                    st.write(", ".join(placeholders))
                
                pdf_texts = []
                for pdf_file in photo_reports:
                    text = processor.extract_text_from_pdf(pdf_file)
                    if text:
                        pdf_texts.append(text)
                        st.success(f"✅ Processed {pdf_file.name} ({len(text)} characters)")
                
                if pdf_texts:
                    extracted_data = processor.extract_data_with_groq(pdf_texts, placeholders, template_type)
                    
                    st.subheader("📊 Extracted Data")
                    
                    # Show extraction statistics
                    non_na_count = sum(1 for v in extracted_data.values() if v != "N/A")
                    extraction_rate = (non_na_count / len(extracted_data)) * 100
                    
                    col1, col2, col3 = st.columns(3)
                    with col1:
                        st.metric("Template Type", template_type.upper())
                    with col2:
                        st.metric("Extraction Rate", f"{extraction_rate:.1f}%")
                    with col3:
                        st.metric("Fields Found", f"{non_na_count}/{len(extracted_data)}")
                    
                    # Create form for editing data
                    with st.form("data_form"):
                        st.subheader("📝 Review and Edit Extracted Data")
                        edited_data = {}
                        
                        # Group fields by category for better organization
                        field_categories = {
                            "Basic Information": [],
                            "Dates": [],
                            "Contact Information": [],
                            "Financial Information": [],
                            "Other": []
                        }
                        
                        for placeholder in placeholders:
                            if any(date_word in placeholder for date_word in ["DATE"]):
                                field_categories["Dates"].append(placeholder)
                            elif any(contact_word in placeholder for contact_word in ["PHONE", "EMAIL", "NAME"]):
                                field_categories["Contact Information"].append(placeholder)
                            elif any(fin_word in placeholder for fin_word in ["COV", "RCV", "ACV", "DEDUCTIBLE", "DEPRECIATION"]):
                                field_categories["Financial Information"].append(placeholder)
                            elif any(basic_word in placeholder for basic_word in ["CLAIM", "POLICY", "FILE", "STREET", "CITY", "STATE", "ZIP"]):
                                field_categories["Basic Information"].append(placeholder)
                            else:
                                field_categories["Other"].append(placeholder)
                        
                        # Display fields by category
                        for category, fields in field_categories.items():
                            if fields:
                                st.markdown(f"**{category}**")
                                cols = st.columns(2)
                                for i, placeholder in enumerate(fields):
                                    with cols[i % 2]:
                                        value = extracted_data.get(placeholder, "N/A")
                                        # Color code based on whether value was found
                                        if value != "N/A":
                                            st.markdown(f"✅ **{placeholder}**")
                                        else:
                                            st.markdown(f"❌ **{placeholder}**")
                                        edited_data[placeholder] = st.text_input(
                                            f"Enter {placeholder}", 
                                            value=value, 
                                            key=placeholder,
                                            help=processor.get_field_descriptions(template_type).get(placeholder, ""),
                                            label_visibility="collapsed"
                                        )
                                st.markdown("---")
                        
                        generate_report = st.form_submit_button("🚀 Generate Report", type="primary")
                    
                    # Handle report generation outside the form
                    if generate_report:
                        template_file.seek(0)
                        output_path = processor.populate_docx_template(template_file, edited_data)
                        
                        if output_path:
                            st.success("✅ Report generated successfully!")
                            
                            # Store the file data in session state for download
                            with open(output_path, 'rb') as file:
                                st.session_state['report_data'] = file.read()
                                st.session_state['report_filename'] = os.path.basename(output_path)
                    
                    # Download button outside the form
                    if 'report_data' in st.session_state:
                        st.download_button(
                            label="📥 Download Generated Report",
                            data=st.session_state['report_data'],
                            file_name=st.session_state['report_filename'],
                            mime="application/vnd.openxmlformats-officedocument.wordprocessingml.document",
                            type="primary"
                        )
                        
                        # Show summary of what was processed
                        st.subheader("📋 Processing Summary")
                        summary_col1, summary_col2 = st.columns(2)
                        with summary_col1:
                            st.info(f"**Template Type:** {template_type.upper()}")
                            st.info(f"**Fields Processed:** {len(placeholders)}")
                            st.info(f"**PDFs Processed:** {len(pdf_texts)}")
                        with summary_col2:
                            st.info(f"**Extraction Rate:** {extraction_rate:.1f}%")
                            st.info(f"**Fields Found:** {non_na_count}")
                            st.info(f"**Output File:** {st.session_state['report_filename']}")
                else:
                    st.error("No text extracted from PDFs")
            else:
                st.error("No placeholders found in template")

if __name__ == "__main__":
    main()