"""MySQL database layer for the encrypted chat server."""

import os
import uuid

import mysql.connector
from dotenv import load_dotenv

from crypto_utils import (
    generate_identity_keypair,
    serialize_ed25519_public,
    generate_signed_prekey,
    serialize_x25519_public,
    generate_one_time_prekeys,
)

load_dotenv()

# Sentinel device_id for self-encrypted copies and legacy (pre-multi-device) rows
SELF_DEVICE_ID = "00000000-0000-0000-0000-000000000000"


def get_connection():
    """Create a new MySQL connection from environment variables."""
    return mysql.connector.connect(
        host=os.getenv("MYSQL_HOST", "localhost"),
        port=int(os.getenv("MYSQL_PORT", "3306")),
        user=os.getenv("MYSQL_USER", "root"),
        password=os.getenv("MYSQL_PASSWORD", ""),
        database=os.getenv("MYSQL_DATABASE", "encrypted_chat"),
    )


def generate_uuid() -> str:
    return str(uuid.uuid4())


# --- Devices ---

def create_device(user_id: str, device_name: str | None = None) -> str:
    """Create a new device for a user. Returns device_id."""
    conn = get_connection()
    try:
        cursor = conn.cursor()
        device_id = generate_uuid()
        cursor.execute(
            "INSERT INTO devices (id, user_id, device_name) VALUES (%s, %s, %s)",
            (device_id, user_id, device_name),
        )
        conn.commit()
        return device_id
    finally:
        conn.close()


def get_user_devices(user_id: str) -> list[dict]:
    """Get all devices for a user."""
    conn = get_connection()
    try:
        cursor = conn.cursor(dictionary=True)
        cursor.execute(
            "SELECT id, user_id, device_name, created_at, last_seen_at "
            "FROM devices WHERE user_id = %s ORDER BY created_at",
            (user_id,),
        )
        return cursor.fetchall()
    finally:
        conn.close()


def get_device(device_id: str) -> dict | None:
    """Get a single device by ID."""
    conn = get_connection()
    try:
        cursor = conn.cursor(dictionary=True)
        cursor.execute(
            "SELECT id, user_id, device_name, created_at, last_seen_at "
            "FROM devices WHERE id = %s",
            (device_id,),
        )
        return cursor.fetchone()
    finally:
        conn.close()


def update_device_last_seen(device_id: str):
    """Update last_seen_at timestamp for a device."""
    conn = get_connection()
    try:
        cursor = conn.cursor()
        cursor.execute(
            "UPDATE devices SET last_seen_at = NOW() WHERE id = %s",
            (device_id,),
        )
        conn.commit()
    finally:
        conn.close()


def delete_device(device_id: str):
    """Delete a device. CASCADE removes its prekeys."""
    conn = get_connection()
    try:
        cursor = conn.cursor()
        cursor.execute("DELETE FROM devices WHERE id = %s", (device_id,))
        # Also clean up prekeys explicitly for device_id column
        cursor.execute("DELETE FROM signed_prekeys WHERE device_id = %s", (device_id,))
        cursor.execute("DELETE FROM one_time_prekeys WHERE device_id = %s", (device_id,))
        conn.commit()
    finally:
        conn.close()


# --- Users ---

def create_user(username: str, email: str, rsa_public_key_pem: str, identity_key: bytes) -> str:
    """Register a new user. Returns user ID."""
    conn = get_connection()
    try:
        cursor = conn.cursor()
        user_id = generate_uuid()
        cursor.execute(
            "INSERT INTO users (id, username, email, rsa_public_key, identity_key) "
            "VALUES (%s, %s, %s, %s, %s)",
            (user_id, username, email, rsa_public_key_pem, identity_key),
        )
        conn.commit()
        return user_id
    finally:
        conn.close()


def get_user_by_email(email: str) -> dict | None:
    """Get user by email."""
    conn = get_connection()
    try:
        cursor = conn.cursor(dictionary=True)
        cursor.execute(
            "SELECT id, username, rsa_public_key, email, identity_key FROM users WHERE email = %s",
            (email,),
        )
        return cursor.fetchone()
    finally:
        conn.close()


def get_user_by_id(user_id: str) -> dict | None:
    """Get user by ID."""
    conn = get_connection()
    try:
        cursor = conn.cursor(dictionary=True)
        cursor.execute(
            "SELECT id, username, rsa_public_key, email, identity_key FROM users WHERE id = %s",
            (user_id,),
        )
        return cursor.fetchone()
    finally:
        conn.close()


def get_user_contacts(user_id: str) -> list[str]:
    """Get all user IDs that share at least one conversation with the given user."""
    conn = get_connection()
    try:
        cursor = conn.cursor()
        cursor.execute(
            "SELECT DISTINCT cm2.user_id "
            "FROM conversation_members cm1 "
            "JOIN conversation_members cm2 ON cm1.conversation_id = cm2.conversation_id "
            "WHERE cm1.user_id = %s AND cm2.user_id != %s",
            (user_id, user_id),
        )
        return [row[0] for row in cursor.fetchall()]
    finally:
        conn.close()


def update_user_rsa_key(user_id: str, rsa_public_key_pem: str):
    """Update user's RSA public key (for login)."""
    conn = get_connection()
    try:
        cursor = conn.cursor()
        cursor.execute("UPDATE users SET rsa_public_key = %s WHERE id = %s", (rsa_public_key_pem, user_id))
        conn.commit()
    finally:
        conn.close()


# --- Pre-keys ---

def store_signed_prekey(user_id: str, spk_id: str, public_key: bytes, signature: bytes,
                        device_id: str | None = None):
    """Store (or replace) a signed pre-key for a user's device."""
    conn = get_connection()
    try:
        cursor = conn.cursor()
        # Remove old SPKs for this user+device
        if device_id:
            cursor.execute("DELETE FROM signed_prekeys WHERE user_id = %s AND device_id = %s",
                           (user_id, device_id))
        else:
            cursor.execute("DELETE FROM signed_prekeys WHERE user_id = %s AND device_id IS NULL",
                           (user_id,))
        cursor.execute(
            "INSERT INTO signed_prekeys (id, user_id, device_id, public_key, signature) "
            "VALUES (%s, %s, %s, %s, %s)",
            (spk_id, user_id, device_id, public_key, signature),
        )
        conn.commit()
    finally:
        conn.close()


def get_signed_prekey(user_id: str, device_id: str | None = None) -> dict | None:
    """Get the current signed pre-key for a user (optionally per device)."""
    conn = get_connection()
    try:
        cursor = conn.cursor(dictionary=True)
        if device_id:
            cursor.execute(
                "SELECT id, public_key, signature, device_id, created_at FROM signed_prekeys "
                "WHERE user_id = %s AND device_id = %s "
                "ORDER BY created_at DESC LIMIT 1",
                (user_id, device_id),
            )
        else:
            cursor.execute(
                "SELECT id, public_key, signature, device_id, created_at FROM signed_prekeys "
                "WHERE user_id = %s ORDER BY created_at DESC LIMIT 1",
                (user_id,),
            )
        return cursor.fetchone()
    finally:
        conn.close()


def store_one_time_prekeys(user_id: str, prekeys: list[dict], device_id: str | None = None):
    """Store a batch of one-time pre-keys. Each dict has {id, public_key (bytes)}."""
    conn = get_connection()
    try:
        cursor = conn.cursor()
        for pk in prekeys:
            cursor.execute(
                "INSERT INTO one_time_prekeys (id, user_id, device_id, public_key) "
                "VALUES (%s, %s, %s, %s)",
                (pk["id"], user_id, device_id, pk["public_key"]),
            )
        conn.commit()
    finally:
        conn.close()


def consume_one_time_prekey(user_id: str, device_id: str | None = None) -> dict | None:
    """Atomically consume one OTP: SELECT FOR UPDATE + DELETE.
    Returns {id, public_key} or None."""
    conn = get_connection()
    try:
        cursor = conn.cursor(dictionary=True)
        conn.start_transaction()
        if device_id:
            cursor.execute(
                "SELECT id, public_key FROM one_time_prekeys "
                "WHERE user_id = %s AND device_id = %s LIMIT 1 FOR UPDATE",
                (user_id, device_id),
            )
        else:
            cursor.execute(
                "SELECT id, public_key FROM one_time_prekeys "
                "WHERE user_id = %s LIMIT 1 FOR UPDATE",
                (user_id,),
            )
        row = cursor.fetchone()
        if row:
            cursor.execute("DELETE FROM one_time_prekeys WHERE id = %s", (row["id"],))
        conn.commit()
        return row
    except Exception:
        conn.rollback()
        raise
    finally:
        conn.close()


def count_one_time_prekeys(user_id: str, device_id: str | None = None) -> int:
    """Count remaining OTPs for a user (optionally per device)."""
    conn = get_connection()
    try:
        cursor = conn.cursor()
        if device_id:
            cursor.execute(
                "SELECT COUNT(*) FROM one_time_prekeys WHERE user_id = %s AND device_id = %s",
                (user_id, device_id),
            )
        else:
            cursor.execute("SELECT COUNT(*) FROM one_time_prekeys WHERE user_id = %s", (user_id,))
        return cursor.fetchone()[0]
    finally:
        conn.close()


def get_key_bundle(user_id: str) -> dict | None:
    """Get complete key bundle for X3DH (single device — legacy compat).

    Returns {identity_key, signed_prekey_id, signed_prekey, spk_signature,
             one_time_prekey_id, one_time_prekey} or None.
    OTP is consumed atomically.
    """
    conn = get_connection()
    try:
        cursor = conn.cursor(dictionary=True)
        # Get user identity key
        cursor.execute("SELECT identity_key FROM users WHERE id = %s", (user_id,))
        user = cursor.fetchone()
        if not user:
            return None

        # Get signed prekey
        cursor.execute(
            "SELECT id, public_key, signature, device_id FROM signed_prekeys WHERE user_id = %s "
            "ORDER BY created_at DESC LIMIT 1",
            (user_id,),
        )
        spk = cursor.fetchone()
        if not spk:
            return None

        # Consume one OTP (may be None) — use transaction for atomicity (H12 fix)
        conn.start_transaction()
        cursor.execute(
            "SELECT id, public_key FROM one_time_prekeys WHERE user_id = %s LIMIT 1 FOR UPDATE",
            (user_id,),
        )
        opk = cursor.fetchone()
        if opk:
            cursor.execute("DELETE FROM one_time_prekeys WHERE id = %s", (opk["id"],))
        conn.commit()

        result = {
            "identity_key": user["identity_key"],
            "signed_prekey_id": spk["id"],
            "signed_prekey": spk["public_key"],
            "spk_signature": spk["signature"],
        }
        if opk:
            result["one_time_prekey_id"] = opk["id"]
            result["one_time_prekey"] = opk["public_key"]
        return result
    except Exception:
        try:
            conn.rollback()
        except Exception:
            pass
        raise
    finally:
        conn.close()


def get_key_bundles_for_user(user_id: str) -> dict | None:
    """Get key bundles for ALL devices of a user. Returns
    {identity_key, device_bundles: [{device_id, signed_prekey_id, signed_prekey_pub,
      spk_signature, opk_id, opk_pub}]} or None.
    Consumes one OPK per device atomically.
    """
    conn = get_connection()
    try:
        cursor = conn.cursor(dictionary=True)
        # Get user identity key
        cursor.execute("SELECT identity_key FROM users WHERE id = %s", (user_id,))
        user = cursor.fetchone()
        if not user:
            return None

        # Get all signed prekeys (one per device, most recent)
        cursor.execute(
            "SELECT id, public_key, signature, device_id FROM signed_prekeys "
            "WHERE user_id = %s ORDER BY created_at DESC",
            (user_id,),
        )
        all_spks = cursor.fetchall()
        if not all_spks:
            return None

        # De-duplicate: keep only the most recent SPK per device_id
        seen_devices = set()
        spks_by_device = []
        for spk in all_spks:
            dev = spk.get("device_id") or "__legacy__"
            if dev not in seen_devices:
                seen_devices.add(dev)
                spks_by_device.append(spk)

        device_bundles = []
        # Commit the implicit transaction from the read-only queries above
        # so we can start an explicit transaction for atomic OPK consumption.
        conn.commit()
        conn.start_transaction()
        for spk in spks_by_device:
            dev_id = spk.get("device_id")
            # Consume one OPK for this device
            if dev_id:
                cursor.execute(
                    "SELECT id, public_key FROM one_time_prekeys "
                    "WHERE user_id = %s AND device_id = %s LIMIT 1 FOR UPDATE",
                    (user_id, dev_id),
                )
            else:
                cursor.execute(
                    "SELECT id, public_key FROM one_time_prekeys "
                    "WHERE user_id = %s AND device_id IS NULL LIMIT 1 FOR UPDATE",
                    (user_id,),
                )
            opk = cursor.fetchone()
            if opk:
                cursor.execute("DELETE FROM one_time_prekeys WHERE id = %s", (opk["id"],))

            bundle = {
                "device_id": dev_id,
                "signed_prekey_id": spk["id"],
                "signed_prekey_pub": spk["public_key"],
                "spk_signature": spk["signature"],
            }
            if opk:
                bundle["opk_id"] = opk["id"]
                bundle["opk_pub"] = opk["public_key"]
            device_bundles.append(bundle)
        conn.commit()

        return {
            "identity_key": user["identity_key"],
            "device_bundles": device_bundles,
        }
    except Exception:
        try:
            conn.rollback()
        except Exception:
            pass
        raise
    finally:
        conn.close()


# --- Conversations ---

def create_conversation(member_user_ids: list[str], joined_at=None, name=None, created_by=None) -> str:
    conn = get_connection()
    try:
        cursor = conn.cursor()
        conv_id = generate_uuid()
        cursor.execute("INSERT INTO conversations (id, name, created_by) VALUES (%s, %s, %s)",
                       (conv_id, name, created_by))
        for uid in member_user_ids:
            cursor.execute(
                "INSERT INTO conversation_members (conversation_id, user_id, joined_at) VALUES (%s, %s, %s)",
                (conv_id, uid, joined_at),
            )
        conn.commit()
        return conv_id
    finally:
        conn.close()


def add_conversation_member(conversation_id: str, user_id: str, joined_at=None):
    conn = get_connection()
    try:
        cursor = conn.cursor()
        cursor.execute(
            "INSERT IGNORE INTO conversation_members (conversation_id, user_id, joined_at) VALUES (%s, %s, %s)",
            (conversation_id, user_id, joined_at),
        )
        conn.commit()
    finally:
        conn.close()


def remove_conversation_member(conversation_id: str, user_id: str):
    conn = get_connection()
    try:
        cursor = conn.cursor()
        cursor.execute(
            "DELETE FROM conversation_members WHERE conversation_id = %s AND user_id = %s",
            (conversation_id, user_id),
        )
        conn.commit()
    finally:
        conn.close()


def count_conversation_members(conversation_id: str) -> int:
    """Count members in a conversation."""
    conn = get_connection()
    try:
        cursor = conn.cursor()
        cursor.execute(
            "SELECT COUNT(*) FROM conversation_members WHERE conversation_id = %s",
            (conversation_id,),
        )
        return cursor.fetchone()[0]
    finally:
        conn.close()


def get_conversation_file_ids(conversation_id: str) -> list[str]:
    """Get all file IDs (images + files) uploaded to a conversation."""
    conn = get_connection()
    try:
        cursor = conn.cursor()
        cursor.execute(
            "SELECT file_id FROM image_uploads WHERE conversation_id = %s",
            (conversation_id,),
        )
        return [row[0] for row in cursor.fetchall()]
    finally:
        conn.close()


def delete_conversation(conversation_id: str):
    """Delete a conversation entirely. CASCADE cleans up members, messages, etc."""
    conn = get_connection()
    try:
        cursor = conn.cursor()
        cursor.execute("DELETE FROM conversations WHERE id = %s", (conversation_id,))
        conn.commit()
    finally:
        conn.close()


def get_conversation_members(conversation_id: str) -> list[dict]:
    conn = get_connection()
    try:
        cursor = conn.cursor(dictionary=True)
        cursor.execute(
            "SELECT u.id, u.username, u.email, u.identity_key FROM conversation_members cm "
            "JOIN users u ON cm.user_id = u.id "
            "WHERE cm.conversation_id = %s",
            (conversation_id,),
        )
        return cursor.fetchall()
    finally:
        conn.close()


def find_direct_conversation(user_id_a: str, user_id_b: str) -> str | None:
    conn = get_connection()
    try:
        cursor = conn.cursor()
        cursor.execute(
            "SELECT cm1.conversation_id FROM conversation_members cm1 "
            "JOIN conversation_members cm2 ON cm1.conversation_id = cm2.conversation_id "
            "WHERE cm1.user_id = %s AND cm2.user_id = %s "
            "AND (SELECT COUNT(*) FROM conversation_members cm3 "
            "     WHERE cm3.conversation_id = cm1.conversation_id) = 2 "
            "LIMIT 1",
            (user_id_a, user_id_b),
        )
        row = cursor.fetchone()
        return row[0] if row else None
    finally:
        conn.close()


def update_conversation_creator(conversation_id: str, new_creator_id: str):
    """Transfer group creator role to another member."""
    conn = get_connection()
    try:
        cursor = conn.cursor()
        cursor.execute(
            "UPDATE conversations SET created_by = %s WHERE id = %s",
            (new_creator_id, conversation_id),
        )
        conn.commit()
    finally:
        conn.close()


def get_conversation(conversation_id: str) -> dict | None:
    """Get conversation by ID."""
    conn = get_connection()
    try:
        cursor = conn.cursor(dictionary=True)
        cursor.execute(
            "SELECT id, created_at, name, created_by, avatar_file FROM conversations WHERE id = %s",
            (conversation_id,),
        )
        return cursor.fetchone()
    finally:
        conn.close()


def update_conversation_avatar(conversation_id: str, avatar_file: str):
    """Set avatar file for a conversation."""
    conn = get_connection()
    try:
        cursor = conn.cursor()
        cursor.execute(
            "UPDATE conversations SET avatar_file = %s WHERE id = %s",
            (avatar_file, conversation_id),
        )
        conn.commit()
    finally:
        conn.close()


def update_conversation_name(conversation_id: str, name: str):
    """Update the name of a conversation."""
    conn = get_connection()
    try:
        cursor = conn.cursor()
        cursor.execute(
            "UPDATE conversations SET name = %s WHERE id = %s",
            (name, conversation_id),
        )
        conn.commit()
    finally:
        conn.close()


def is_conversation_member(conversation_id: str, user_id: str) -> bool:
    conn = get_connection()
    try:
        cursor = conn.cursor()
        cursor.execute(
            "SELECT 1 FROM conversation_members WHERE conversation_id = %s AND user_id = %s",
            (conversation_id, user_id),
        )
        return cursor.fetchone() is not None
    finally:
        conn.close()


def list_user_conversations(user_id: str) -> list[dict]:
    conn = get_connection()
    try:
        cursor = conn.cursor(dictionary=True)
        cursor.execute(
            "SELECT c.id, c.created_at, c.name, c.created_by, c.avatar_file FROM conversations c "
            "JOIN conversation_members cm ON c.id = cm.conversation_id "
            "WHERE cm.user_id = %s ORDER BY c.created_at DESC",
            (user_id,),
        )
        convs = cursor.fetchall()
        for conv in convs:
            cursor.execute(
                "SELECT u.id AS user_id, u.username, u.email FROM conversation_members cm "
                "JOIN users u ON cm.user_id = u.id "
                "WHERE cm.conversation_id = %s",
                (conv["id"],),
            )
            conv["members"] = cursor.fetchall()
        return convs
    finally:
        conn.close()


# --- Group Invitations ---

def create_invitation(conversation_id: str, user_id: str, invited_by: str):
    """Create a pending group invitation."""
    conn = get_connection()
    try:
        cursor = conn.cursor()
        inv_id = generate_uuid()
        cursor.execute(
            "INSERT IGNORE INTO group_invitations (id, conversation_id, user_id, invited_by) "
            "VALUES (%s, %s, %s, %s)",
            (inv_id, conversation_id, user_id, invited_by),
        )
        conn.commit()
    finally:
        conn.close()


def get_pending_invitations(user_id: str) -> list[dict]:
    """Get all pending invitations for a user, joined with conversation and inviter info."""
    conn = get_connection()
    try:
        cursor = conn.cursor(dictionary=True)
        cursor.execute(
            "SELECT gi.id, gi.conversation_id, gi.invited_by, gi.created_at, "
            "c.name AS conversation_name, u.username AS invited_by_username "
            "FROM group_invitations gi "
            "JOIN conversations c ON gi.conversation_id = c.id "
            "JOIN users u ON gi.invited_by = u.id "
            "WHERE gi.user_id = %s "
            "ORDER BY gi.created_at DESC",
            (user_id,),
        )
        return cursor.fetchall()
    finally:
        conn.close()


def delete_invitation(conversation_id: str, user_id: str):
    """Delete a pending invitation."""
    conn = get_connection()
    try:
        cursor = conn.cursor()
        cursor.execute(
            "DELETE FROM group_invitations WHERE conversation_id = %s AND user_id = %s",
            (conversation_id, user_id),
        )
        conn.commit()
    finally:
        conn.close()


def has_pending_invitation(conversation_id: str, user_id: str) -> bool:
    """Check if a user has a pending invitation for a conversation."""
    conn = get_connection()
    try:
        cursor = conn.cursor()
        cursor.execute(
            "SELECT 1 FROM group_invitations WHERE conversation_id = %s AND user_id = %s",
            (conversation_id, user_id),
        )
        return cursor.fetchone() is not None
    finally:
        conn.close()


# --- Messages ---

def store_message(
    conversation_id: str,
    sender_id: str,
    ratchet_header: bytes,
    recipients: list[dict],
    x3dh_header: bytes | None = None,
    sender_chain_id: bytes | None = None,
    sender_chain_n: int | None = None,
    image_file_id: str | None = None,
    sender_device_id: str | None = None,
) -> str:
    """Store an encrypted message with per-recipient ciphertext.

    recipients: [{user_id, encrypted_content (bytes), nonce (bytes),
                  device_id (str, optional), ratchet_header (bytes, optional),
                  x3dh_header (bytes, optional)}]
    """
    conn = get_connection()
    try:
        cursor = conn.cursor()
        msg_id = generate_uuid()
        cursor.execute(
            "INSERT INTO messages (id, conversation_id, sender_id, sender_device_id, "
            "ratchet_header, x3dh_header, sender_chain_id, sender_chain_n, image_file_id) "
            "VALUES (%s, %s, %s, %s, %s, %s, %s, %s, %s)",
            (msg_id, conversation_id, sender_id, sender_device_id, ratchet_header,
             x3dh_header, sender_chain_id, sender_chain_n, image_file_id),
        )
        for r in recipients:
            device_id = r.get("device_id", SELF_DEVICE_ID)
            cursor.execute(
                "INSERT INTO message_recipients (message_id, user_id, device_id, "
                "encrypted_content, nonce, ratchet_header, x3dh_header) "
                "VALUES (%s, %s, %s, %s, %s, %s, %s)",
                (msg_id, r["user_id"], device_id, r["encrypted_content"], r["nonce"],
                 r.get("ratchet_header"), r.get("x3dh_header")),
            )
        conn.commit()
        return msg_id
    finally:
        conn.close()


def get_messages(conversation_id: str, user_id: str, limit: int = 50, offset: int = 0,
                 device_id: str | None = None) -> list[dict]:
    """Get messages for a user in a conversation, JOINing their per-recipient ciphertext.

    If device_id is set, returns rows where mr.device_id matches OR is the sentinel
    (self-encrypted / legacy). This ensures both device-specific and self-encrypted
    copies are returned.
    """
    conn = get_connection()
    try:
        cursor = conn.cursor(dictionary=True)
        if device_id:
            cursor.execute(
                "SELECT m.id, m.conversation_id, m.sender_id, m.sender_device_id, "
                "m.ratchet_header, m.x3dh_header, "
                "m.sender_chain_id, m.sender_chain_n, m.created_at, m.deleted_at, m.image_file_id, "
                "mr.encrypted_content, mr.nonce, mr.device_id AS mr_device_id, "
                "mr.ratchet_header AS mr_ratchet_header, mr.x3dh_header AS mr_x3dh_header "
                "FROM messages m "
                "JOIN message_recipients mr ON m.id = mr.message_id AND mr.user_id = %s "
                "  AND (mr.device_id = %s OR mr.device_id = %s) "
                "JOIN conversation_members cm ON cm.conversation_id = m.conversation_id AND cm.user_id = %s "
                "WHERE m.conversation_id = %s AND (cm.joined_at IS NULL OR m.created_at >= cm.joined_at) "
                "ORDER BY m.created_at DESC LIMIT %s OFFSET %s",
                (user_id, device_id, SELF_DEVICE_ID, user_id, conversation_id, limit, offset),
            )
        else:
            cursor.execute(
                "SELECT m.id, m.conversation_id, m.sender_id, m.sender_device_id, "
                "m.ratchet_header, m.x3dh_header, "
                "m.sender_chain_id, m.sender_chain_n, m.created_at, m.deleted_at, m.image_file_id, "
                "mr.encrypted_content, mr.nonce, mr.device_id AS mr_device_id, "
                "mr.ratchet_header AS mr_ratchet_header, mr.x3dh_header AS mr_x3dh_header "
                "FROM messages m "
                "JOIN message_recipients mr ON m.id = mr.message_id AND mr.user_id = %s "
                "JOIN conversation_members cm ON cm.conversation_id = m.conversation_id AND cm.user_id = %s "
                "WHERE m.conversation_id = %s AND (cm.joined_at IS NULL OR m.created_at >= cm.joined_at) "
                "ORDER BY m.created_at DESC LIMIT %s OFFSET %s",
                (user_id, user_id, conversation_id, limit, offset),
            )
        return cursor.fetchall()
    finally:
        conn.close()


def get_message_conversation(message_id: str) -> str | None:
    conn = get_connection()
    try:
        cursor = conn.cursor()
        cursor.execute("SELECT conversation_id FROM messages WHERE id = %s", (message_id,))
        row = cursor.fetchone()
        return row[0] if row else None
    finally:
        conn.close()


def get_message_sender(message_id: str) -> str | None:
    conn = get_connection()
    try:
        cursor = conn.cursor()
        cursor.execute("SELECT sender_id FROM messages WHERE id = %s", (message_id,))
        row = cursor.fetchone()
        return row[0] if row else None
    finally:
        conn.close()


# --- Group Sender Keys ---

def store_sender_key(conversation_id: str, sender_id: str, chain_id: bytes,
                     device_id: str | None = None):
    """Store or update a sender key chain ID for a group member's device."""
    conn = get_connection()
    try:
        cursor = conn.cursor()
        dev = device_id or SELF_DEVICE_ID
        cursor.execute(
            "REPLACE INTO group_sender_keys (conversation_id, sender_id, device_id, chain_id) "
            "VALUES (%s, %s, %s, %s)",
            (conversation_id, sender_id, dev, chain_id),
        )
        conn.commit()
    finally:
        conn.close()


def get_sender_key(conversation_id: str, sender_id: str,
                   device_id: str | None = None) -> dict | None:
    conn = get_connection()
    try:
        cursor = conn.cursor(dictionary=True)
        dev = device_id or SELF_DEVICE_ID
        cursor.execute(
            "SELECT chain_id, created_at FROM group_sender_keys "
            "WHERE conversation_id = %s AND sender_id = %s AND device_id = %s",
            (conversation_id, sender_id, dev),
        )
        return cursor.fetchone()
    finally:
        conn.close()


# --- Read Receipts ---

def mark_messages_read(conversation_id: str, user_id: str, message_ids: list[str]):
    if not message_ids:
        return
    conn = get_connection()
    try:
        cursor = conn.cursor()
        for mid in message_ids:
            cursor.execute(
                "INSERT IGNORE INTO message_reads (message_id, user_id) VALUES (%s, %s)",
                (mid, user_id),
            )
        conn.commit()
    finally:
        conn.close()


def get_unread_counts(user_id: str) -> dict[str, int]:
    """Return {conversation_id: unread_count} for all conversations the user is in."""
    conn = get_connection()
    try:
        cursor = conn.cursor(dictionary=True)
        cursor.execute(
            "SELECT m.conversation_id, COUNT(*) AS cnt "
            "FROM messages m "
            "JOIN message_recipients mr ON mr.message_id = m.id AND mr.user_id = %s "
            "LEFT JOIN message_reads mrd ON mrd.message_id = m.id AND mrd.user_id = %s "
            "WHERE m.sender_id != %s AND m.deleted_at IS NULL AND mrd.message_id IS NULL "
            "GROUP BY m.conversation_id",
            (user_id, user_id, user_id),
        )
        return {row["conversation_id"]: row["cnt"] for row in cursor.fetchall()}
    finally:
        conn.close()


def get_message_read_status(message_ids: list[str]) -> dict:
    if not message_ids:
        return {}
    conn = get_connection()
    try:
        cursor = conn.cursor(dictionary=True)
        placeholders = ",".join(["%s"] * len(message_ids))
        cursor.execute(
            f"SELECT mr.message_id, mr.user_id, mr.read_at "
            f"FROM message_reads mr "
            f"WHERE mr.message_id IN ({placeholders})",
            tuple(message_ids),
        )
        result = {}
        for row in cursor.fetchall():
            mid = row["message_id"]
            if mid not in result:
                result[mid] = []
            result[mid].append({
                "user_id": row["user_id"],
                "read_at": row["read_at"].isoformat() if hasattr(row["read_at"], "isoformat") else str(row["read_at"]),
            })
        return result
    finally:
        conn.close()


# --- Delete ---

def soft_delete_message(message_id: str, sender_id: str) -> dict | None:
    """Soft-delete a message if sender matches. Returns {'image_file_id': ...} or None."""
    conn = get_connection()
    try:
        cursor = conn.cursor(dictionary=True)
        cursor.execute(
            "SELECT sender_id, image_file_id FROM messages WHERE id = %s AND deleted_at IS NULL",
            (message_id,),
        )
        row = cursor.fetchone()
        if not row or row["sender_id"] != sender_id:
            return None
        cursor.execute(
            "UPDATE messages SET deleted_at = NOW() WHERE id = %s",
            (message_id,),
        )
        # Clear per-recipient ciphertext
        cursor.execute(
            "UPDATE message_recipients SET encrypted_content = %s WHERE message_id = %s",
            (b"", message_id),
        )
        conn.commit()
        return {"image_file_id": row.get("image_file_id")}
    finally:
        conn.close()


def set_message_image_file_id(message_id: str, file_id: str):
    conn = get_connection()
    try:
        cursor = conn.cursor()
        cursor.execute(
            "UPDATE messages SET image_file_id = %s WHERE id = %s",
            (file_id, message_id),
        )
        conn.commit()
    finally:
        conn.close()


# --- Image Uploads ---

def create_image_upload(file_id: str, conversation_id: str, uploader_id: str, file_size: int):
    conn = get_connection()
    try:
        cursor = conn.cursor()
        cursor.execute(
            "INSERT INTO image_uploads (file_id, conversation_id, uploader_id, file_size) "
            "VALUES (%s, %s, %s, %s)",
            (file_id, conversation_id, uploader_id, file_size),
        )
        conn.commit()
    finally:
        conn.close()


def complete_image_upload(file_id: str):
    conn = get_connection()
    try:
        cursor = conn.cursor()
        cursor.execute(
            "UPDATE image_uploads SET completed = TRUE WHERE file_id = %s",
            (file_id,),
        )
        conn.commit()
    finally:
        conn.close()


def get_image_upload(file_id: str) -> dict | None:
    conn = get_connection()
    try:
        cursor = conn.cursor(dictionary=True)
        cursor.execute(
            "SELECT file_id, conversation_id, uploader_id, file_size, completed, created_at "
            "FROM image_uploads WHERE file_id = %s",
            (file_id,),
        )
        return cursor.fetchone()
    finally:
        conn.close()


def delete_image_upload(file_id: str):
    conn = get_connection()
    try:
        cursor = conn.cursor()
        cursor.execute("DELETE FROM image_uploads WHERE file_id = %s", (file_id,))
        conn.commit()
    finally:
        conn.close()


# --- User Profiles ---

def create_default_profile(user_id: str):
    """Create a default profile for a new user."""
    conn = get_connection()
    try:
        cursor = conn.cursor()
        cursor.execute(
            "INSERT IGNORE INTO user_profiles (user_id) VALUES (%s)",
            (user_id,),
        )
        conn.commit()
    finally:
        conn.close()


def get_user_profile(user_id: str, viewer_id: str | None = None) -> dict | None:
    """Get user profile joined with user info. Respects visibility if viewer is different user."""
    conn = get_connection()
    try:
        cursor = conn.cursor(dictionary=True)
        cursor.execute(
            "SELECT u.id AS user_id, u.username, u.email, u.created_at, "
            "p.phone, p.phone_visible, p.email_visible, p.location, "
            "p.location_visible, p.avatar_file, p.updated_at "
            "FROM users u LEFT JOIN user_profiles p ON u.id = p.user_id "
            "WHERE u.id = %s",
            (user_id,),
        )
        row = cursor.fetchone()
        if not row:
            return None
        # If viewing someone else's profile, apply visibility rules
        if viewer_id and viewer_id != user_id:
            if not row.get("email_visible"):
                row["email"] = None
            if not row.get("phone_visible"):
                row["phone"] = None
            if not row.get("location_visible"):
                row["location"] = None
        return row
    finally:
        conn.close()


def update_user_profile(user_id: str, **fields):
    """Upsert user profile fields. Allowed: phone, phone_visible, email_visible,
    location, location_visible, avatar_file."""
    allowed = {"phone", "phone_visible", "email_visible", "location",
               "location_visible", "avatar_file"}
    filtered = {k: v for k, v in fields.items() if k in allowed}
    if not filtered:
        return
    conn = get_connection()
    try:
        cursor = conn.cursor()
        # Upsert: insert default then update
        cursor.execute(
            "INSERT IGNORE INTO user_profiles (user_id) VALUES (%s)",
            (user_id,),
        )
        set_clause = ", ".join(f"{k} = %s" for k in filtered)
        values = list(filtered.values()) + [user_id]
        cursor.execute(
            f"UPDATE user_profiles SET {set_clause} WHERE user_id = %s",
            values,
        )
        conn.commit()
    finally:
        conn.close()


def batch_reencrypt_messages(user_id: str, updates: list[dict]):
    """Batch re-encrypt message_recipients rows with self-encryption key data.

    Each update: {message_id, encrypted_content (bytes), nonce (bytes)}.
    Sets ratchet_header to '{"self":true}' and clears x3dh_header.
    Only updates rows belonging to user_id with sentinel device_id (self-encrypted copies).
    """
    if not updates:
        return
    conn = get_connection()
    try:
        cursor = conn.cursor()
        self_header = b'{"self":true}'
        for u in updates:
            cursor.execute(
                "UPDATE message_recipients "
                "SET encrypted_content = %s, nonce = %s, ratchet_header = %s, x3dh_header = NULL "
                "WHERE message_id = %s AND user_id = %s AND device_id = %s",
                (u["encrypted_content"], u["nonce"], self_header, u["message_id"],
                 user_id, SELF_DEVICE_ID),
            )
        conn.commit()
    finally:
        conn.close()


# --- Phantom Users ---

def create_phantom_user(email: str) -> dict:
    """Create a phantom user with valid crypto keys for X3DH.

    Phantom users have rsa_public_key = 'PHANTOM' as a marker.
    Returns user dict: {id, username, email, identity_key}.
    """
    username = email.split("@")[0]
    user_id = generate_uuid()

    # Generate real crypto keys so X3DH works on the client side
    ik_private, ik_public = generate_identity_keypair()
    ik_public_bytes = serialize_ed25519_public(ik_public)

    spk = generate_signed_prekey(ik_private)
    spk_pub_bytes = serialize_x25519_public(spk["public"])
    spk_sig = spk["signature"]

    opks = generate_one_time_prekeys(count=5)

    conn = get_connection()
    try:
        cursor = conn.cursor()
        cursor.execute(
            "INSERT INTO users (id, username, email, rsa_public_key, identity_key) "
            "VALUES (%s, %s, %s, %s, %s)",
            (user_id, username, email, "PHANTOM", ik_public_bytes),
        )
        cursor.execute(
            "INSERT INTO signed_prekeys (id, user_id, public_key, signature) VALUES (%s, %s, %s, %s)",
            (spk["id"], user_id, spk_pub_bytes, spk_sig),
        )
        for opk in opks:
            cursor.execute(
                "INSERT INTO one_time_prekeys (id, user_id, public_key) VALUES (%s, %s, %s)",
                (opk["id"], user_id, serialize_x25519_public(opk["public"])),
            )
        conn.commit()
        return {"id": user_id, "username": username, "email": email, "identity_key": ik_public_bytes}
    finally:
        conn.close()


def is_phantom_user(user_id: str) -> bool:
    """Check if a user is a phantom (rsa_public_key == 'PHANTOM')."""
    conn = get_connection()
    try:
        cursor = conn.cursor()
        cursor.execute("SELECT rsa_public_key FROM users WHERE id = %s", (user_id,))
        row = cursor.fetchone()
        return row is not None and row[0] == "PHANTOM"
    finally:
        conn.close()


def delete_phantom_user(user_id: str):
    """Delete a phantom user. CASCADE removes signed_prekeys, one_time_prekeys,
    conversation_members, message_recipients, etc."""
    conn = get_connection()
    try:
        cursor = conn.cursor()
        cursor.execute(
            "DELETE FROM users WHERE id = %s AND rsa_public_key = %s",
            (user_id, "PHANTOM"),
        )
        conn.commit()
    finally:
        conn.close()


def upgrade_phantom_user(phantom_id: str, username: str, rsa_public_key_pem: str,
                         identity_key: bytes) -> str | None:
    """Upgrade a phantom user to a real user in-place.

    Preserves user_id and all FK references (conversation_members, group_invitations, etc.).
    Deletes phantom's server-generated prekeys (real user will upload own on first login).
    Returns phantom_id as the new user_id, or None if phantom no longer exists.
    """
    conn = get_connection()
    try:
        cursor = conn.cursor()
        cursor.execute(
            "UPDATE users SET username = %s, rsa_public_key = %s, identity_key = %s "
            "WHERE id = %s AND rsa_public_key = 'PHANTOM'",
            (username, rsa_public_key_pem, identity_key, phantom_id),
        )
        if cursor.rowcount == 0:
            conn.rollback()
            return None
        # Remove phantom's server-generated crypto keys — real user uploads own
        cursor.execute("DELETE FROM signed_prekeys WHERE user_id = %s", (phantom_id,))
        cursor.execute("DELETE FROM one_time_prekeys WHERE user_id = %s", (phantom_id,))
        conn.commit()
        return phantom_id
    finally:
        conn.close()


def get_all_phantom_user_ids() -> set[str]:
    """Return set of all phantom user IDs (for server startup cache)."""
    conn = get_connection()
    try:
        cursor = conn.cursor()
        cursor.execute("SELECT id FROM users WHERE rsa_public_key = %s", ("PHANTOM",))
        return {row[0] for row in cursor.fetchall()}
    finally:
        conn.close()


def cleanup_stale_phantoms(max_age_days: int = 30) -> int:
    """Delete phantom users older than max_age_days with no active conversations with real users."""
    conn = get_connection()
    try:
        cursor = conn.cursor()
        # Two-step: SELECT ids first, then DELETE.
        # MySQL error 1093: can't DELETE from table referenced in subquery.
        cursor.execute("""
            SELECT u.id FROM users u
            WHERE u.rsa_public_key = 'PHANTOM'
              AND u.created_at < DATE_SUB(NOW(), INTERVAL %s DAY)
              AND NOT EXISTS (
                  SELECT 1 FROM conversation_members cm1
                  JOIN conversation_members cm2 ON cm1.conversation_id = cm2.conversation_id
                  JOIN users u2 ON cm2.user_id = u2.id
                  WHERE cm1.user_id = u.id
                    AND u2.rsa_public_key != 'PHANTOM'
              )
        """, (max_age_days,))
        ids = [row[0] for row in cursor.fetchall()]
        if not ids:
            return 0
        cursor.execute(
            "DELETE FROM users WHERE id IN (%s)" % ",".join(["%s"] * len(ids)),
            ids,
        )
        deleted = cursor.rowcount
        conn.commit()
        return deleted
    finally:
        conn.close()


def remove_conversation_member_atomic(conversation_id: str, user_id: str) -> bool:
    """Remove member and return True if actually removed (row existed). M6 TOCTOU fix."""
    conn = get_connection()
    try:
        cursor = conn.cursor()
        cursor.execute(
            "DELETE FROM conversation_members WHERE conversation_id = %s AND user_id = %s",
            (conversation_id, user_id),
        )
        conn.commit()
        return cursor.rowcount > 0
    finally:
        conn.close()


def get_stale_uploads(max_age_seconds: int = 3600) -> list[dict]:
    conn = get_connection()
    try:
        cursor = conn.cursor(dictionary=True)
        cursor.execute(
            "SELECT file_id FROM image_uploads "
            "WHERE completed = FALSE AND created_at < DATE_SUB(NOW(), INTERVAL %s SECOND)",
            (max_age_seconds,),
        )
        return cursor.fetchall()
    finally:
        conn.close()