import asyncio
import json
from copy import deepcopy
from datetime import datetime
from typing import List

import pymysql
import pymysql.cursors

from .... import logger
from ....constants import Tables, User
from ....typings import JWTSubject, OptionalRow, UserID
from ....utils.sql import placeholders

logger = logger.new_logger(__name__)


class Users:
    """Class for interacting with the users table."""

    _null_placeholder = "__null__"

    def __init__(self, client):
        self.client = client

    async def create(self, subject: JWTSubject, fields: dict = None) -> int:
        """Create user for a subject.
        Args:
            subject: JWT subject.
            fields: Fields to insert, excluding subject. Any field not provided
                will be set to NULL or the default value as specified in the
                database schema, if applicable.
        """
        if not fields:
            fields = {}
        fields[User.SUBJECT] = subject  # Override subject in case it's present
        # in fields.

        # Remove any fields that are None, as they will be set to NULL or the default
        # value in the database schema (if the schema allows it).
        non_null_columns = [
            column
            for column in [
                User.SUBJECT,
                User.NICKNAME,
                User.NAME,
                User.EMAIL,
                User.TIMEZONE,
                User.LOCALE,
            ]
            if fields.get(column) is not None
        ]
        sql_columns = ", ".join(f"`{col}`" for col in non_null_columns)
        query = (
            f"INSERT INTO {Tables.USERS} "
            f"({sql_columns}) "
            f"VALUES ({placeholders(non_null_columns)})"
        )
        try:
            params = tuple(fields.get(column) for column in non_null_columns)
            last_row_id = await self.client.db.execute_query_nofetch(
                query, params, lastrowid=True
            )
            return last_row_id
        except pymysql.err.IntegrityError as ex:
            if User.NICKNAME in str(ex).lower():
                # Nickname taken, try again with a partially random one
                new_nickname = (
                    f"{fields.get(User.NICKNAME)}_{int(datetime.now().timestamp())}"
                )
                fields[User.NICKNAME] = new_nickname
                await self.create(subject, fields)
            else:
                raise RuntimeError("Error while creating user") from ex

    async def _get(
        self,
        subject: JWTSubject = None,
        user_id: UserID = None,
    ) -> OptionalRow:
        """Get user by where clause.
        Args:
            subject: JWT subject.
            user_id: User ID.
        Returns:
            User data or None if the user does not exist.
        """
        if sum(v is not None for v in [subject, user_id]) != 1:
            raise ValueError("Exactly one of subject or user_id must be provided")

        scopes = User.SCOPES

        where_clause = f"{User.SUBJECT} = %s" if subject else f"{User.ID} = %s"

        query = f"""
            SELECT
                U.*,
            IFNULL(
                JSON_OBJECTAGG(
                    -- Provide default values in case of NULL because user has no scopes
                    IFNULL(UPS.scope, '{self._null_placeholder}'),
                    IFNULL(UPL.level, '{self._null_placeholder}')
                ),
                JSON_OBJECT()
            ) AS {scopes}
            FROM
                {Tables.USERS} U
            LEFT JOIN
                {Tables.USERS_PERMISSIONS} UP ON U.id = UP.user
            LEFT JOIN
                {Tables.USERS_PERMISSIONS_SCOPES} UPS ON UP.scope = UPS.id
            LEFT JOIN
                {Tables.USERS_PERMISSIONS_LEVELS} UPL ON UP.level = UPL.id
            WHERE
                U.{where_clause}
            GROUP BY
                U.id;
        """
        passed_value = next(v for v in [subject, user_id] if v is not None)
        user_raw = await self.client.db.execute_query_fetchone(query, (passed_value,))
        user = self._parse_scopes([user_raw])[0] if user_raw else user_raw
        return user

    async def get(self, subject: JWTSubject) -> OptionalRow:
        """Get user by subject.
        Args:
            subject: JWT subject.
        Returns:
            User data or None if the user does not exist.
        """
        logger.debug("Getting user with subject: {}", subject)
        user = await self._get(subject=subject)
        return user

    async def get_by_id(self, user_id: UserID) -> OptionalRow:
        """Get user by user ID.
        Args:
            user_id: User ID.
        Returns:
            User data or None if the user does not exist.
        """
        logger.debug("Getting user with ID: {}", user_id)
        user = await self._get(user_id=user_id)
        return user

    async def list(
        self, limit: int = None, offset: int = None
    ) -> tuple[int, list[OptionalRow]]:
        """List users with total count.

        Args:
            limit: Number of users to return.
            offset: Offset to start from.

        Returns:
            Tuple containing total user count and list of users.
        """
        scopes = User.SCOPES
        include_offset = limit and offset
        query = f"""
            SELECT
                U.id,
                U.locale,
                U.timezone,
                U.nickname,
                U.name,
                U.email,
                IFNULL(
                JSON_OBJECTAGG(
                    -- Provide default values in case of NULL because user has no scopes
                    IFNULL(UPS.scope, '{self._null_placeholder}'),
                    IFNULL(UPL.level, '{self._null_placeholder}')
                ),
                JSON_OBJECT()
            ) AS {scopes}
            FROM
                {Tables.USERS} U
            LEFT JOIN
                {Tables.USERS_PERMISSIONS} UP ON U.id = UP.user
            LEFT JOIN
                {Tables.USERS_PERMISSIONS_SCOPES} UPS ON UP.scope = UPS.id
            LEFT JOIN
                {Tables.USERS_PERMISSIONS_LEVELS} UPL ON UP.level = UPL.id
            GROUP BY
                U.id
            {'LIMIT %s' if limit else ''}
            {'OFFSET %s' if include_offset else ''};
        """
        params = []
        if limit:
            params.append(limit)
        if include_offset:  # Offset requires limit
            params.append(offset)

        logger.debug("Listing users with limit: %s and offset: %s", limit, offset)

        total_count_helper_col = "total_count"

        async def fetch_users():
            """Async function to fetch users and remove NULL scopes."""
            users = await self.client.db.execute_query_fetchall(query, tuple(params))
            return self._parse_scopes(users)

        # Fetch total count and users in parallel
        result, _total_count_obj = await asyncio.gather(
            *[
                fetch_users(),
                self.client.db.execute_query_fetchone(
                    f"SELECT COUNT(*) AS {total_count_helper_col} FROM {Tables.USERS};"
                ),
            ]
        )

        total_count = _total_count_obj[total_count_helper_col]

        return total_count, result

    async def _update(
        self,
        subject: JWTSubject = None,
        user_id: UserID = None,
        fields: dict = None,
    ) -> None:
        """Partially update user data.
        Args:
            subject: JWT subject.
            fields: Fields to update.
        """
        if fields is None:
            raise ValueError("fields must be provided")
        if sum(v is not None for v in [subject, user_id]) != 1:
            raise ValueError("Exactly one of subject or user_id must be provided")

        where_clause = f"{User.SUBJECT} = %s" if subject else f"{User.ID} = %s"
        set_clause: str = ", ".join(f"{field} = %s" for field in fields.keys())

        query = f"UPDATE {Tables.USERS} SET {set_clause} WHERE {where_clause}"

        passed_value = next(v for v in [subject, user_id] if v is not None)
        params = tuple(fields.values()) + (passed_value,)

        await self.client.db.execute_query_nofetch(query, params)

    async def update(self, subject: JWTSubject, fields: dict) -> None:
        """Partially update user data.
        Args:
            subject: JWT subject.
            fields: Fields to update.
        """
        await self._update(subject=subject, fields=fields)

    async def update_by_id(self, user_id: UserID, fields: dict) -> None:
        """Partially update user data (by user ID).
        Args:
            user_id: User ID.
            fields: Fields to update.
        """
        await self._update(user_id=user_id, fields=fields)

    def _parse_scopes(self, users: List[OptionalRow]) -> List[OptionalRow]:
        """Parse scopes from the scopes dictionary of the users.
        Args:
            users: List of users, which must have a dict
                containing scopes as a JSON string.
        Returns:
            List of users with NULL values removed from the scopes dictionary.
        """
        cache = {}  # Cache parsed scopes to avoid re-parsing

        for user in users:
            scopes = user[User.SCOPES]
            if scopes not in cache:
                cache[scopes] = {
                    scope: level
                    for scope, level in json.loads(scopes).items()
                    if self._null_placeholder not in {scope, level}
                }
            # Avoid reference issues by always creating a new dict. Otherwise, the
            # same dict would be shared among all users with the same scopes.
            user[User.SCOPES] = deepcopy(cache[scopes])

        return users


class Permissions:
    """Class for interacting with the users_permissions table."""

    def __init__(self, client):
        self.client = client

    async def get_meta(self) -> dict[str, dict]:
        """Get metadata for permissions.

        Returns:
            Metadata for permissions, including all possible levels and scopes,
            as stored in the database.
        """
        query_levels = f"SELECT id, level FROM {Tables.USERS_PERMISSIONS_LEVELS}"
        query_scopes = f"""
            SELECT id, scope, description
            FROM {Tables.USERS_PERMISSIONS_SCOPES}
        """
        fetch_all = self.client.db.execute_query_fetchall
        levels, scopes = await asyncio.gather(
            *[
                fetch_all(query_levels),
                fetch_all(query_scopes),
            ]
        )
        meta = {
            "levels": {level["id"]: level["level"] for level in levels},
            "scopes": {scope["id"]: scope for scope in scopes},
        }
        return meta

    async def upsert(self, user_id: UserID, scopes: dict) -> None:
        """Upsert and/or delete permissions for a user.

        Any scope not existing for the user will be inserted, whereas existing
        scopes will have their levels updated.

        Any scope with a None level will be deleted, if it exists.

        Args:
            user_id: User ID.
            scopes: Scopes to update. E.g. {"MODULES.STATISTICS": "READ"}
        """
        conn = await self.client.db.get_connection()
        with conn:  # Atomic transaction
            try:
                with conn.cursor() as cur:
                    # Get scope IDs to replace scope names provided in the request
                    query_get_scope_ids = f"""
                        SELECT id, scope
                        FROM {Tables.USERS_PERMISSIONS_SCOPES}
                        WHERE scope IN ({placeholders(scopes.keys())})
                    """
                    cur.execute(query_get_scope_ids, tuple(scopes.keys()))
                    scope_ids = {row["scope"]: row["id"] for row in cur.fetchall()}

                    # Delete existing permissions related to the user and scopes
                    # provided in the request
                    query_delete = f"""
                        DELETE FROM {Tables.USERS_PERMISSIONS}
                        WHERE user = %s
                        AND scope IN ({placeholders(scopes.keys())})
                    """
                    params = (user_id, *[scope_ids[scope] for scope in scopes])
                    cur.execute(query_delete, params)

                    # Get level IDs to replace level names provided in the request
                    query_get_level_ids = f"""
                        SELECT id, level
                        FROM {Tables.USERS_PERMISSIONS_LEVELS}
                        WHERE level IN ({placeholders(scopes.values())})
                    """
                    cur.execute(query_get_level_ids, tuple(scopes.values()))
                    level_ids = {row["level"]: row["id"] for row in cur.fetchall()}

                    # Insert new permissions
                    query_insert = f"""
                        INSERT INTO {Tables.USERS_PERMISSIONS}
                        (user, scope, level, date_created_utc)
                        VALUES (%s, %s, %s, UTC_TIMESTAMP())
                    """
                    # Skip NULL levels, as user intention is to
                    # remove the scope
                    params = [
                        (user_id, scope_ids[scope], level_ids[level])
                        for scope, level in scopes.items()
                        if level is not None
                    ]
                    cur.executemany(query_insert, params)
                    conn.commit()

            except Exception as ex:
                logger.error("Exception occurred: {}", ex)
                conn.rollback()  # Rollback in case of an exception