import Cookies from "js-cookie";
import { getApiBaseUrl } from "@/components/main/BaseUrl";
import {
  decryptText,
  encryptText,
  generateRoomKeyRaw,
  importPublicKeySpkiB64,
  importRoomAesKey,
  unwrapRoomKey,
  wrapRoomKeyForRecipient,
} from "./crypto";
import { E2EE_LOCKED_PREVIEW, isE2eeEnvelope } from "./envelope";
import { ensureIdentityKeys } from "./keys";
import { b64ToBuf, bufToB64, idbGet, idbSet } from "./storage";

type RoomKeysResponse = {
  e2ee_enabled: boolean;
  myWrappedKey: string | null;
  memberPublicKeys: { userId: string; publicKey: string | null }[];
  missingWrapUserIds: string[];
};

const roomAesCache = new Map<string, CryptoKey>();
const roomRawCache = new Map<string, Uint8Array>();

function authHeaders(): HeadersInit {
  const token = Cookies.get("token");
  return {
    "Content-Type": "application/json",
    ...(token ? { Authorization: `Bearer ${token}` } : {}),
  };
}

async function fetchRoomKeys(roomId: string): Promise<RoomKeysResponse | null> {
  const base = getApiBaseUrl();
  const res = await fetch(`${base}/chat-rooms/${roomId}/e2ee-keys`, {
    headers: authHeaders(),
  });
  if (!res.ok) return null;
  return (await res.json()) as RoomKeysResponse;
}

async function putRoomKeys(
  roomId: string,
  body: {
    e2ee_enabled?: boolean;
    wrapped_keys?: { userId: string; wrappedKey: string }[];
  }
): Promise<void> {
  const base = getApiBaseUrl();
  await fetch(`${base}/chat-rooms/${roomId}/e2ee-keys`, {
    method: "PUT",
    headers: authHeaders(),
    body: JSON.stringify(body),
  });
}

async function loadCachedRoomRaw(
  myUserId: string,
  roomId: string
): Promise<Uint8Array | null> {
  const mem = roomRawCache.get(roomId);
  if (mem) return mem;
  const b64 = await idbGet<string>(`roomkey:${myUserId}:${roomId}`);
  if (!b64) return null;
  const raw = new Uint8Array(b64ToBuf(b64));
  roomRawCache.set(roomId, raw);
  return raw;
}

async function persistRoomRaw(
  myUserId: string,
  roomId: string,
  raw: Uint8Array
): Promise<CryptoKey> {
  roomRawCache.set(roomId, raw);
  await idbSet(`roomkey:${myUserId}:${roomId}`, bufToB64(raw));
  const aes = await importRoomAesKey(raw);
  roomAesCache.set(roomId, aes);
  return aes;
}

export async function ensureRoomAesKey(
  myUserId: string,
  roomId: string
): Promise<CryptoKey | null> {
  const cached = roomAesCache.get(roomId);
  if (cached) return cached;

  const { privateKey } = await ensureIdentityKeys(myUserId);
  const info = await fetchRoomKeys(roomId);
  if (!info) return null;

  let raw = await loadCachedRoomRaw(myUserId, roomId);

  if (!raw && info.myWrappedKey) {
    try {
      raw = await unwrapRoomKey(info.myWrappedKey, privateKey);
      await persistRoomRaw(myUserId, roomId, raw);
    } catch {
      raw = null;
    }
  }

  if (!raw && !info.e2ee_enabled) {
    raw = await generateRoomKeyRaw();
    await persistRoomRaw(myUserId, roomId, raw);
    const wraps: { userId: string; wrappedKey: string }[] = [];
    for (const m of info.memberPublicKeys) {
      if (!m.publicKey) continue;
      try {
        const pub = await importPublicKeySpkiB64(m.publicKey);
        wraps.push({
          userId: m.userId,
          wrappedKey: await wrapRoomKeyForRecipient(raw, pub),
        });
      } catch {
        /* skip */
      }
    }
    if (!wraps.some((w) => w.userId === myUserId)) {
      const myPub = info.memberPublicKeys.find((m) => m.userId === myUserId);
      if (myPub?.publicKey) {
        const pub = await importPublicKeySpkiB64(myPub.publicKey);
        wraps.push({
          userId: myUserId,
          wrappedKey: await wrapRoomKeyForRecipient(raw, pub),
        });
      }
    }
    await putRoomKeys(roomId, { e2ee_enabled: true, wrapped_keys: wraps });
    return roomAesCache.get(roomId) || (await persistRoomRaw(myUserId, roomId, raw));
  }

  if (raw) {
    // Sync wraps for members missing keys
    if (info.missingWrapUserIds?.length) {
      const wraps: { userId: string; wrappedKey: string }[] = [];
      for (const uid of info.missingWrapUserIds) {
        const m = info.memberPublicKeys.find((x) => x.userId === uid);
        if (!m?.publicKey) continue;
        try {
          const pub = await importPublicKeySpkiB64(m.publicKey);
          wraps.push({
            userId: uid,
            wrappedKey: await wrapRoomKeyForRecipient(raw, pub),
          });
        } catch {
          /* skip */
        }
      }
      if (wraps.length) {
        await putRoomKeys(roomId, { e2ee_enabled: true, wrapped_keys: wraps });
      }
    }
    return persistRoomRaw(myUserId, roomId, raw);
  }

  return null;
}

export async function encryptRoomContent(
  myUserId: string,
  roomId: string,
  plaintext: string
): Promise<string | null> {
  if (!plaintext) return plaintext;
  if (isE2eeEnvelope(plaintext)) return plaintext;
  const key = await ensureRoomAesKey(myUserId, roomId);
  if (!key) return null;
  return encryptText(key, plaintext);
}

export async function decryptRoomContent(
  myUserId: string,
  roomId: string,
  content: string
): Promise<string> {
  if (!isE2eeEnvelope(content)) return content;
  const key = await ensureRoomAesKey(myUserId, roomId);
  if (!key) return E2EE_LOCKED_PREVIEW;
  const pt = await decryptText(key, content);
  return pt ?? E2EE_LOCKED_PREVIEW;
}

export async function decryptRoomMessageList<
  T extends {
    content?: string;
    replyTo?: { content?: string; [k: string]: unknown } | null;
  }
>(
  myUserId: string,
  roomId: string,
  messages: T[]
): Promise<T[]> {
  await ensureRoomAesKey(myUserId, roomId);
  const out: T[] = [];
  for (const msg of messages) {
    let next = msg;
    if (msg.content && isE2eeEnvelope(msg.content)) {
      const decrypted = await decryptRoomContent(myUserId, roomId, msg.content);
      next = { ...next, content: decrypted };
    }
    if (next.replyTo?.content && isE2eeEnvelope(next.replyTo.content)) {
      const replyContent = await decryptRoomContent(
        myUserId,
        roomId,
        next.replyTo.content
      );
      next = {
        ...next,
        replyTo: { ...next.replyTo, content: replyContent },
      };
    }
    out.push(next);
  }
  return out;
}
