#!/usr/bin/env python3
import json
import subprocess

spreadsheet_id = '1uVd-xZFF4TEQGqtvw9z6W8fffeaifPCoLsek83GmEoQ'
account = 'jonny@jonnyshannon.com'

# Download current USA sheet (which has all companies)
print("Downloading USA sheet (master list)...")
result = subprocess.run([
    'gog', 'sheets', 'get',
    spreadsheet_id,
    'USA!A:Z',
    '--account', account,
    '--json'
], capture_output=True, text=True)

if result.returncode != 0:
    print(f"Error downloading: {result.stderr}")
    exit(1)

data = json.loads(result.stdout)
all_rows = data['values']
header = all_rows[0]

print(f"Loaded {len(all_rows)-1} companies from USA sheet")

# Define country mappings
country_sheets = {
    'USA': ['USA', 'US', 'UNITED STATES'],
    'South Africa': ['ZA', 'SOUTH AFRICA'],
    'Australia': ['AU', 'AUSTRALIA'],
    'New Zealand': ['NZ', 'NEW ZEALAND'],
    'Canada': ['CA', 'CANADA'],
    'UK': ['UK', 'GB', 'UNITED KINGDOM', 'GREAT BRITAIN']
}

# Filter companies for each country
for sheet_name, country_codes in country_sheets.items():
    # Filter rows
    filtered = [header]
    for row in all_rows[1:]:
        if len(row) > 2:
            country = (row[2] or '').strip().upper()
            if country in country_codes:
                filtered.append(row)
    
    print(f"\n{sheet_name}: {len(filtered)-1} companies")
    
    if len(filtered) > 1:
        # Clear sheet
        print(f"  Clearing {sheet_name}...")
        subprocess.run([
            'gog', 'sheets', 'clear',
            spreadsheet_id,
            f'{sheet_name}!A:Z',
            '--account', account
        ], capture_output=True)
        
        # Update sheet
        print(f"  Uploading data...")
        result = subprocess.run([
            'gog', 'sheets', 'update',
            spreadsheet_id,
            f'{sheet_name}!A1',
            '--values-json', json.dumps(filtered),
            '--input', 'USER_ENTERED',
            '--account', account
        ], capture_output=True, text=True)
        
        if result.returncode == 0:
            print(f"  ✅ {sheet_name} updated!")
        else:
            print(f"  ❌ Error: {result.stderr}")
    else:
        print(f"  ⚠️  No companies found")

# Handle "Other Countries" - everything that's NOT in the listed countries
print(f"\n--- Other Countries ---")
excluded_codes = []
for codes in country_sheets.values():
    excluded_codes.extend(codes)

other = [header]
for row in all_rows[1:]:
    if len(row) > 2:
        country = (row[2] or '').strip().upper()
        if country and country not in excluded_codes:
            other.append(row)

print(f"Other Countries: {len(other)-1} companies")

if len(other) > 1:
    print("  Clearing Other Countries...")
    subprocess.run([
        'gog', 'sheets', 'clear',
        spreadsheet_id,
        'Other Countries!A:Z',
        '--account', account
    ], capture_output=True)
    
    print("  Uploading data...")
    result = subprocess.run([
        'gog', 'sheets', 'update',
        spreadsheet_id,
        'Other Countries!A1',
        '--values-json', json.dumps(other),
        '--input', 'USER_ENTERED',
        '--account', account
    ], capture_output=True, text=True)
    
    if result.returncode == 0:
        print("  ✅ Other Countries updated!")
    else:
        print(f"  ❌ Error: {result.stderr}")

print(f"\n🔗 https://docs.google.com/spreadsheets/d/{spreadsheet_id}/edit")