#!/usr/bin/env python3
import json
import os
import re
import openpyxl
from openpyxl.utils import get_column_letter
from direct_xml_update import update_excel_with_direct_xml

def update_excel_variables(excel_path):
    """
    Update the Variables sheet in the Excel file with values from config.json
    
    Args:
        excel_path (str): Path to the Excel file to update
    
    Returns:
        bool: True if successful, False otherwise
    """
    # Define paths
    script_dir = os.path.dirname(os.path.abspath(__file__))
    config_path = os.path.join(script_dir, 'config.json')
    
    try:
        # Load config.json
        with open(config_path, 'r') as f:
            config = json.load(f)
            user_data = config.get('user_data', {})
        
        # Load Excel workbook
        print(f"Opening Excel file: {excel_path}")
        wb = openpyxl.load_workbook(excel_path)
        
        # Try to access the Variables sheet
        try:
            # First try by name
            sheet = wb['Variables']
        except KeyError:
            # If not found by name, try to access the last sheet
            sheet_names = wb.sheetnames
            if sheet_names:
                print(f"Variables sheet not found by name. Using last sheet: {sheet_names[-1]}")
                sheet = wb[sheet_names[-1]]
            else:
                print("No sheets found in the workbook")
                return False
        
        # Map config variables to Excel cells based on the provided mapping
        cell_mappings = {
            'B2': user_data.get('store_name', ''),
            'B31': user_data.get('starting_date', ''),
            'B32': user_data.get('duration', 36),
            'B37': user_data.get('open_days_per_month', 0),
            
            # Convenience store type
            'H37': user_data.get('convenience_store_type', {}).get('stores_number', 0),
            'C37': user_data.get('convenience_store_type', {}).get('monthly_transactions', 0),
            # Convert boolean to 1/0 for has_digital_screens
            'I37': 1 if user_data.get('convenience_store_type', {}).get('has_digital_screens', False) else 0,
            'J37': user_data.get('convenience_store_type', {}).get('screen_count', 0),
            'K37': user_data.get('convenience_store_type', {}).get('screen_percentage', 0),
            # Convert boolean to 1/0 for has_in_store_radio
            'M37': 1 if user_data.get('convenience_store_type', {}).get('has_in_store_radio', False) else 0,
            'N37': user_data.get('convenience_store_type', {}).get('radio_percentage', 0),
            
            # Supermarket store type
            'H38': user_data.get('supermarket_store_type', {}).get('stores_number', 0),
            'C38': user_data.get('supermarket_store_type', {}).get('monthly_transactions', 0),
            # Convert boolean to 1/0 for has_digital_screens
            'I38': 1 if user_data.get('supermarket_store_type', {}).get('has_digital_screens', False) else 0,
            'J38': user_data.get('supermarket_store_type', {}).get('screen_count', 0),
            'K38': user_data.get('supermarket_store_type', {}).get('screen_percentage', 0),
            # Convert boolean to 1/0 for has_in_store_radio
            'M38': 1 if user_data.get('supermarket_store_type', {}).get('has_in_store_radio', False) else 0,
            'N38': user_data.get('supermarket_store_type', {}).get('radio_percentage', 0),
            
            # Hypermarket store type
            'H39': user_data.get('hypermarket_store_type', {}).get('stores_number', 0),
            'C39': user_data.get('hypermarket_store_type', {}).get('monthly_transactions', 0),
            # Convert boolean to 1/0 for has_digital_screens
            'I39': 1 if user_data.get('hypermarket_store_type', {}).get('has_digital_screens', False) else 0,
            'J39': user_data.get('hypermarket_store_type', {}).get('screen_count', 0),
            'K39': user_data.get('hypermarket_store_type', {}).get('screen_percentage', 0),
            # Convert boolean to 1/0 for has_in_store_radio
            'M39': 1 if user_data.get('hypermarket_store_type', {}).get('has_in_store_radio', False) else 0,
            'N39': user_data.get('hypermarket_store_type', {}).get('radio_percentage', 0),
            
            # On-site channels
            'B43': user_data.get('website_visitors', 0),
            'B44': user_data.get('app_users', 0),
            'B45': user_data.get('loyalty_users', 0),
            
            # Off-site channels
            'B49': user_data.get('facebook_followers', 0),
            'B50': user_data.get('instagram_followers', 0),
            'B51': user_data.get('google_views', 0),
            'B52': user_data.get('email_subscribers', 0),
            'B53': user_data.get('sms_users', 0),
            'B54': user_data.get('whatsapp_contacts', 0)
        }
        
        # Update the cells
        for cell_ref, value in cell_mappings.items():
            try:
                sheet[cell_ref] = value
                print(f"Updated {cell_ref} with value: {value}")
            except Exception as e:
                print(f"Error updating cell {cell_ref}: {e}")
        
        # Update sheet names - replace {store_name} with actual store name
        store_name = user_data.get('store_name', '')
        if store_name:
            # Dictionary to store old sheet name to new sheet name mappings
            sheet_name_mapping = {}
            
            # Make a copy of the sheet names to avoid modifying during iteration
            sheet_names = wb.sheetnames.copy()
            for sheet_name in sheet_names:
                if '{store_name}' in sheet_name:
                    new_sheet_name = sheet_name.replace('{store_name}', store_name)
                    # Get the sheet by its old name
                    sheet = wb[sheet_name]
                    # Set the new title
                    sheet.title = new_sheet_name
                    # Store the mapping
                    sheet_name_mapping[sheet_name] = new_sheet_name
                    print(f"Renamed sheet '{sheet_name}' to '{new_sheet_name}'")
            
            # Use direct XML modification to replace all instances of {store_name}
            print("Using direct XML modification to update all formulas...")
            modified_file = update_excel_with_direct_xml(excel_path, store_name)
            
            if modified_file and os.path.exists(modified_file):
                # Use the modified file instead of the original
                print(f"Using modified file: {modified_file}")
                # Copy the modified file back to the original location
                import shutil
                shutil.copy2(modified_file, excel_path)
                # Remove the modified file
                os.remove(modified_file)
                # Reload the workbook to get the changes
                wb = openpyxl.load_workbook(excel_path)
        
        # Save the workbook
        wb.save(excel_path)
        print(f"Excel file updated successfully: {excel_path}")
        return True
        
    except Exception as e:
        print(f"Error updating Excel file: {e}")
        return False

def update_formula_references(workbook, sheet_name_mapping):
    """
    Update formula references in the Graphics sheet to point to the renamed sheets
    
    Args:
        workbook: The openpyxl workbook object
        sheet_name_mapping: Dictionary mapping old sheet names to new sheet names
    """
    try:
        # Get the Graphics sheet
        graphics_sheet = workbook['Graphics']
        print("Updating formula references in Graphics sheet...")
        
        # Track the number of updates
        updates_count = 0
        total_formulas = 0
        
        # Create a dictionary to track which cells have been processed
        processed_cells = {}
        
        # First pass: Identify all cells with formulas that reference the forecast sheets
        formula_cells = []
        
        # Iterate through all cells in the Graphics sheet
        for row in graphics_sheet.iter_rows():
            for cell in row:
                # Check if the cell contains a formula
                if cell.data_type == 'f' and cell.value and isinstance(cell.value, str):
                    total_formulas += 1
                    original_formula = cell.value
                    
                    # Check if the formula references any of the renamed sheets
                    references_renamed_sheet = False
                    for old_name in sheet_name_mapping.keys():
                        if old_name in original_formula:
                            references_renamed_sheet = True
                            break
                    
                    if references_renamed_sheet:
                        formula_cells.append((cell, original_formula))
        
        print(f"Found {len(formula_cells)} cells with formulas referencing renamed sheets (out of {total_formulas} total formulas)")
        
        # Second pass: Update the formulas
        for cell, original_formula in formula_cells:
            updated_formula = original_formula
            formula_updated = False
            
            # Check if the formula references any of the renamed sheets
            for old_name, new_name in sheet_name_mapping.items():
                # Pattern 1: Sheet name in single quotes: 'Sheet Name'!
                if f"'{old_name}'!" in updated_formula:
                    updated_formula = updated_formula.replace(f"'{old_name}'!", f"'{new_name}'!")
                    formula_updated = True
                
                # Pattern 2: Sheet name without quotes: SheetName!
                # We need to be careful with this pattern to avoid partial matches
                pattern = re.compile(f"(^|[^'])({re.escape(old_name)})!")
                if pattern.search(updated_formula):
                    updated_formula = pattern.sub(f"\\1{new_name}!", updated_formula)
                    formula_updated = True
                
                # Pattern 3: Sheet name in INDIRECT function: INDIRECT("'Sheet Name'!...")
                if f"INDIRECT(\"'{old_name}'!" in updated_formula:
                    updated_formula = updated_formula.replace(
                        f"INDIRECT(\"'{old_name}'!", 
                        f"INDIRECT(\"'{new_name}'!"
                    )
                    formula_updated = True
                
                # Pattern 4: Sheet name in double quotes: "Sheet Name"!
                if f"\"{old_name}\"!" in updated_formula:
                    updated_formula = updated_formula.replace(f"\"{old_name}\"!", f"\"{new_name}\"!")
                    formula_updated = True
                
                # Pattern 5: Simple text replacement for any remaining instances
                # Only do this if we're sure it's a sheet reference
                if old_name in updated_formula and not formula_updated:
                    updated_formula = updated_formula.replace(old_name, new_name)
                    formula_updated = True
            
            # If the formula was updated, set it back to the cell
            if formula_updated:
                try:
                    cell.value = updated_formula
                    updates_count += 1
                    print(f"Updated formula in cell {cell.coordinate}: {original_formula} -> {updated_formula}")
                except Exception as e:
                    print(f"Error updating formula in cell {cell.coordinate}: {e}")
        
        # Special handling for specific cells that might have been missed
        # These are known cells with complex formulas referencing the forecast sheets
        special_cells = ['C4', 'H4', 'O4', 'U4', 'AA4', 'AG4', 
                         'C5', 'H5', 'O5', 'U5', 'AA5', 'AG5',
                         'C6', 'H6', 'O6', 'U6', 'AA6', 'AG6',
                         'C25', 'H25', 'O25', 'U25', 'AA25', 'AG25',
                         'C26', 'H26', 'O26', 'U26', 'AA26', 'AG26',
                         'C27', 'H27', 'O27', 'U27', 'AA27', 'AG27']
        
        # Print all cells in the Graphics sheet that have formulas
        print("\nDiagnostic: Checking all cells with formulas in Graphics sheet:")
        formula_cells_diagnostic = []
        
        # Special check for rows 25-27
        print("\nDiagnostic: Checking cells in rows 25-27:")
        
        # Define columns to check, focusing on G through AG
        columns_to_check = ['C', 'G', 'H', 'I', 'J', 'K', 'L', 'M', 'N', 'O', 'P', 'Q', 'R', 'S', 'T', 'U', 'V', 'W', 'X', 'Y', 'Z', 
                           'AA', 'AB', 'AC', 'AD', 'AE', 'AF', 'AG']
        
        # Cells with ArrayFormula objects that need special handling
        array_formula_cells = []
        
        # Check each cell in rows 25-27 for all specified columns
        for row_num in range(25, 28):  # rows 25, 26, 27
            for col in columns_to_check:
                cell_coord = f"{col}{row_num}"
                try:
                    cell = graphics_sheet[cell_coord]
                    # Print cell information regardless of whether it has a formula
                    if cell.value is not None:
                        value_type = cell.data_type
                        value_preview = str(cell.value)[:50] if cell.value else "None"
                        print(f"Row {row_num}, Cell {cell_coord}: Type={value_type}, Value={value_preview}...")
                        
                        # Check if it's a formula that references our sheets
                        if cell.data_type == 'f':
                            if isinstance(cell.value, str) and any(old_name in cell.value for old_name in sheet_name_mapping.keys()):
                                formula_cells_diagnostic.append(f"{cell_coord}: {cell.value[:50]}...")
                                # Add this cell to our special handling
                                if cell_coord not in special_cells:
                                    special_cells.append(cell_coord)
                            # Check for ArrayFormula objects
                            elif str(cell.value).startswith("<openpyxl.worksheet.formula.ArrayFormula"):
                                # These are the cells we need to handle specially
                                array_formula_cells.append(cell_coord)
                                if cell_coord not in special_cells:
                                    special_cells.append(cell_coord)
                except Exception as e:
                    print(f"Error accessing cell {cell_coord}: {e}")
                    pass
        
        # Special handling for ArrayFormula cells
        if array_formula_cells:
            print(f"\nFound {len(array_formula_cells)} cells with ArrayFormula objects that need special handling:")
            for cell_coord in array_formula_cells:
                print(f"  {cell_coord}")
            
            # Handle array formula cells by directly modifying the worksheet's _cells dictionary
            # This is a workaround since openpyxl doesn't provide a direct API for modifying ArrayFormula objects
            print("\nAttempting to update array formulas by directly accessing the worksheet's XML...")
            
            # Get the worksheet's underlying XML
            try:
                # First, save the workbook to a temporary file
                temp_dir = os.path.dirname(excel_path) if excel_path else os.getcwd()
                temp_file = os.path.join(temp_dir, "_temp_for_xml_edit.xlsx")
                wb.save(temp_file)
                
                # Then reload it with a different library that can modify the XML directly
                from zipfile import ZipFile
                import xml.etree.ElementTree as ET
                
                # Define XML namespaces
                namespaces = {
                    'main': 'http://schemas.openxmlformats.org/spreadsheetml/2006/main'
                }
                
                # Extract and modify the sheet XML
                with ZipFile(temp_file, 'r') as zip_ref:
                    # Find the sheet XML file for the Graphics sheet
                    sheet_id = None
                    for i, sheet in enumerate(wb.worksheets):
                        if sheet.title == 'Graphics':
                            sheet_id = i + 1
                            break
                    
                    if sheet_id:
                        sheet_xml_path = f'xl/worksheets/sheet{sheet_id}.xml'
                        sheet_xml = zip_ref.read(sheet_xml_path).decode('utf-8')
                        
                        # Parse the XML
                        root = ET.fromstring(sheet_xml)
                        
                        # Find all formula cells and replace text
                        formula_count = 0
                        for formula_elem in root.findall('.//main:f', namespaces):
                            formula_text = formula_elem.text
                            if formula_text and any(old_name in formula_text for old_name in sheet_name_mapping.keys()):
                                # Replace the sheet names in the formula text
                                for old_name, new_name in sheet_name_mapping.items():
                                    formula_text = formula_text.replace(old_name, new_name)
                                formula_elem.text = formula_text
                                formula_count += 1
                        
                        print(f"Updated {formula_count} formulas in the XML")
                        
                        # Write the modified XML back
                        # Note: This would require rebuilding the XLSX file, which is complex
                        # For now, we'll just note that we need to handle these cells
                        print("XML modification would be required to update array formulas")
                        
                # Clean up the temporary file
                if os.path.exists(temp_file):
                    os.remove(temp_file)
                    
            except Exception as e:
                print(f"Error attempting to modify XML: {e}")
                
            print("\nWarning: Array formulas in cells {', '.join(array_formula_cells)} may still reference old sheet names.")
            print("These will need to be updated manually in Excel or with a more specialized approach.")
            
            # Add a warning message
            print("\nIMPORTANT: The Excel file may still show recovery errors due to array formulas.")
            print("To fix this, open the file in Excel, accept the recovery, and save it.")
            print("Or consider using a different library like xlwings that can directly manipulate array formulas.")
        
        # Check all cells for formulas referencing the renamed sheets
        for row in graphics_sheet.iter_rows():
            for cell in row:
                if cell.data_type == 'f' and cell.value and isinstance(cell.value, str):
                    formula = cell.value
                    if any(old_name in formula for old_name in sheet_name_mapping.keys()):
                        if not any(cell.coordinate == f_cell.split(':')[0] for f_cell in formula_cells_diagnostic):
                            formula_cells_diagnostic.append(f"{cell.coordinate}: {formula[:50]}...")
        
        print(f"\nFound {len(formula_cells_diagnostic)} cells with formulas referencing forecast sheets:")
        for cell_info in formula_cells_diagnostic:
            print(f"  {cell_info}")
        
        for cell_coord in special_cells:
            if cell_coord not in processed_cells:
                try:
                    cell = graphics_sheet[cell_coord]
                    if cell.data_type == 'f' and cell.value and isinstance(cell.value, str):
                        original_formula = cell.value
                        updated_formula = original_formula
                        formula_updated = False
                        
                        for old_name, new_name in sheet_name_mapping.items():
                            if old_name in updated_formula:
                                updated_formula = updated_formula.replace(old_name, new_name)
                                formula_updated = True
                        
                        if formula_updated:
                            cell.value = updated_formula
                            updates_count += 1
                            print(f"Special handling: Updated formula in cell {cell_coord}: {original_formula} -> {updated_formula}")
                except Exception as e:
                    print(f"Error in special handling for cell {cell_coord}: {e}")
        
        print(f"Updated {updates_count} formula references in Graphics sheet")
        
        # Check if there are any remaining formulas with the old sheet names
        remaining_formulas = []
        for row in graphics_sheet.iter_rows():
            for cell in row:
                if cell.data_type == 'f' and cell.value and isinstance(cell.value, str):
                    formula = cell.value
                    if any(old_name in formula for old_name in sheet_name_mapping.keys()):
                        remaining_formulas.append(f"{cell.coordinate}: {formula[:50]}...")
        
        if remaining_formulas:
            print(f"Warning: Found {len(remaining_formulas)} formulas that still reference old sheet names:")
            for formula in remaining_formulas:
                print(f"  {formula}")
        else:
            print("All formulas have been successfully updated!")
        
        return updates_count
    except KeyError:
        print("Graphics sheet not found in workbook")
        return 0
    except Exception as e:
        print(f"Error updating formula references: {e}")
        return 0

if __name__ == "__main__":
    # For testing purposes
    import sys
    if len(sys.argv) > 1:
        excel_path = sys.argv[1]
        update_excel_variables(excel_path)
    else:
        print("Please provide the path to the Excel file as an argument")