import asyncio
import json
import os
import traceback

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

from src import config
from src.query import BASE_GET_ASSET_FINALS


async def main() -> None:
    batch_size = int(os.getenv("BATCH_SIZE", "1"))
    num_workers = int(os.getenv("CONCURRENCY", "3"))
    env = config.ENVIRONMENT

    rows = get_input_rows()
    if not rows:
        print("No asset finals found for the supplied input.")
        return

    # chunk data into batches
    batches = [rows[x : x + batch_size] for x in range(0, len(rows), batch_size)]

    # lambda function to invoke
    function_name = f"{env}-lambda-assets-hive-ai-detection"

    # process batches concurrently
    queue: asyncio.Queue = asyncio.Queue()
    for batch in batches:
        queue.put_nowait(batch)
    await asyncio.gather(
        *[
            asyncio.create_task(process_queue(function_name, queue))
            for _ in range(num_workers)
        ]
    )


async def process_queue(function_name: str, queue: asyncio.Queue) -> None:
    session = aioboto3.Session()
    while not queue.empty():
        asset_finals = await queue.get()
        config = AioConfig(
            read_timeout=900, retries={"max_attempts": 10, "mode": "standard"}
        )
        async with session.client("lambda", config=config) as lambda_client:
            try:
                payload = {
                    "eventSource": "custom",
                    "assets": [
                        {"ASSET_FINAL_ID": x["id"], "DURATION_MS": x["duration"]}
                        for x in asset_finals
                    ],
                }
                try:
                    response = await lambda_client.invoke(
                        FunctionName=function_name,
                        InvocationType="RequestResponse",
                        Payload=json.dumps(payload).encode("utf-8"),
                    )
                    print(response)
                except lambda_client.exceptions.TooManyRequestsException:
                    print("Rate limit exceeded, re-queuing batch...")
                    await asyncio.sleep(60)
                    await queue.put(asset_finals)
                    continue
                except botocore.exceptions.ClientError as e:
                    error_code = e.response["Error"]["Code"]
                    if error_code == "ExpiredTokenException":
                        print("AWS token has expired.")
                        exit(1)
                    raise
                except botocore.exceptions.NoCredentialsError:
                    print("AWS credentials not found.")
                    exit(1)
            except Exception:
                traceback.print_exc()
        queue.task_done()


def get_input_rows() -> list[dict[str, int]]:
    return load_rows_from_snowflake()


def load_rows_from_snowflake() -> list[dict[str, int]]:
    sql, params = get_lookup_query_and_params()
    snowflake_conn = get_snowflake_connection()
    cursor = snowflake_conn.cursor(DictCursor)
    try:
        cursor.execute(sql, params)
        rows = cursor.fetchall()
    finally:
        cursor.close()
        snowflake_conn.close()
    print(f"Query returned {len(rows)} rows")
    return [
        {
            "id": int(row["ID"]),
            "duration": int(row["DURATION"]),
        }
        for row in rows
    ]


def get_lookup_query_and_params() -> tuple[str, list[int]]:
    label_ids = os.getenv("LABEL_ID")
    artist_ids = os.getenv("ARTIST_ID")
    subaccount_ids = os.getenv("SUBACCOUNT_ID")
    product_ids = os.getenv("PRODUCT_ID")

    if label_ids:
        id_list = [int(x.strip()) for x in label_ids.split(",")]
        placeholders = ", ".join(["%s"] * len(id_list))
        sql = BASE_GET_ASSET_FINALS + f" AND pj.vendor_id IN ({placeholders});"
        return sql, id_list

    if artist_ids:
        id_list = [int(x.strip()) for x in artist_ids.split(",")]
        placeholders = ", ".join(["%s"] * len(id_list))
        sql = BASE_GET_ASSET_FINALS + f" AND pj.artist_id IN ({placeholders});"
        return sql, id_list

    if subaccount_ids:
        id_list = [int(x.strip()) for x in subaccount_ids.split(",")]
        placeholders = ", ".join(["%s"] * len(id_list))
        sql = BASE_GET_ASSET_FINALS + f" AND pj.subaccount_id IN ({placeholders});"
        return sql, id_list

    if product_ids:
        id_list = [int(x.strip()) for x in product_ids.split(",")]
        placeholders = ", ".join(["%s"] * len(id_list))
        sql = BASE_GET_ASSET_FINALS + f" AND r.release_id IN ({placeholders});"
        return sql, id_list

    raise ValueError("No input source provided.")


def get_snowflake_connection() -> snowflake.connector.SnowflakeConnection:
    snowflake_connection = None
    try:
        if not config.SNOWFLAKE_WAREHOUSE:
            raise ValueError(
                "SNOWFLAKE_WAREHOUSE is empty or not set. "
                "Set the environment variable before running this script."
            )

        snowflake_connection = snowflake.connector.connect(
            user=config.SNOWFLAKE_USER,
            account=config.SNOWFLAKE_ACCOUNT,
            role=config.SNOWFLAKE_ROLE,
            authenticator="SNOWFLAKE_JWT",
            private_key=config.SNOWFLAKE_PRIVATE_KEY,
        )

        cursor = snowflake_connection.cursor()
        try:
            cursor.execute(f"USE WAREHOUSE {config.SNOWFLAKE_WAREHOUSE}")
        except Exception as ex:
            raise RuntimeError(
                "Snowflake warehouse is invalid or inaccessible for this role. "
                f"warehouse={config.SNOWFLAKE_WAREHOUSE}, "
                f"role={config.SNOWFLAKE_ROLE}, "
                f"environment={config.ENVIRONMENT}. Original error: {ex}"
            ) from ex

        try:
            if config.SNOWFLAKE_DATABASE:
                cursor.execute(f"USE DATABASE {config.SNOWFLAKE_DATABASE}")
            if config.SNOWFLAKE_SCHEMA:
                cursor.execute(f"USE SCHEMA {config.SNOWFLAKE_SCHEMA}")
        except Exception as ex:
            raise RuntimeError(
                "Snowflake database/schema is invalid or inaccessible for this role. "
                f"database={config.SNOWFLAKE_DATABASE}, "
                f"schema={config.SNOWFLAKE_SCHEMA}, "
                f"role={config.SNOWFLAKE_ROLE}. Original error: {ex}"
            ) from ex

        try:
            cursor.execute("SELECT CURRENT_WAREHOUSE()")
            warehouse_row = cursor.fetchone()
            if warehouse_row is None:
                raise RuntimeError(
                    "Unable to read CURRENT_WAREHOUSE() from Snowflake session"
                )
            current_warehouse = warehouse_row[0]
            cursor.execute(
                "SELECT CURRENT_ROLE(), CURRENT_DATABASE(), CURRENT_SCHEMA()"
            )
            context_row = cursor.fetchone()
            if context_row is None:
                raise RuntimeError(
                    "Unable to read Snowflake role/database/schema context"
                )
            current_role, current_database, current_schema = context_row
        finally:
            cursor.close()

    except Exception as ex:
        print(f"Failed to connect Snowflake database. {ex}")
        exit(1)

    print(
        "Snowflake database connected successfully. "
        f"Active warehouse: {current_warehouse}; "
        f"role={current_role}; database={current_database}; schema={current_schema}"
    )
    return snowflake_connection