import argparse
import boto3
from datetime import datetime
import json
import base64
import sys


def invoke_lambda(function_name, payload_dict, region="us-east-1"):
    print(f"Invoking Lambda: {function_name}")
    print(f"Payload: {payload_dict}")
    client = boto3.client("lambda", region_name=region)

    response = client.invoke(
        FunctionName=function_name,
        InvocationType="RequestResponse",
        Payload=json.dumps(payload_dict),
        LogType="Tail"
    )

    # Read the response payload
    response_payload = response['Payload'].read().decode('utf-8')
    print("\n=== Lambda Response ===")
    print(response_payload)

    # Decode and print logs (if available)
    log_result = response.get('LogResult')
    if log_result:
        decoded_logs = base64.b64decode(log_result).decode('utf-8')
        print("\n=== Lambda Logs ===")
        print(decoded_logs)
    else:
        print("No logs returned from Lambda.")

    if response.get("FunctionError"):
        print("Lambda function returned an error.")
        sys.exit(1)


def parse_args():
    parser = argparse.ArgumentParser(
        description="Invoke a Lambda function with a datetime payload"
    )
    parser.add_argument(
        "env",
        choices=["qa", "prod"],
        help="Target environment (qa or prod)"
    )
    parser.add_argument(
        "datetime",
        type=str,
        help="Datetime in format YYYY-MM-DD HH:MM:SS"
    )
    return parser.parse_args()


def validate_datetime(dt_str):
    try:
        return datetime.strptime(dt_str, "%Y-%m-%d %H:%M:%S")
    except ValueError:
        print("Invalid datetime format. Expected 'YYYY-MM-DD HH:MM:SS'")
        sys.exit(1)
        

if __name__ == "__main__":
    args = parse_args()
    validate_datetime(args.datetime)
    
    lambda_function_name = f"{args.env}-lambda-populate-hfa-pending-request"

    payload = {
        "datetime": args.datetime
    }

    invoke_lambda(lambda_function_name, payload)

    print("\n Lambda invocation completed successfully.")
    sys.exit(0)