import { query, mutation } from "./_generated/server";
import { v } from "convex/values";
import { hasPermission } from "./roles";
import { AUDIT_ACTIONS, logAudit } from "./audit";
import { getPublicStorageUrl } from "./storageUrl";

/**
 * Rotate the symmetric key for a DM channel. Inserts a brand-new
 * versioned row for each participant — existing rows are left alone
 * so previously-encrypted messages remain decryptable.
 *
 * The caller proves they're a DM participant by passing their own
 * userId; the server cross-checks against `dmParticipants` for the
 * channel. Every recipient userId in `entries` must also be a
 * participant — no leaking keys to random users.
 *
 * The new rows are tagged with `maxExistingVersion + 1`.
 */
export const rotateDMKey = mutation({
  args: {
    channelId: v.id("channels"),
    initiatorUserId: v.id("userProfiles"),
    entries: v.array(
      v.object({
        userId: v.id("userProfiles"),
        encryptedKeyBundle: v.string(),
      }),
    ),
  },
  returns: v.object({ keyVersion: v.number() }),
  handler: async (ctx, args) => {
    const channel = await ctx.db.get(args.channelId);
    if (!channel) throw new Error("Channel not found");
    if (channel.type !== "dm") {
      throw new Error("rotateDMKey is only supported for DM channels");
    }

    // Verify every (initiator + entries) userId is in dmParticipants.
    const participants = await ctx.db
      .query("dmParticipants")
      .withIndex("by_channel", (q) => q.eq("channelId", args.channelId))
      .collect();
    const participantSet = new Set(participants.map((p) => p.userId as string));
    if (!participantSet.has(args.initiatorUserId as unknown as string)) {
      throw new Error("Not a participant in this DM");
    }
    for (const entry of args.entries) {
      if (!participantSet.has(entry.userId as unknown as string)) {
        throw new Error("Target userId is not a participant in this DM");
      }
    }

    // Find the current max keyVersion for this channel. New rows go
    // one above that. If no rows exist yet, start at 2 so legacy
    // messages tagged version 1 still hit their original key.
    const existing = await ctx.db
      .query("channelKeys")
      .withIndex("by_channel", (q) => q.eq("channelId", args.channelId))
      .collect();
    const maxVersion = existing.reduce(
      (m, k) => (k.keyVersion > m ? k.keyVersion : m),
      0,
    );
    const newVersion = maxVersion + 1;

    for (const entry of args.entries) {
      await ctx.db.insert("channelKeys", {
        channelId: args.channelId,
        userId: entry.userId,
        encryptedKeyBundle: entry.encryptedKeyBundle,
        keyVersion: newVersion,
      });
    }

    return { keyVersion: newVersion };
  },
});

// Batch upsert encrypted key bundles
export const uploadKeys = mutation({
  args: {
    keys: v.array(
      v.object({
        channelId: v.id("channels"),
        userId: v.id("userProfiles"),
        encryptedKeyBundle: v.string(),
        keyVersion: v.number(),
      })
    ),
  },
  returns: v.object({ success: v.boolean(), count: v.number() }),
  handler: async (ctx, args) => {
    for (const keyData of args.keys) {
      if (!keyData.channelId || !keyData.userId || !keyData.encryptedKeyBundle) {
        continue;
      }

      // Check if exists (upsert)
      const existing = await ctx.db
        .query("channelKeys")
        .withIndex("by_channel_and_user", (q) =>
          q.eq("channelId", keyData.channelId).eq("userId", keyData.userId)
        )
        .unique();

      if (existing) {
        await ctx.db.patch(existing._id, {
          encryptedKeyBundle: keyData.encryptedKeyBundle,
          keyVersion: keyData.keyVersion,
        });
      } else {
        await ctx.db.insert("channelKeys", {
          channelId: keyData.channelId,
          userId: keyData.userId,
          encryptedKeyBundle: keyData.encryptedKeyBundle,
          keyVersion: keyData.keyVersion,
        });
      }
    }

    return { success: true, count: args.keys.length };
  },
});

// Get user's encrypted key bundles (reactive!)
export const getKeysForUser = query({
  args: { userId: v.id("userProfiles") },
  returns: v.array(
    v.object({
      channel_id: v.id("channels"),
      encrypted_key_bundle: v.string(),
      key_version: v.number(),
    })
  ),
  handler: async (ctx, args) => {
    const keys = await ctx.db
      .query("channelKeys")
      .withIndex("by_user", (q) => q.eq("userId", args.userId))
      .collect();

    return keys.map((k) => ({
      channel_id: k.channelId,
      encrypted_key_bundle: k.encryptedKeyBundle,
      key_version: k.keyVersion,
    }));
  },
});

/**
 * Admin-only query: list users who don't have a `channelKeys` row for
 * the given (non-DM) channel. Powers the Channel Settings → Access
 * panel where an admin grants missing keys to users who joined via a
 * broken invite (only received one channel's key instead of all).
 *
 * The actor themselves is filtered out — you can't be missing your own
 * key from your own POV, and the UI never needs to grant to self.
 * Ghosts and users without a public key are filtered (can't log in
 * anyway / nothing to encrypt against).
 */
export const getUsersMissingChannelKey = query({
  args: {
    actorId: v.id("userProfiles"),
    channelId: v.id("channels"),
  },
  returns: v.array(
    v.object({
      userId: v.id("userProfiles"),
      username: v.string(),
      displayName: v.union(v.string(), v.null()),
      avatarUrl: v.union(v.string(), v.null()),
      userPublicKey: v.string(),
    }),
  ),
  handler: async (ctx, args) => {
    if (!(await hasPermission(ctx, args.actorId, "manage_channels"))) {
      throw new Error("Forbidden");
    }
    const channel = await ctx.db.get(args.channelId);
    if (!channel) throw new Error("Channel not found");
    if (channel.type === "dm") {
      throw new Error("grantChannelAccess is not supported for DM channels");
    }

    const existing = await ctx.db
      .query("channelKeys")
      .withIndex("by_channel", (q) => q.eq("channelId", args.channelId))
      .collect();
    const have = new Set(existing.map((k) => k.userId as unknown as string));

    const users = await ctx.db.query("userProfiles").collect();
    const missing = users.filter(
      (u) =>
        !!u.publicIdentityKey &&
        !u.isGhost &&
        (u._id as unknown as string) !== (args.actorId as unknown as string) &&
        !have.has(u._id as unknown as string),
    );

    const results = [];
    for (const u of missing) {
      let avatarUrl: string | null = null;
      if (u.avatarStorageId) {
        avatarUrl = await getPublicStorageUrl(ctx, u.avatarStorageId);
      }
      results.push({
        userId: u._id,
        username: u.username,
        displayName: u.displayName ?? null,
        avatarUrl,
        userPublicKey: u.publicIdentityKey,
      });
    }
    return results;
  },
});

/**
 * Admin-gated single-user grant. The admin's client has already
 * decrypted its own channel-key bundle and re-encrypted the key against
 * the target user's RSA public key — this mutation just upserts the
 * row, re-checks auth, and writes an audit entry. Kept separate from
 * `uploadKeys` on purpose: that path is used unauthenticated during
 * register/invite-accept for the caller's own keys, and loosening it
 * to accept cross-user writes would let any client grant themselves
 * access to any channel.
 */
export const grantChannelAccess = mutation({
  args: {
    actorId: v.id("userProfiles"),
    channelId: v.id("channels"),
    userId: v.id("userProfiles"),
    encryptedKeyBundle: v.string(),
    keyVersion: v.number(),
  },
  returns: v.object({ success: v.boolean() }),
  handler: async (ctx, args) => {
    if (!(await hasPermission(ctx, args.actorId, "manage_channels"))) {
      throw new Error("Forbidden");
    }
    const channel = await ctx.db.get(args.channelId);
    if (!channel) throw new Error("Channel not found");
    if (channel.type === "dm") {
      throw new Error("grantChannelAccess is not supported for DM channels");
    }
    const target = await ctx.db.get(args.userId);
    if (!target) throw new Error("Target user not found");
    if (target.isGhost || !target.publicIdentityKey) {
      throw new Error("Target user can't receive keys");
    }

    const existing = await ctx.db
      .query("channelKeys")
      .withIndex("by_channel_and_user", (q) =>
        q.eq("channelId", args.channelId).eq("userId", args.userId),
      )
      .unique();
    if (existing) {
      await ctx.db.patch(existing._id, {
        encryptedKeyBundle: args.encryptedKeyBundle,
        keyVersion: args.keyVersion,
      });
    } else {
      await ctx.db.insert("channelKeys", {
        channelId: args.channelId,
        userId: args.userId,
        encryptedKeyBundle: args.encryptedKeyBundle,
        keyVersion: args.keyVersion,
      });
    }

    await logAudit(ctx, {
      actorId: args.actorId,
      action: AUDIT_ACTIONS.KEYS_GRANT,
      targetType: "channel",
      targetId: args.channelId as unknown as string,
      targetName: channel.name,
      metadata: { userId: args.userId, username: target.username },
    });

    return { success: true };
  },
});