"""sync asset files from Production to QA environment."""

import argparse
import asyncio
import os

import aioboto3
import snowflake.connector
from aiobotocore.config import AioConfig
from botocore.exceptions import ClientError
from snowflake.connector import DictCursor

from src.query import GET_AUDIO_FILES, GET_VIDEO_FILES


async def main():
    """Entrypoint."""
    args = get_args()

    # setup snowflake connection
    snowflake_user = os.environ["SNOWFLAKE_USER"]
    snowflake_db = os.environ["SNOWFLAKE_DB"]
    snowflake_warehouse = os.environ["SNOWFLAKE_WAREHOUSE"]
    snowflake_auth_type = os.environ["SNOWFLAKE_AUTH_TYPE"]
    snowflake_private_key_dir = os.environ.get("SNOWFLAKE_PRIVATE_KEY_DIR")
    snowflake_private_key_file = os.environ.get("SNOWFLAKE_PRIVATE_KEY_FILE")
    snowflake_private_key_file_pwd = os.environ.get("SNOWFLAKE_PRIVATE_KEY_FILE_PWD")

    snowflake_account = os.environ["SNOWFLAKE_ACCOUNT"]
    if snowflake_auth_type not in ("externalbrowser", "snowflake"):
        print(f"Invalid 'SNOWFLAKE_AUTH_TYPE' of '{snowflake_auth_type}'")
        exit(1)

    conn_args = {
        "user": snowflake_user,
        "authenticator": snowflake_auth_type,
        "account": snowflake_account,
    }
    if snowflake_auth_type == "snowflake":
        if snowflake_private_key_dir and snowflake_private_key_file:
            conn_args["private_key_file"] = os.path.join(
                snowflake_private_key_dir, snowflake_private_key_file
            )
        if snowflake_private_key_file_pwd:
            conn_args["private_key_file_pwd"] = snowflake_private_key_file_pwd
    snowflake_conn = snowflake.connector.connect(**conn_args)
    cursor = snowflake_conn.cursor(DictCursor)
    cursor.execute(f"USE DATABASE {snowflake_db}")
    cursor.execute(f"USE WAREHOUSE {snowflake_warehouse}")

    # handle args
    num_workers = args.workers
    params = {
        "upc": args.upc,
        "vendor_id": args.vendor_id,
        "gt_release_id": args.gt_release_id,
        "lt_release_id": args.lt_release_id,
    }

    # query snowflake for assets
    cursor.execute(GET_AUDIO_FILES, params)
    audio_rows = cursor.fetchall()
    cursor.execute(GET_VIDEO_FILES, params)
    video_rows = cursor.fetchall()
    rows = sorted(audio_rows + video_rows, key=lambda r: r["RELEASE_ID"])
    cursor.close()
    snowflake_conn.close()
    print(f"""
        Query returned:
          {len(audio_rows)} audio files
          {len(video_rows)} video files
          {len(rows)} total
    """)

    # queue up assets for sync
    queue = asyncio.Queue()
    for row in rows:
        queue.put_nowait((aioboto3.Session(), row))

    # copy assets to QA
    try:
        await asyncio.gather(
            *[
                asyncio.create_task(process_queue_with_copy(queue))
                for _ in range(num_workers)
            ]
        )
    except ClientError as e:
        if e.response["Error"]["Code"] == "RequestTimeTooSkewed":
            print(
                f"\nToo many concurrent requests caused a time skew error. "
                f"Try re-running with fewer workers (current: --workers {num_workers})."
            )
            exit(1)
        raise
    print("Script completed.")


def get_args():
    """Read and validate CLI args."""
    parser = argparse.ArgumentParser(
        formatter_class=argparse.ArgumentDefaultsHelpFormatter
    )

    parser.add_argument(
        "--upc",
        type=int,
        default=0,
        help="UPC that needs to sync assets on QA",
    )
    parser.add_argument(
        "--vendor_id",
        type=int,
        default=0,
        help="LabelID that needs to sync all assets on QA",
    )
    parser.add_argument(
        "--workers",
        type=int,
        default=20,
        help="Number of parallel workers to download with",
    )
    parser.add_argument(
        "--gt-release-id",
        type=int,
        default=0,
        help="Filter RELEASES to rows with RELEASE_ID greater than this value",
    )
    parser.add_argument(
        "--lt-release-id",
        type=int,
        default=0,
        help="Filter RELEASES to rows with RELEASE_ID less than this value",
    )
    args = parser.parse_args()

    # validate upc or vendor_id selection
    if not args.upc and not args.vendor_id:
        parser.print_help()
        print("\nMust pass at least one argument (either upc or vendor_id)!")
        exit(1)

    # validate release id range
    if args.gt_release_id and args.lt_release_id:
        if args.gt_release_id >= args.lt_release_id:
            parser.print_help()
            print("\n--gt-release-id must be less than --lt-release-id!")
            exit(1)

    return args


async def process_queue_with_copy(queue):
    """Pop from queue and copy to QA bucket."""
    while not queue.empty():
        data = await queue.get()
        row = data[1]
        row_title = "_".join(
            str(x)
            for x in [row["RELEASE_ID"], row["UPC"], row["ASSET_TYPE"]]
            if x is not None
        )
        # Copy to destination QA bucket
        await copy_s3_object(data[0], row_title, row["S3_BUCKET"], row["S3_FILENAME"])

        queue.task_done()


async def copy_s3_object(s3_session, row_title, s3_bucket, s3_file_path):
    """Copy assets from PROD S3 bucket to QA S3 bucket."""
    destination_bucket = s3_bucket.replace("prod-", "qa-")
    destination_path = s3_file_path

    async with s3_session.client("s3", config=AioConfig(read_timeout=300)) as client:
        # Check if destination already exists; skip copy if so
        try:
            await client.head_object(Bucket=destination_bucket, Key=destination_path)
            print(
                f"{row_title:<40} : Skipping {s3_file_path} — already exists in "
                f"{destination_bucket}"
            )
            return
        except client.exceptions.ClientError as e:
            if e.response["Error"]["Code"] != "404":
                raise

        print(
            f"{row_title:<40} : Copying {s3_file_path} file from {s3_bucket} "
            f"to {destination_bucket}"
        )

        # Copy object using S3 managed copy (handles multipart for files > 5 GB)
        copy_source = {"Bucket": s3_bucket, "Key": s3_file_path}
        await client.copy(copy_source, destination_bucket, destination_path)