import os
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 get_sjreport_batch_ids(supabase: Client, week_start_date: str, batch_size: int = 1000):
    response = supabase.rpc("get_sjreport_update_batch_ids", {
        "p_week_start_date": week_start_date,
        "p_batch_size": batch_size
    }).execute()
    if hasattr(response, 'data') and response.data:
        return response.data
    return []

def update_sjreport_batch(supabase: Client, week_start_date: str, batch_ids: list):
    if not batch_ids:
        return 0
    response = supabase.rpc("update_sjreport_batch_by_ids", {
        "p_week_start_date": week_start_date,
        "p_ids": batch_ids
    }).execute()
    if hasattr(response, 'data') and response.data:
        return response.data
    return 0

def batch_update_sjreport_with_product_ids(supabase: Client, week_start_date: str, batch_size: int = 1000):
    total_updated = 0
    while True:
        batch_ids = get_sjreport_batch_ids(supabase, week_start_date, batch_size)
        if not batch_ids:
            break
        updated = update_sjreport_batch(supabase, week_start_date, batch_ids)
        total_updated += updated if isinstance(updated, int) else 0
        logging.info(f"Updated {updated} records in this batch, total updated: {total_updated}")
        if not batch_ids or (isinstance(updated, int) and updated < batch_size):
            break
    logging.info(f"Batch update complete. Total records updated: {total_updated}")
    return total_updated

def invoke_create_tasks(lambda_client: boto3.client, week_start_date: str) -> None:
    """Invoke CreateTasksBulk (S3 + Step Function path for task creation)."""
    function_arn = os.environ.get('CREATE_TASKS_BULK_ARN')
    if not function_arn:
        raise ValueError('CREATE_TASKS_BULK_ARN environment variable is not set')
    try:
        payload = {"week_start_date": week_start_date}
        lambda_client.invoke(
            FunctionName=function_arn,
            InvocationType='Event',
            Payload=json.dumps(payload)
        )
    except Exception as e:
        logging.error(f"Error invoking CreateTasksBulk: {str(e)}")
        raise

def update_upload_status(supabase: Client, record_id: str, error: str = '') -> None:
    """Update the upload record status with error information"""
    try:
        update_data = {
            "error": error,
            "processed": True
        }
        supabase.table("sjreport_uploads").update(update_data).eq("id", record_id).execute()
    except Exception as e:
        logging.error(f"Error updating upload status: {str(e)}")
        raise

def lambda_handler(event, context):
    lambda_client = boto3.client('lambda')
    week_start_date = event.get('week_start_date')
    record_id = event.get('record_id')
    batch_size = 500
    
    if not week_start_date or not record_id:
        return {
            'statusCode': 400,
            'body': json.dumps({'error': 'week_start_date and record_id are required'})
        }
    
    try:
        supabase = get_supabase_client()
        total_updated = 0; # remove temp
        # total_updated = batch_update_sjreport_with_product_ids(supabase, week_start_date, batch_size)
        invoke_create_tasks(lambda_client, week_start_date)
        
        # Update status on success
        update_upload_status(supabase, record_id)
        
        return {
            'statusCode': 200,
            'body': json.dumps({'message': f'Batch update complete. Total records updated: {total_updated}'})
        }
    except Exception as e:
        error_message = f"Error in batch update Lambda: {str(e)}"
        logging.error(error_message)
        
        # Update status with error
        try:
            update_upload_status(supabase, record_id, error_message)
        except Exception as update_error:
            logging.error(f"Failed to update error status: {str(update_error)}")
        
        return {
            'statusCode': 500,
            'body': json.dumps({'error': error_message})
        }