"""
Insert one batch of task records from S3 into Supabase task table.
Invoked by Step Functions Create Tasks workflow; each invocation handles one S3 object.
"""
import os
import json
import boto3
from supabase import create_client, Client

# SSM for Supabase config
ssm = boto3.client('ssm')
s3 = boto3.client('s3')

def get_parameter(name: str) -> str:
    """Get a parameter from AWS SSM Parameter Store."""
    response = ssm.get_parameter(Name=name, WithDecryption=True)
    return response['Parameter']['Value']

def lambda_handler(event, context):
    """
    Event: { "s3_bucket": str, "s3_key": str, "week_start_date": str }
    Reads JSON array from S3 and upserts into Supabase task table.
    """
    bucket = event['s3_bucket']
    key = event['s3_key']

    # Get batch from S3
    try:
        resp = s3.get_object(Bucket=bucket, Key=key)
        batch = json.loads(resp['Body'].read().decode('utf-8'))
    except Exception as e:
        return {'statusCode': 500, 'error': f'S3 read failed: {str(e)}', 's3_key': key}

    if not batch:
        return {'statusCode': 200, 'inserted': 0, 's3_key': key}

    url = get_parameter('/supabase/url')
    key_supabase = get_parameter('/supabase/anon')
    supabase: Client = create_client(url, key_supabase)
    # Allow long-running statement and wait for locks (55P03 = lock timeout)
    supabase.postgrest.rpc("raw_sql", {"query": "SET statement_timeout = '20min'; SET lock_timeout = '20min';"})

    # Upsert in chunks to hold locks for less time and reduce lock contention
    chunk_size = 500
    try:
        for i in range(0, len(batch), chunk_size):
            chunk = batch[i : i + chunk_size]
            supabase.table('task').upsert(
                chunk,
                on_conflict='chain,store_name,product_id,week_start_date'
            ).execute()
    except Exception as e:
        return {'statusCode': 500, 'error': str(e), 's3_key': key, 'batch_size': len(batch)}

    return {'statusCode': 200, 'inserted': len(batch), 's3_key': key}