import json
from supabase import create_client, Client
import logging
import boto3

logging.basicConfig(level=logging.INFO)

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 clear_old_task_records(supabase: Client):
    """Call the PostgreSQL function to clear old task records"""
    response = supabase.rpc("clear_old_task_records").execute()
    return response.data

def lambda_handler(event, context):
    try:
        supabase = get_supabase_client()
        
        # Call the PostgreSQL function to clear old records
        results = clear_old_task_records(supabase)
        
        if not results:
            logging.info("No store records found")
            return {
                'statusCode': 200,
                'body': json.dumps({
                    'message': 'No store records found',
                    'stores_processed': 0
                })
            }
        
        # Process results from the function
        total_deleted = 0
        processed_stores = 0
        deleted_stores = 0
        
        for result in results:
            store_code = result['store_code']
            max_date = result['max_week_start_date']
            deleted_count = result['deleted_count']
            status = result['status']
            
            processed_stores += 1
            total_deleted += deleted_count
            
            if status == 'DELETED':
                deleted_stores += 1
                logging.info(f"Processed store: {store_code}, max date: {max_date}, deleted records: {deleted_count}")
            else:
                logging.info(f"Processed store: {store_code}, max date: {max_date}, no records to delete")
        
        return {
            'statusCode': 200,
            'body': json.dumps({
                'message': 'Deletion process completed',
                'stores_processed': processed_stores,
                'stores_with_deletions': deleted_stores,
                'total_records_deleted': total_deleted,
                'results': results
            })
        }
        
    except Exception as e:
        logging.error(f"Error in deletion Lambda: {str(e)}")
        return {
            'statusCode': 500,
            'body': json.dumps({
                'error': str(e)
            })
        }