"""
Operations related to the audits_rows table in the database.
"""

import json
import uuid
from asyncio import Lock
from collections import defaultdict
from dataclasses import dataclass
from math import isnan

from .... import logger
from ....constants import DBColumns, Tables
from ....typings import AuditGroupID, AuditID, RowID, Rows

logger = logger.new_logger(__name__)


@dataclass(slots=True)
class Row:
    """Data class for an audit row."""

    audit: AuditID
    idx: RowID
    type: str
    data: dict


class AuditRows:
    """Class for interacting with the audits_rows table in the database."""

    def __init__(self, client):
        self.client = client
        self.rows: list[Row] = []
        self._rows_lock = Lock()

    async def add(
        self, audit_id: AuditID, row_idx: RowID, row_type: str, row_data: dict
    ) -> None:
        """Add a row to the list of rows to save to the database. The INSERT
        operation is performed when the save method is called, thus avoiding
        more than one database operation.

        Args:
            audit_id: Audit ID the row belongs to.
            row_idx: Row index in the audit, to prevent duplicates.
            row_type: Row type (e.g. "CARVED_OUT").
            row_data: Row data (e.g. {"upc": "123456789012", "isrc": "US1234567890",
            ...}).
        """
        async with self._rows_lock:
            self.rows.append(Row(audit_id, row_idx, row_type, row_data))

    async def get_by_audit_group_id(self, audit_group_id: AuditGroupID) -> Rows:
        """Get audit rows by audit group ID.

        Args:
            audit_group_id: Audit group ID for which to get the rows.
        """
        query = f"""
        SELECT
            AR.audit,
            AR.row_idx,
            ART.type,
            AR.data
        FROM {Tables.AUDITS_ROWS} AR
        LEFT JOIN {Tables.AUDITS} A ON A.id = AR.audit
        LEFT JOIN {Tables.AUDITS_ROWS_ROWTYPES} ARR ON AR.id = ARR.row_id
        LEFT JOIN {Tables.AUDITS_ROWTYPES} ART ON ARR.rowtype_id = ART.id
        WHERE A.`group` = %s
        """
        result = await self.client.db.execute_query_fetchall(query, (audit_group_id,))
        for row in result:
            row["data"] = json.loads(row["data"])
        return list(result)

    async def save(self) -> None:
        """Save rows to the database. Calling this method will clear the
        rows list to avoid duplicate insert attempts.
        """
        rows_grouped_by_type = await self._prepare_rows_for_save()

        temp_audits_rows: str = f"temp_audits_rows_{uuid.uuid4().hex}"

        conn = await self.client.db.get_connection()
        with conn:
            with conn.cursor() as cur:
                try:
                    # Use staging table for insertion to prevent unnecessary
                    # round trips to the database and to avoid potentially
                    # interfering with other transactions of the same type.
                    cur.execute(f"""
                        CREATE TEMPORARY TABLE {temp_audits_rows}
                        LIKE {Tables.AUDITS_ROWS};
                    """)

                    cur.execute(f"""
                        ALTER TABLE {temp_audits_rows}
                        ADD COLUMN type TEXT NOT NULL;
                    """)

                    temp_data = [
                        (row.audit, row.idx, _prepare_data(row.data), row_type)
                        for row_type, rows in rows_grouped_by_type.items()
                        for row in rows
                    ]
                    cur.executemany(
                        f"""
                        INSERT INTO {temp_audits_rows} (audit, row_idx, data, type)
                        VALUES (%s, %s, %s, %s);
                        """,
                        temp_data,
                    )

                    cur.execute(f"""
                        INSERT INTO {Tables.AUDITS_ROWTYPES} (type)
                        SELECT DISTINCT type FROM {temp_audits_rows}
                        ON DUPLICATE KEY UPDATE id=LAST_INSERT_ID({Tables.AUDITS_ROWTYPES}.id);
                        """)

                    cur.execute(f"""
                        INSERT INTO {Tables.AUDITS_ROWS} (audit, row_idx, data)
                        SELECT audit, row_idx, data FROM {temp_audits_rows};
                        """)

                    cur.execute(f"""
                        INSERT INTO {Tables.AUDITS_ROWS_ROWTYPES} (row_id, rowtype_id)
                        SELECT AR.id, ART.id
                        FROM {Tables.AUDITS_ROWS} AR
                        JOIN {temp_audits_rows} TAR ON AR.audit = TAR.audit
                          AND AR.row_idx = TAR.row_idx
                        JOIN {Tables.AUDITS_ROWTYPES} ART ON TAR.type = ART.type;
                        """)

                    conn.commit()

                    cur.execute(f"DROP TEMPORARY TABLE {temp_audits_rows};")

                except Exception as ex:
                    conn.rollback()
                    logger.error(
                        f"Transaction rolled back due to error inserting "
                        f"audit rows: {ex}"
                    )
                    raise ex

    async def _prepare_rows_for_save(self) -> dict[str, list[Row]]:
        """Prepare rows for saving to the database, removing them from the
        queue in the process. Returns a dictionary with rows grouped by type.

        Returns:
            Dictionary with rows grouped by type. Example:
            {
                "type A": [Row(...), Row(...)],
                "type B": [Row(...), Row(...)],
                "type C": [Row(...), Row(...)],
            }
        """
        rows_grouped_by_type: dict = defaultdict(list)
        async with self._rows_lock:
            extracted_rows = self.rows
            self.rows = []
        for row in extracted_rows:
            # Remove ROW_IDX from row data, if present, to prevent redundancy with
            # the row_idx argument. Creating a new dict to avoid mutating the
            # original one.
            if DBColumns.ROW_IDX in row.data:
                row.data = {k: v for k, v in row.data.items() if k != DBColumns.ROW_IDX}
            # Group rows by type to avoid multiple insertions of the same type
            rows_grouped_by_type[row.type].append(row)
        return rows_grouped_by_type


def _prepare_data(data: dict) -> str:
    """Prepare data for insertion into the database. Returns a stringified JSON.

    Args:
        data: Data to prepare (i.e. dict representing a row's data).
    """
    for key, value in data.items():
        # Convert sets and tuples to lists (JSON does not support sets and tuples)
        if isinstance(value, (set, tuple)):
            data[key] = list(value)
        # Convert NaN to None (which will become JSON null)
        if isinstance(value, float) and isnan(value):
            data[key] = None
    return json.dumps(data)