import os
import pandas as pd
from docx import Document
import streamlit as st
import tempfile

def populate_docx_template(excel_file, template_file, temp_dir, selected_indices):
    """
    Populate template documents for selected rows, preserving watermarks
    """
    # Load Excel data
    df = pd.read_excel(excel_file)
    
    generated_files = []
    
    # Process each selected row in the Excel file
    for idx, index in enumerate(selected_indices):
        row = df.iloc[index]
        
        # Load the template for each row (this preserves all formatting including footers/watermarks)
        doc = Document(template_file)
        
        # Extract data from the Excel row
        client_name = str(row.get('Client Name ', '')) if pd.notna(row.get('Client Name ', '')) else ''
        matter_overview = str(row.get('Matter/Case Overview', '')) if pd.notna(row.get('Matter/Case Overview', '')) else ''
        description = str(row.get('Description  (in 500 words)', '')) if pd.notna(row.get('Description  (in 500 words)', '')) else ''
        value = str(row.get('Value ', '')) if pd.notna(row.get('Value ', '')) else ''
        lead_partners = str(row.get('Lead partners', '')) if pd.notna(row.get('Lead partners', '')) else ''
        team_partners = str(row.get('Team Members', '')) if pd.notna(row.get('Team Members', '')) else ''
        dates = str(row.get('Current Status (ongoing / completed )', '')) if pd.notna(row.get('Current Status (ongoing / completed )', '')) else ''
        publishable = str(row.get('Confidential / Publishable ', '')) if pd.notna(row.get('Confidential / Publishable ', '')) else ''

        if len(doc.tables) > 0:
            table = doc.tables[0]

            # Client Name
            table.cell(2, 0).text = client_name
            
            # Matter description
            matter_text = f"{matter_overview}\n\n{description}"
            table.cell(4, 0).text = matter_text
            
            # Deal value
            table.cell(5, 1).text = value
           
            # Lead partner(s)
            table.cell(12, 0).text = lead_partners
            table.cell(16, 0).text = team_partners
            
            table.cell(22, 2).text = dates

        # Create a temporary filename with index prefix to maintain order
        temp_filename = os.path.join(temp_dir, f"{idx:03d}_{client_name}.docx")

        # Save the populated document to the temporary directory
        doc.save(temp_filename)
        generated_files.append(temp_filename)
    
    return generated_files

def combine_documents_with_watermark(file_paths, output_file):
    """
    Combine multiple documents while preserving watermarks and footers
    """
    if not file_paths:
        return None
    
    # Create a new combined document
    combined_doc = Document(file_paths[0])  # Start with the first document to preserve all formatting
    
    # For each additional file
    for file_path in file_paths[1:]:
        # Add a page break
        combined_doc.add_page_break()
        
        # Open the document
        doc = Document(file_path)
        
        # Copy all content from the document
        for para in doc.paragraphs:
            p = combined_doc.add_paragraph()
            for run in para.runs:
                p.add_run(run.text).bold = run.bold
                
        # Copy all tables
        for table in doc.tables:
            t = combined_doc.add_table(rows=len(table.rows), cols=len(table.columns))
            for i, row in enumerate(table.rows):
                for j, cell in enumerate(row.cells):
                    if i < len(t.rows) and j < len(t.rows[i].cells):
                        t.rows[i].cells[j].text = cell.text
    
    # Save the combined document
    combined_doc.save(output_file)
    return output_file

def app():
    st.title("Legal 500 Document Generator")

    # Upload the Excel file
    excel_file = st.file_uploader("Upload your Excel file", type=["xlsx", "xls"])
    
    if excel_file is not None:
        # Read Excel file for selection interface
        df = pd.read_excel(excel_file)
        
        st.subheader("Select Records to Include")
        
        # Create tabs similar to benchmark code
        tab1, tab2 = st.tabs(["Publishable", "Confidential"])
        
        selected_indices = []
        
        # Publishable tab
        with tab1:
            st.write("Select publishable records (marked as 'No' in Excel):")
            
            # Filter to show only publishable records (where "No" means publishable)
            publishable_rows = []
            for i, row in df.iterrows():
                conf_val = row.get('Confidential / Publishable ', '')
                if pd.notna(conf_val) and (str(conf_val).lower() == 'no' or 'publishable' in str(conf_val).lower()):
                    publishable_rows.append(i)
            
            # Show publishable records
            for i in publishable_rows:
                row = df.iloc[i]
                client = row.get('Client Name ', '')
                matter = row.get('Matter/Case Overview', '')
                label = f"{client} - {matter}"
                
                if st.checkbox(label, key=f"pub_{i}"):
                    if i not in selected_indices:
                        selected_indices.append(i)
        
        # Confidential tab
        with tab2:
            st.write("Select confidential records (marked as 'Yes' in Excel):")
            
            # Filter to show only confidential records (where "Yes" means confidential)
            confidential_rows = []
            for i, row in df.iterrows():
                conf_val = row.get('Confidential / Publishable ', '')
                if pd.notna(conf_val) and (str(conf_val).lower() == 'yes' or 'confidential' in str(conf_val).lower()):
                    confidential_rows.append(i)
            
            # Show confidential records
            for i in confidential_rows:
                row = df.iloc[i]
                client = row.get('Client Name ', '')
                matter = row.get('Matter/Case Overview', '')
                label = f"{client} - {matter}"
                
                if st.checkbox(label, key=f"conf_{i}"):
                    if i not in selected_indices:
                        selected_indices.append(i)
        
        # Generate button
        if st.button("Generate Documents"):
            if selected_indices:
                with tempfile.TemporaryDirectory() as temp_dir:
                    # Use a hardcoded template path as in your original code
                    template_file = r"c:\Users\Administrator\Desktop\Legal 500.doc"
                    output_file = os.path.join(temp_dir, "combined_output.docx")
                    
                    # Progress indicator
                    progress_bar = st.progress(0)
                    status = st.empty()
                    
                    # Generate individual documents
                    status.text("Creating individual documents...")
                    generated_files = populate_docx_template(excel_file, template_file, temp_dir, selected_indices)
                    progress_bar.progress(50)
                    
                    # Combine documents with watermark preserved
                    status.text("Combining documents with watermarks...")
                    final_doc = combine_documents_with_watermark(generated_files, output_file)
                    progress_bar.progress(100)
                    
                    # Provide download button
                    if final_doc and os.path.exists(final_doc):
                        status.text("Documents generated successfully!")
                        with open(final_doc, "rb") as f:
                            st.download_button(
                                label="Download Combined Document",
                                data=f,
                                file_name="Legal_500_Combined.docx",
                                mime="application/vnd.openxmlformats-officedocument.wordprocessingml.document"
                            )
                    else:
                        st.error("Failed to generate combined document.")
            else:
                st.warning("Please select at least one record to generate documents.")

if __name__ == "__main__":
    app()    in this the first row is filled succesfullly ..after taht in second row i wanted to fill the same doc nd then append it to the page1 ....but after row 1 it is simply pasting text on blank white oages not in the form document .....i wanted to retian the watermark in fooetr on each page\