import os
import json
from datetime import datetime
from supabase import create_client, Client
import pandas as pd
from io import StringIO
import logging
import boto3

logging.basicConfig(level=logging.INFO)

# Initialize Lambda client
LAMBDA_CLIENT = boto3.client('lambda')
CLEAR_SJREPORT_OLD_DATA_ARN = os.environ['CLEAR_SJREPORT_OLD_DATA_ARN']

def get_parameter(name: str) -> str:
    ssm = boto3.client('ssm')
    try:
        response = ssm.get_parameter(Name=name, WithDecryption=True)
        return response['Parameter']['Value']
    except Exception as e:
        logging.error(f"Error getting parameter {name}: {str(e)}")
        raise

def get_supabase_client() -> Client:
    url: str = get_parameter('/supabase/url')
    key: str = get_parameter('/supabase/anon')
    supabase: Client = create_client(url, key)
    supabase.postgrest.rpc("raw_sql", {"query": "SET statement_timeout = '0';"})
    return supabase

def lambda_handler(event, context):
    try:
        supabase = get_supabase_client()
        batch_size = 10000
        offset = 0
        all_data = []
        while True:
            response = supabase.rpc('get_old_sjreport_rows_paged', {'p_limit': batch_size, 'p_offset': offset}).execute()
            data = response.data
            if not data:
                break
            all_data.extend(data)
            offset += batch_size
            logging.info(f"Fetched {len(data)} records, total so far: {len(all_data)}")
        if not all_data:
            logging.info("No data to backup - all records have the maximum week_start_date")
            return {
                'statusCode': 200,
                'body': json.dumps({
                    'message': 'No data to backup - all records have the maximum week_start_date',
                    'records_backed_up': 0
                })
            }
        df = pd.DataFrame(all_data)
        csv_buffer = StringIO()
        df.to_csv(csv_buffer, index=False)
        bucket_name = "uploads/backups"
        timestamp = datetime.now().strftime('%Y%m%d_%H%M%S')
        file_name = f"sjreport_old_dates_backup_{timestamp}.csv"
        supabase.storage.from_(bucket_name).upload(
            file_name,
            csv_buffer.getvalue().encode('utf-8')
        )
        record_count = len(all_data)
        
        # Invoke sjreport_clear_old_data Lambda function
        payload = json.dumps(event)
        response = LAMBDA_CLIENT.invoke(
            FunctionName=CLEAR_SJREPORT_OLD_DATA_ARN,
            InvocationType='Event',  # Asynchronous invocation
            Payload=payload
        )
        
        return {
            'statusCode': 200,
            'body': json.dumps({
                'message': 'Backup completed successfully and clear old data triggered',
                'records_backed_up': record_count,
                'file_name': file_name,
                'clear_old_data_triggered': True
            })
        }
    except Exception as e:
        logging.error(f"Error in backup Lambda: {str(e)}")
        return {
            'statusCode': 500,
            'body': json.dumps({
                'error': str(e)
            })
        }