# fastapi_server/user_mgmt.py
from . import db_setup as db
import bcrypt
import asyncio
from aiomysql import DictCursor

async def fetch_user_info(arg: str, type: str):
    """Fetches user info.
        This kind of query is supposed to either return ONE user or nothing.
    """
    if db.mode == "mysql":
        async with db.pool.acquire() as conn:
            async with conn.cursor(DictCursor) as cur:
                await cur.execute("SELECT * FROM users WHERE %s = %s", (type, arg))
                return await cur.fetchone()
    elif db.mode == "sqlite":
        async with db.file.connect() as conn:
            async with conn.execute("SELECT * FROM users WHERE ? = ?", (type, arg)) as cursor:
                return await cursor.fetchone()
    return None


async def create_user(first_name: str, last_name: str, email: str, password: str):
    """Creates a new user."""
    hashed_password = bcrypt.hashpw(password.encode(), bcrypt.gensalt())
    if db.db.mode == "mysql":
        async with db.db.pool.acquire() as conn:
            async with conn.cursor() as cur:
                await cur.execute("INSERT INTO users (`first_name`,`last_name`,`email`,`password`) VALUES (%s,%s,%s,%s)", (first_name, last_name, email, hashed_password.decode()))
                await conn.commit()
                return True
    elif db.db.mode == "sqlite":
        async with db.db.file.connect() as conn:
            await conn.execute("INSERT INTO users (`first_name`,`last_name`,`email`,`password`) VALUES (?, ?, ?, ?)", (first_name, last_name, email, hashed_password.decode()))
            await conn.commit()
            return True
    return False


async def login_user(email: str, password: str):
    """Logs in a user."""
    user_data = await fetch_user_info(email, "email")
    if user_data is None:
        return None  # User doesn't exist
    else:
        if bcrypt.checkpw(password.encode(), user_data.get("password").encode()):
            return user_data  # Login successful
        else:
            return None  # Incorrect password
        

async def edit_user(user_id: int, updates: dict):
    """Edits an existing user."""
    if db.mode == "mysql":
        placeholders = ", ".join(["%s"] * len(updates))
        values = list(updates.values())
        query = f"UPDATE users SET {', '.join([f'{key} = %s' for key in updates.keys()])} WHERE id = %s"
        values.append(user_id)
        async with db.pool.acquire() as conn:
            async with conn.cursor() as cur:
                await cur.execute(query, values)
                await conn.commit()
                return True
    elif db.mode == "sqlite":
        placeholders = ", ".join(["?"] * len(updates))
        values = list(updates.values())
        values.append(user_id)
        query = f"UPDATE users SET {', '.join([f'{key} = ?' for key in updates.keys()])} WHERE id = ?"
        async with db.file.connect() as conn:
            await conn.execute(query, values)
            await conn.commit()
            return True
    return False