import boto3
from pandas import DataFrame
from io import BytesIO
from json import dumps
from asyncio import create_task
from botocore.exceptions import ClientError
from gooder_ai.s3.types import UploadParams, UploadFileParams
from gooder_ai import globals


# Convert dataframe to an in-memory csv sheet Ref: https://pandas.pydata.org/docs/reference/api/pandas.DataFrame.to_csv.html#pandas.DataFrame.to_csv
# Using buffer reference: https://stackoverflow.com/a/40540730
# Using BytesIO instead of StringIO because, s3 upload file function uses BytesIO
async def transform_dataframe_to_csv_file_object(data: DataFrame) -> BytesIO:
    csv_buffer = BytesIO()
    data.to_csv(csv_buffer, index=False)
    csv_buffer.seek(
        0
    )  # When we write something in buffer, file pointer is moved at the end of file. Here we are resetting the file pointer position to start after writing to file. Ref: https://stackoverflow.com/a/53485819
    return csv_buffer


# Convert config dictionary to a string to an in-memory json Ref: https://docs.python.org/3/library/json.html#basic-usage
# Using buffer reference: https://docs.python.org/3/library/io.html#binary-i-o
# Using BytesIO instead of StringIO because, s3 upload file function uses BytesIO
async def transform_config_to_json_file_object(config: dict) -> BytesIO:
    json_string = dumps(config)
    json_buffer = BytesIO(json_string.encode("utf-8"))
    json_buffer.seek(
        0
    )  # When we write something in buffer, file pointer is moved at the end of file. Here we are resetting the file pointer position to start after writing to file. Ref: https://stackoverflow.com/a/53485819
    return json_buffer


async def upload_handler(input: UploadParams) -> str | None:
    s3_client = input["s3_client"]
    identity_id = input["identity_id"]
    data = input["data"]
    mode = input["mode"]
    bucket_name = input["bucket_name"]
    file_name = input["file_name"]

    path: str | None = None
    parsed_data: BytesIO | None = None

    if isinstance(data, DataFrame):
        # The path to upload file is referenced from amplify. Ref: https://docs.amplify.aws/gen1/javascript/build-a-backend/storage/path/#using-protected-accesslevel
        path = f"{mode}/{identity_id}/data/{file_name}.csv"
        parsed_data = await transform_dataframe_to_csv_file_object(data)
    else:
        path = f"{mode}/{identity_id}/config/{file_name}.json"
        parsed_data = await transform_config_to_json_file_object(data)

    try:
        s3_client.upload_fileobj(parsed_data, bucket_name, path)
    except ClientError as e:
        print(f"Error uploading file: {e}")
        path = None
    return path


async def upload_files(input: UploadFileParams) -> dict[str, str | None]:
    aws_access_key_id = input["aws_access_key_id"]
    aws_secret_access_key = input["aws_secret_access_key"]
    aws_session_token = input["aws_session_token"]
    identity_id = input["identity_id"]
    data = input["data"]
    file_name = input["file_name"]
    config = input["config"]
    mode = input["mode"]

    bucket_name = input["bucket_name"]

    # Reference for using temporary credentials : https://docs.aws.amazon.com/IAM/latest/UserGuide/id_credentials_temp_use-resources.html
    s3_client = boto3.client(
        "s3",
        aws_access_key_id=aws_access_key_id,
        aws_secret_access_key=aws_secret_access_key,
        aws_session_token=aws_session_token,
    )
    # We use async keyword to define a co-routine if we need to execute a co-routine in the background we need to create a task of it.
    # Here we need both of these upload file to run in background that's why we are creating a tasks of them.
    # Ref: https://stackoverflow.com/questions/57966935/asyncio-task-vs-coroutine
    upload_csv_task = create_task(
        upload_handler(
            {
                "identity_id": identity_id,
                "s3_client": s3_client,
                "data": data,
                "file_name": file_name,
                "mode": mode,
                "bucket_name": bucket_name,
            }
        )
    )
    upload_config_task = create_task(
        upload_handler(
            {
                "bucket_name": bucket_name,
                "file_name": file_name,
                "data": config,
                "identity_id": identity_id,
                "mode": mode,
                "s3_client": s3_client,
            }
        )
    )

    csv_path = await upload_csv_task
    config_path = await upload_config_task

    output = {"csv_path": csv_path, "config_path": config_path}
    return output
