import openpyxl

def update_excel_ampacity(file_path, headers_row_idx, ampacity_cols_identifiers, target_cross_section="1x1200/35"):
    print(f"Updating {file_path}...")
    wb = openpyxl.load_workbook(file_path)
    ws = wb.active
    
    # openpyxl is 1-indexed for rows and columns
    headers = [cell.value for cell in ws[headers_row_idx]]
    
    # Identify column indices for ampacity (0-indexed locally for easier row access)
    col_indices = []
    for identifier in ampacity_cols_identifiers:
        if isinstance(identifier, int):
            col_indices.append(identifier)
        else:
            try:
                # list.index returns 0-indexed position
                col_indices.append(headers.index(identifier))
            except ValueError:
                print(f"Warning: Could not find column '{identifier}' in {file_path}")

    # Find row index for "Number of cores and cross-section" or use index 8
    cs_col_idx = 8
    try:
        cs_col_idx = headers.index("Number of cores and cross-section")
    except ValueError:
        pass

    rows_updated = 0
    # ws.iter_rows returns 1-indexed rows
    for row in ws.iter_rows(min_row=headers_row_idx + 1):
        # row is a tuple of cells, so row[cs_col_idx] is 0-indexed access to the tuple
        if str(row[cs_col_idx].value).strip() == target_cross_section:
            for col_idx in col_indices:
                row[col_idx].value = "On Request"
            rows_updated += 1
            
    wb.save(file_path)
    print(f"Updated {rows_updated} rows in {file_path}")

# File 1: medium-voltage-KM.xlsx
update_excel_ampacity(
    'data/excel/medium-voltage-KM.xlsx',
    1, # Headers are in first row (1-indexed)
    [
        'Current ratings in air, trefoil*',
        'Current ratings in air, flat*',
        'Current ratings in ground, trefoil*',
        'Current ratings in ground, flat*'
    ]
)

# File 2: medium-voltage-KM 170126.xlsx
update_excel_ampacity(
    'data/excel/medium-voltage-KM 170126.xlsx',
    1, # Indices 39 and 41 were from a 0-indexed JSON representation
    [39, 41] 
)