#!/usr/bin/env python3
"""
Script to add email column to all CSV rate files
"""

import csv
import os

# Company email mapping
COMPANY_EMAILS = {
    'MSC': 'bookings@msc.com',
    'SSC Consolidation': 'bookings@sscconsolidation.com',
    'ECU Worldwide': 'bookings@ecuworldwide.com',
    'TCC Logistics': 'bookings@tcclogistics.com',
    'NVO Consolidation': 'bookings@nvoconsolidation.com',
    'Test Maritime Express': 'bookings@testmaritime.com'
}

csv_dir = 'apps/backend/src/infrastructure/storage/csv-storage/rates'

# Process each CSV file
for filename in os.listdir(csv_dir):
    if not filename.endswith('.csv'):
        continue

    filepath = os.path.join(csv_dir, filename)
    print(f'Processing {filename}...')

    # Read existing data
    rows = []
    with open(filepath, 'r', encoding='utf-8') as f:
        reader = csv.DictReader(f)
        fieldnames = reader.fieldnames

        # Check if email column already exists
        if 'companyEmail' in fieldnames:
            print(f'  - Email column already exists, skipping')
            continue

        # Add email column header
        new_fieldnames = list(fieldnames)
        # Insert email after companyName
        company_name_index = new_fieldnames.index('companyName')
        new_fieldnames.insert(company_name_index + 1, 'companyEmail')

        # Read all rows and add email
        for row in reader:
            company_name = row['companyName']
            company_email = COMPANY_EMAILS.get(company_name, f'bookings@{company_name.lower().replace(" ", "")}.com')
            row['companyEmail'] = company_email
            rows.append(row)

    # Write back with new column
    with open(filepath, 'w', encoding='utf-8', newline='') as f:
        writer = csv.DictWriter(f, fieldnames=new_fieldnames)
        writer.writeheader()
        writer.writerows(rows)

    print(f'  - Added companyEmail column ({len(rows)} rows updated)')

print('\nDone! All CSV files updated.')
