import Cookies from "js-cookie";
import { getApiBaseUrl } from "@/components/main/BaseUrl";
import {
  exportPrivateKeyJwk,
  exportPublicKeySpkiB64,
  generateIdentityKeyPair,
  importPrivateKeyJwk,
  importPublicKeySpkiB64,
} from "./crypto";
import { idbGet, idbSet } from "./storage";

type StoredIdentity = {
  userId: string;
  privateJwk: JsonWebKey;
  publicKeyB64: string;
};

let memoryPair: { userId: string; privateKey: CryptoKey; publicKeyB64: string } | null =
  null;

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

async function publishPublicKey(publicKeyB64: string): Promise<void> {
  const base = getApiBaseUrl();
  await fetch(`${base}/account/e2ee/public-key`, {
    method: "PUT",
    headers: authHeaders(),
    body: JSON.stringify({ publicKey: publicKeyB64 }),
  });
}

export async function ensureIdentityKeys(userId: string): Promise<{
  privateKey: CryptoKey;
  publicKeyB64: string;
}> {
  const uid = String(userId);
  if (memoryPair?.userId === uid) {
    return {
      privateKey: memoryPair.privateKey,
      publicKeyB64: memoryPair.publicKeyB64,
    };
  }

  const storeKey = `identity:${uid}`;
  const stored = await idbGet<StoredIdentity>(storeKey);

  if (stored?.privateJwk && stored.publicKeyB64) {
    const privateKey = await importPrivateKeyJwk(stored.privateJwk);
    memoryPair = {
      userId: uid,
      privateKey,
      publicKeyB64: stored.publicKeyB64,
    };
    try {
      await publishPublicKey(stored.publicKeyB64);
    } catch {
      /* publish best-effort */
    }
    return { privateKey, publicKeyB64: stored.publicKeyB64 };
  }

  const pair = await generateIdentityKeyPair();
  const publicKeyB64 = await exportPublicKeySpkiB64(pair.publicKey);
  const privateJwk = await exportPrivateKeyJwk(pair.privateKey);
  await idbSet(storeKey, {
    userId: uid,
    privateJwk,
    publicKeyB64,
  } satisfies StoredIdentity);

  memoryPair = { userId: uid, privateKey: pair.privateKey, publicKeyB64 };
  try {
    await publishPublicKey(publicKeyB64);
  } catch {
    /* publish best-effort */
  }
  return { privateKey: pair.privateKey, publicKeyB64 };
}

export async function fetchPeerPublicKey(
  peerUserId: string
): Promise<string | null> {
  const base = getApiBaseUrl();
  const res = await fetch(
    `${base}/account/e2ee/public-key/${encodeURIComponent(peerUserId)}`,
    { headers: authHeaders() }
  );
  if (!res.ok) return null;
  const data = (await res.json()) as { publicKey?: string | null };
  return data.publicKey || null;
}

export async function importPeerPublicKey(
  peerUserId: string
): Promise<CryptoKey | null> {
  const b64 = await fetchPeerPublicKey(peerUserId);
  if (!b64) return null;
  return importPublicKeySpkiB64(b64);
}
