import json
import boto3
import psycopg
from botocore.exceptions import ClientError
import os


def get_db_credentials(secret_name):
    secrets_client = boto3.client("secretsmanager")
    try:
        response = secrets_client.get_secret_value(SecretId=secret_name)
        secret = json.loads(response["SecretString"])
        return secret
    except ClientError as e:
        print(f"Error retrieving secret: {e}")
        raise e


def handler(event, context):
    """Updates all existing entries in bots_onetimefreeusage to have free mobile v4 gen limit of 20"""
    env = os.environ["env"]
    secret_name = f"studio-api-{env}-envs"
    credentials = get_db_credentials(secret_name)
    try:
        with psycopg.connect(credentials["DATABASE_URL"]) as connection:
            with connection.cursor() as cursor:
                batch_size = 10000
                last_user_id = 0
                total_updated = 0
                while True:
                    # First get the next batch of user_ids
                    cursor.execute(
                        """
                        SELECT user_id 
                        FROM bots_onetimefreeusage 
                        WHERE user_id > %s 
                        ORDER BY user_id 
                        LIMIT %s
                    """,
                        (last_user_id, batch_size),
                    )

                    user_ids = cursor.fetchall()
                    if not user_ids:
                        break  # No more records to process
                    # Create the update query with the exact user_ids from this batch
                    user_id_list = [uid[0] for uid in user_ids]
                    placeholders = ",".join(["%s"] * len(user_id_list))

                    cursor.execute(
                        f"""
                        UPDATE bots_onetimefreeusage 
                        SET free_mobile_v4_gens_limit = 20 
                        WHERE user_id IN ({placeholders})
                    """,
                        user_id_list,
                    )
                    rows_affected = cursor.rowcount
                    total_updated += rows_affected
                    last_user_id = user_ids[-1][
                        0
                    ]  # Get the last user_id from this batch
                    print(
                        f"Batch completed: Updated {rows_affected} rows. Last user_id: {last_user_id}"
                    )
                    connection.commit()  # Commit after each batch
                print(f"Total rows updated: {total_updated}")
                return {
                    "statusCode": 200,
                    "body": json.dumps(f"Successfully updated {total_updated} rows"),
                }
    except Exception as e:
        print(f"Error: {e}")
        return {"statusCode": 500, "body": json.dumps(f"Error: {str(e)}")}