import os
import pandas as pd
from docx import Document
from docx.shared import RGBColor
from docxcompose.composer import Composer
import streamlit as st
import tempfile

def populate_docx_template(excel_file, template_files, temp_dir, selected_indices):
    """
    Generate one full Word document per selected Excel row from the template,
    filling all relevant cells in the table while preserving formatting.
    
    template_files should be a dict with 'publishable' and 'confidential' keys
    """
    df = pd.read_excel(excel_file)

    generated_files = []

    for idx, row_idx in enumerate(selected_indices):
        row = df.iloc[row_idx]
        
        # Determine if the record is confidential or publishable
        conf_status = str(row.get('Confidential / Publishable ', '')).lower()
        is_confidential = (conf_status == 'yes' or 'confidential'  in conf_status)
        
        # Select the appropriate template
        template_key = 'confidential' if is_confidential else 'publishable'
        doc = Document(template_files[template_key])

        # Assuming first table holds the data fields.
        if not doc.tables:
            st.error("Template file appears to have no tables.")
            return []

        table = doc.tables[0]

        # Prepare data - try different column name variations for description
        serial_number = str(row.get('S No. ', '').strip() if pd.notna(row.get('S No. ', ''))else '')
        client_name = str(row.get('Client Name ', '').strip() if pd.notna(row.get('Client Name ', '')) else '')
        matter_overview = str(row.get('Matter/Case Overview', '').strip() if pd.notna(row.get('Matter/Case Overview', '')) else '')
        
        # Try different possible column names for description
        desc_col_names = ['Description  (in 500 words) ', 'Description  (in 500 words)', 'Description (in 500 words)', 'Description']
        description = ""
        for col_name in desc_col_names:
            if col_name in df.columns:
                description = str(row.get(col_name, '').strip() if pd.notna(row.get(col_name, '')) and hasattr(row.get(col_name, ''), 'strip') else row.get(col_name, ''))
                break
                
        role_col_names = ['Role of S&A', 'Role of SA', 'Role']
        role_of_sa = ""
        for col_name in role_col_names:
            if col_name in df.columns:
                role_of_sa = str(row.get(col_name, '').strip() if pd.notna(row.get(col_name, '')) and hasattr(row.get(col_name, ''), 'strip') else row.get(col_name, ''))
                break
                
        sig_col_names = ['The significance, impact and outcome of the case', 'Significance', 'Impact']
        significance = ""
        for col_name in sig_col_names:
            if col_name in df.columns:
                significance = str(row.get(col_name, '').strip() if pd.notna(row.get(col_name, '')) and hasattr(row.get(col_name, ''), 'strip') else row.get(col_name, ''))
                break
                
        value = str(row.get('Value ', '').strip() if pd.notna(row.get('Value ', '')) and hasattr(row.get('Value ', ''), 'strip') else row.get('Value ', ''))
        cross_border = str(row.get('Is this a cross-border matter? If yes, please indicate the jurisdictions involved.', '').strip() if pd.notna(row.get('Is this a cross-border matter? If yes, please indicate the jurisdictions involved.', '')) and hasattr(row.get('Is this a cross-border matter? If yes, please indicate the jurisdictions involved.', ''), 'strip') else row.get('Is this a cross-border matter? If yes, please indicate the jurisdictions involved.', ''))
        lead_partners = str(row.get('Lead partners', '').strip() if pd.notna(row.get('Lead partners', '')) and hasattr(row.get('Lead partners', ''), 'strip') else row.get('Lead partners', ''))
        team_partners = str(row.get('Team Members', '').strip() if pd.notna(row.get('Team Members', '')) and hasattr(row.get('Team Members', ''), 'strip') else row.get('Team Members', ''))
        
        other_firms = str(row.get('Other Firms/counsels Involved ', '').strip() if pd.notna(row.get('Other Firms/counsels Involved ', '')) and hasattr(row.get('Other Firms/counsels Involved ', ''), 'strip') else row.get('Other Firms/counsels Involved ', ''))

        # Fix for datetime objects - properly handle the Current Status field
        current_status = row.get('Current Status (ongoing / completed )', '')
        dates = str(current_status) if pd.notna(current_status) else ''

        try:
            # Fill standard cells first
            table.cell(2, 0).text = client_name
            if is_confidential:
                for paragraph in table.cell(2, 0).paragraphs:
                    for run in paragraph.runs:
                        run.font.color.rgb = RGBColor(255, 0, 0)
                        
            # Fill matter overview
            overview_cell = table.cell(4, 0)
            overview_cell.text = ""
            
            # Create a new paragraph for the overview header
            overview_header = overview_cell.add_paragraph()
            header_run = overview_header.add_run("Overview")
            header_run.bold = True
            if is_confidential:
                header_run.font.color.rgb = RGBColor(255, 0, 0)
            
            # Create a new paragraph for the description text
            desc_para = overview_cell.add_paragraph()
            desc_run = desc_para.add_run(description)
            if is_confidential:
                desc_run.font.color.rgb = RGBColor(255, 0, 0)
            
            # Add a line break and the Role of S&A header
            role_header = overview_cell.add_paragraph()
            role_header.add_run("\n")  # Add extra line break
            role_run = role_header.add_run("Role of S&A")
            role_run.bold = True
            if is_confidential:
                role_run.font.color.rgb = RGBColor(255, 0, 0)
            
            # Add the role text
            role_para = overview_cell.add_paragraph()
            role_text_run = role_para.add_run(role_of_sa)
            if is_confidential:
                role_text_run.font.color.rgb = RGBColor(255, 0, 0)
            
            # Add a line break and Significance header
            sig_header = overview_cell.add_paragraph()
            sig_header.add_run("\n")  # Add extra line break
            sig_run = sig_header.add_run("Significance & Impact of the Matter")
            sig_run.bold = True
            if is_confidential:
                sig_run.font.color.rgb = RGBColor(255, 0, 0)
            
            # Add the significance text
            sig_para = overview_cell.add_paragraph()
            sig_text_run = sig_para.add_run(significance)
            if is_confidential:
                sig_text_run.font.color.rgb = RGBColor(255, 0, 0)
                
            # Fill the remaining standard cells
            table.cell(5, 1).text = value
            table.cell(20,2).txt = other_firms
            table.cell(8,1).text = cross_border
            table.cell(12, 0).text = lead_partners
            
            # Add "S&A law offices" to cell (12,1)
            table.cell(12, 2).text = "S&A law offices"
            
            table.cell(16, 0).text = team_partners
            
            # Add "S&A law offices" to cell (16,1)
            table.cell(16, 2).text = "S&A law offices"
            
            table.cell(22, 3).text = dates
            
            # Apply red formatting to all other cells if confidential
            if is_confidential:
                for cell_pos in [(5, 1),(8,1), (12, 0), (12, 2), (16, 0), (16, 2), (22, 3)]:
                    cell = table.cell(cell_pos[0], cell_pos[1])
                    for paragraph in cell.paragraphs:
                        for run in paragraph.runs:
                            run.font.color.rgb = RGBColor(255, 0, 0)

        except IndexError as e:
            st.error("Template table structure mismatch. Please check your template cell indices.")
            return []
        except Exception as e:
            st.error(f"Error processing document: {str(e)}")
            return []

        # Sanitize client name for filename
        safe_client_name = "".join(c for c in client_name if c.isalnum() or c in (" ", "_")).rstrip()
        safe_client_name = safe_client_name.replace(" ", "_") or "client"

        # Add prefix to indicate confidential/publishable status
        prefix = "CONF" if is_confidential else "PUB"
        outpath = os.path.join(temp_dir, f"{idx:03d}_{prefix}_{safe_client_name}.docx")
        doc.save(outpath)
        generated_files.append(outpath)

    return generated_files

def combine_documents_with_watermark(file_paths, output_file):
    """
    Combine multiple .docx files preserving full formatting like tables, headers, footers, watermarks.
    """
    if not file_paths:
        return None

    combined_doc = Document(file_paths[0])
    composer = Composer(combined_doc)

    for fp in file_paths[1:]:
        doc = Document(fp)
        composer.append(doc)

    composer.save(output_file)
    return output_file

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

    # Upload Excel
    excel_file = st.file_uploader("Upload Excel file with data", type=["xlsx", "xls"])

    # Template file inputs
    st.subheader("Template Files")
    col1, col2 = st.columns(2)
    
    with col1:
        publishable_template_path = st.text_input(
            "Publishable Template Path", 
            value="C:\\Users\\Administrator\\Desktop\\Legal 500.doc"
        )
    
    with col2:
        confidential_template_path = st.text_input(
            "Confidential Template Path", 
            value="C:\\Users\\Administrator\\Desktop\\Legal 500..doc"
        )
        
    if excel_file is not None:
        # Load the Excel for showing selections
        df = pd.read_excel(excel_file)

        st.subheader("Select Records to Include")

        tab1, tab2 = st.tabs(["Publishable", "Confidential"])
        selected_indices = []

        with tab1:
            st.write("Select publishable records (marked as 'No' or 'Publishable'):")
            publishable_rows = []
            for i, row in df.iterrows():
                pval = row.get('Confidential / Publishable ', '')
                if pd.notna(pval) and (str(pval).lower() == 'no' or 'publishable' in str(pval).lower()):
                    publishable_rows.append(i)
            for i in publishable_rows:
                row = df.iloc[i]
                # Display both client name and matter value for easier identification
                serial_number =row.get('S No. ', '')
                value = row.get('Value ', '')
                label = f"{row.get('S No.','')}-{row.get('Client Name ', '')} -  ({value})"
                if st.checkbox(label, key=f"pub_{i}"):
                    if i not in selected_indices:
                        selected_indices.append(i)

        with tab2:
            st.write("Select confidential records (marked as 'Yes', 'Confidential', or 'Maintain Confidentiality'):")
            confidential_rows = []
            for i, row in df.iterrows():
                pval = row.get('Confidential / Publishable ', '')
                if pd.notna(pval) and (str(pval).lower() == 'yes' or 
                                       'confidential' in str(pval).lower() or 
                                       'maintain confidentiality' in str(pval).lower()):
                    confidential_rows.append(i)
            for i in confidential_rows:
                row = df.iloc[i]
                # Display both client name and matter value for easier identification
                serial_number =row.get('S No.', '')
                value = row.get('Value ', '')
                label = f"{row.get('S No.','')}-{row.get('Client Name ', '')} -  ({value})"
                if st.checkbox(label, key=f"conf_{i}"):
                    if i not in selected_indices:
                        selected_indices.append(i)


        if st.button("Generate Documents"):
            if selected_indices:
                with tempfile.TemporaryDirectory() as temp_dir:
                    # Validate template paths
                    if not os.path.exists(publishable_template_path):
                        st.error(f"Publishable template file not found: {publishable_template_path}")
                        return
                    
                    if not os.path.exists(confidential_template_path):
                        st.error(f"Confidential template file not found: {confidential_template_path}")
                        return
                    
                    # Create template dictionary
                    template_files = {
                        'publishable': publishable_template_path,
                        'confidential': confidential_template_path
                    }

                    progress_bar = st.progress(0)
                    status = st.empty()

                    status.text("Populating templates per selected record...")
                    
                    with st.spinner("Processing..."):
                        generated_files = populate_docx_template(excel_file, template_files, temp_dir, selected_indices)
                    
                    progress_bar.progress(50)

                    if not generated_files:
                        st.error("No documents generated. Check template and data.")
                        return

                    status.text("Combining generated documents (preserving watermark and formatting)...")
                    combined_path = os.path.join(temp_dir, "Legal_500_Combined.docx")
                    combine_documents_with_watermark(generated_files, combined_path)
                    progress_bar.progress(100)

                    if os.path.exists(combined_path):
                        status.text("Documents generated successfully!")
                        with open(combined_path, "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("Select at least one record to generate documents.")

if __name__ == "__main__":
    app()   in this in the selection box part i wanted the serial number to be displayed also and same as in excel but from this im not getting seril number brfore the client nama and valuelable.....also i want if we select 4 punlishable  and  2 non ...so in cell as we have publishable matter append the number with this ...1 2 3 4 ....and for non oubishable also 1 2 3 4