import boto3
import gzip
import csv
import zipfile
import io
import pymysql
import sys
from datetime import datetime
from dotenv import load_dotenv

load_dotenv()

from config import RDS_DB_CONFIG  # noqa: E402

if not sys.argv[1]:
    print("No report run UUID provided")
    sys.exit(1)

REPORT_RUN_UUID = sys.argv[1]

BUCKET = "dev-collaborators-reports-test"

# Fields in order returned by the query
INPUT_FIELDS = [
    "txn_id",
    "collaboratorid",
    "statement_period_id",
    "statement_period_name",
    "transaction_type",
    "upc",
    "isrc",
    "cd_id",
    "track_jd",
    "trackname",
    "quantity",
    "account_payee_currency",
    "track_count",
    "ratetype",
    "countryname",
    "storename",
    "labelname",
    "productname",
    "artistname",
    "royalty_basis",
    "splitrate_sum",
    "collaborator_split",
    "collaborator_share",
]

client = boto3.client("s3")
rds_connection = pymysql.connect(
    **RDS_DB_CONFIG, cursorclass=pymysql.cursors.DictCursor
)

print("Fetching input file")

# Get the GZIP compressed input file
compressed_input_object = client.get_object(
    Bucket=BUCKET, Key=f"{REPORT_RUN_UUID}.csv.gz"
)
compressed_input_file = compressed_input_object["Body"]

# Create tuple representing creation date as required by ZipInfo
now = datetime.now()
created_date = (
    now.year,
    now.month,
    now.day,
    now.hour,
    now.minute,
    now.second,
    now.microsecond,
)

# Open the GZIP compressed input file for reading
with gzip.open(compressed_input_file, mode="rt") as input_file:
    input_reader = csv.DictReader(
        input_file, fieldnames=INPUT_FIELDS, delimiter=",", quotechar='"'
    )

    collaborator_id = None

    output_file = None
    output_zip = None
    output_csv = None
    output_writer = None
    total = None
    currency = None

    def finalise_output():
        output_csv.close()
        output_zip.close()

        # Reset the stream position ready for uploading
        output_file.seek(0)

        print(f"Uploading {collaborator_id}.zip")

        # Upload to S3
        client.upload_fileobj(
            output_file,
            BUCKET,
            f"{REPORT_RUN_UUID}/{collaborator_id}.zip",
        )

        with rds_connection.cursor() as cursor:
            cursor.execute(
                """
                    update report set
                        status = 'GENERATED',
                        file_location = %s,
                        amount = %s,
                        currency = %s,
                        generated_datetime = now()
                    where
                        report_run_uuid = %s
                        and collaborator_id = %s
                """,
                (
                    f"s3://{BUCKET}/{REPORT_RUN_UUID}/{collaborator_id}.zip",
                    total,
                    currency,
                    REPORT_RUN_UUID,
                    collaborator_id,
                ),
            )

        rds_connection.commit()

    for row in input_reader:
        # If this is a new collaborator or the first row
        if (
            collaborator_id != row["collaboratorid"]
            and row["collaboratorid"] != "\\\\N"
        ):
            # If there is a previous collaborator
            if collaborator_id:
                # Close and upload the previous output file
                finalise_output()

            # Set the new collaborator ID
            collaborator_id = row["collaboratorid"]
            total = row["collaborator_share"]
            currency = row["account_payee_currency"]

            # Create a new output file
            output_file = io.BytesIO()
            output_zip = zipfile.ZipFile(output_file, "w")

            output_csv_metadata = zipfile.ZipInfo(
                filename=f"{collaborator_id}.xls", date_time=created_date
            )
            output_csv_metadata.compress_type = zipfile.ZIP_DEFLATED

            output_csv = output_zip.open(output_csv_metadata, "w")
            output_csv_text = io.TextIOWrapper(output_csv, encoding="utf-8")

            output_writer = csv.writer(
                output_csv_text, delimiter="\t", quotechar='"', quoting=csv.QUOTE_ALL
            )

            # Write the summary section
            output_writer.writerow(("Report Run UUID", REPORT_RUN_UUID))
            output_writer.writerow(("Collaborator ID", collaborator_id))
            output_writer.writerow(("Collaborator Total", row["collaborator_share"]))
            output_writer.writerow(("Period ID", row["statement_period_id"]))
            output_writer.writerow(("Period Name", row["statement_period_name"]))
            output_writer.writerow(())

            # Write the column headers
            output_writer.writerow(INPUT_FIELDS)
        else:
            # This is a row for the current collaborator so just write it
            output_writer.writerow(row[key] for key in INPUT_FIELDS)

    # Close and upload the final output file
    finalise_output()