import sys
import requests
import pandas as pd
import json
import re
import argparse

def get_excel_download_url(google_sheet_url):
    """
    Convert a Google Sheets URL to its Excel export URL.
    """
    # Extract the Sheet ID to construct the download URL
    match = re.search(r'/d/([a-zA-Z0-9-_]+)', google_sheet_url)
    if not match:
        raise ValueError('Invalid Google Sheets URL')
    sheet_id = match.group(1)
    # Export the entire spreadsheet as an Excel file (includes all sheets)
    return f'https://docs.google.com/spreadsheets/d/{sheet_id}/export?format=xlsx&id={sheet_id}'

def download_excel(url, output_path):
    """
    Download the Excel file from the given URL.
    """
    response = requests.get(url)
    if response.status_code != 200:
        raise Exception(f'Failed to download file: {response.status_code}')
    with open(output_path, 'wb') as f:
        f.write(response.content)

def convert_excel_to_json(excel_path, json_path, sheet_name):
    """
    Read the Excel file, extract columns from the specified sheet, and save as JSON.
    """
    df = pd.read_excel(excel_path, sheet_name=sheet_name)
    # Ensure columns exist
    required_columns = [
        'Placeholder',
        'Laconic Address',
        'Total LPS Allocation',
        'Lock (months)',
        'Vest (months)'
    ]
    for col in required_columns:
        if col not in df.columns:
            raise Exception(f'Missing required column: {col}')

    result = {}
    for _, row in df.iterrows():
        placeholder = str(row['Placeholder']) if not pd.isna(row['Placeholder']) else ''

        laconic_address = str(row['Laconic Address']) if not pd.isna(row['Laconic Address']) else ''

        # Use laconic_address as key if placeholder is missing or empty
        key = placeholder if placeholder and placeholder.lower() != 'nan' else laconic_address

        # Skip the row if both 'Placeholder' and 'Laconic Address' are missing or invalid
        if not key or key.lower() == 'nan':
            continue

        entry = {
            'total_lps_allocation': row['Total LPS Allocation'] if not pd.isna(row['Total LPS Allocation']) else None,
            'lock_months': row['Lock (months)'] if not pd.isna(row['Lock (months)']) else None,
            'vest_months': row['Vest (months)'] if not pd.isna(row['Vest (months)']) else None,
            'laconic_address': row['Laconic Address'] if not pd.isna(row['Laconic Address']) else None
        }

        result[key] = entry

    with open(json_path, 'w') as f:
        json.dump(result, f, indent=2)

def main():
    parser = argparse.ArgumentParser(description='Generate LPS distribution JSON from Excel or Google Sheet')
    parser.add_argument('--input', '-i', required=True, help='Input: Google Sheet URL or local Excel file path')
    parser.add_argument('--output', '-o', default='distribution.json', help='Output JSON file path (default: distribution.json)')
    parser.add_argument('--sheet', '-s', default='Genesis Allocation', help='Sheet name to read (default: Genesis Allocation)')
    
    args = parser.parse_args()
    
    if args.input.startswith('https://'):
        excel_url = get_excel_download_url(args.input)
        excel_path = 'sheet.xlsx'
        print(f'Downloading Excel file from: {excel_url}')
        download_excel(excel_url, excel_path)
    else:
        excel_path = args.input
        print(f'Using Excel file at path: {excel_path}')
    
    print(f'Converting Excel to JSON (sheet: {args.sheet})...')
    convert_excel_to_json(excel_path, args.output, args.sheet)
    print(f'JSON saved to {args.output}')

if __name__ == '__main__':
    main()