import os
import pickle
from google.oauth2 import service_account
from googleapiclient.discovery import build

SHEET_ID = ''
RANGE_NAME = 'AL'
KEY_PATH = '/nfs/data/$user/projects/crash/diffusion/diffusers/examples/textual_inversion/files/autoexplogger-ccfefb7f10a9.json'

class AutoExpLogger():
    def __init__(self, sheet_id=SHEET_ID, range_name='Attack', key_path=KEY_PATH):
        ## TODO under dev ## write them in files for safety

        # The ID of the Google Sheet you want to access
        self.sheet_id = sheet_id

        # The range of cells you want to modify (e.g., 'Sheet1!A1:B2')
        self.range_name = f'{RANGE_NAME}-{range_name}'

        # The path to your service account JSON key file
        self.key_path = key_path

        ## build service
        # Load the service account credentials
        credentials = service_account.Credentials.from_service_account_file(key_path, scopes=['https://www.googleapis.com/auth/spreadsheets'])

        # Build the service for accessing the Google Sheets API
        self.service = build('sheets', 'v4', credentials=credentials)

    def write(self, prompt_id, exp_id, num_evals, avg_loss):  ## values: 2D array
        ## find next empty row id
        checker_range_name = f'{self.range_name}!C1:C1000000'
        result = self.service.spreadsheets().values().get(spreadsheetId=self.sheet_id, range=checker_range_name).execute()
        data = result['values']
        
        try:
            next_row_id = data.index([exp_id]) + 1
        except ValueError:
            next_row_id = len(data) + 1

        start_col = chr(ord('F') + prompt_id)
        res_range_name = f'{self.range_name}!{start_col}{next_row_id}:Z{next_row_id}'
        sts_range_name = f'{self.range_name}!A{next_row_id}:Z{next_row_id}'
        cfg_range_name = f'{self.range_name}!C{next_row_id}:Z{next_row_id}'
        bgt_range_name = f'{self.range_name}!E{next_row_id}:Z{next_row_id}'
        print(f'Write to row {next_row_id}')
        
        ## write status
        body = {
            'values': [['done']]
        }
        result = self.service.spreadsheets().values().update(spreadsheetId=self.sheet_id,
                                                             range=sts_range_name,
                                                             valueInputOption='RAW', body=body).execute()


        ## write exp congfig
        body = {
            'values': [[exp_id]]
        }
        result = self.service.spreadsheets().values().update(spreadsheetId=self.sheet_id,
                                                             range=cfg_range_name,
                                                             valueInputOption='RAW', body=body).execute()

        ## write budget
        body = {
            'values': [[num_evals]]
        }
        result = self.service.spreadsheets().values().update(spreadsheetId=self.sheet_id,
                                                             range=bgt_range_name,
                                                             valueInputOption='RAW', body=body).execute()

        ## write numerical results
        body = {
            'values': [[avg_loss]]
        }
        result = self.service.spreadsheets().values().update(spreadsheetId=self.sheet_id,
                                                             range=res_range_name,
                                                             valueInputOption='RAW', body=body).execute()
        print('Cells updated: %d' % result.get('updatedCells'))
        return