// Browser state: identity, pins, contacts and read markers in localStorage;
// sealed envelopes in IndexedDB, opened only on demand, so nothing at rest
// is plaintext (the master secret excepted — it is the user's browser profile).

import { aeadDecrypt, aeadEncrypt, hex, hkdfSha256, randomBytes, unhex, utf8Bytes } from "./crypto.js";
import { ID_LEN, MAX_CHAIN, b32decode, b32encode, identityFromSeed, identitySeed, tokenFor } from "./proto.js";

const KEY = "gsmol";

export const TIER_MAIN = 0, TIER_REQUESTS = 1;

// Every getter reads the whole blob and some are called per message row, so the
// parse is cached. `derived` holds the identities that follow from it, which
// cost an Ed25519 scalar multiplication each.
let cached = null, derived = null;

function load() {
  if (cached) return cached;
  try {
    cached = JSON.parse(localStorage.getItem(KEY) || "{}");
  } catch {
    cached = {};
  }
  return cached;
}

function save(state) {
  localStorage.setItem(KEY, JSON.stringify(state));
  cached = state;
  derived = null;
}

// Another tab writing this key leaves our copy stale.
window.addEventListener("storage", (event) => {
  if (event.key === KEY || event.key === null) cached = derived = null;
});

function update(fn) {
  const state = load();
  const next = fn(state);
  save(next ?? state);
}

// --- identity (§2) -------------------------------------------------------------

export function master() {
  const raw = load().master;
  return raw ? unhex(raw) : null;
}

// The rotation index (§7) of the identity currently in use.
export function rotations() {
  return load().rotations ?? 0;
}

export function identity() {
  return identities()[0] ?? null;
}

// §7: every key rotated away from is re-derivable from the master, since mail
// sealed to a superseded key is readable with nothing else.
export function identities() {
  if (derived) return derived;
  const m = master();
  if (!m) return [];
  derived = [];
  for (let n = rotations(); n >= 0; n--) derived.push(identityFromSeed(identitySeed(m, n)));
  return derived;
}

function bindMaster(newMaster, rotationIndex, syncOk) {
  if (master()) throw new Error("an identity already exists; rotate it instead");
  update(state => {
    state.master = hex(newMaster);
    state.rotations = rotationIndex;
    state.syncOk = syncOk;
  });
  setCursor(0, new Uint8Array(ID_LEN));
}

// A fresh identity: rotation index 0, and an empty accepted set is already
// complete, so it may sync.
export function setIdentity(newMaster) {
  bindMaster(newMaster, 0, true);
}

// §2: recovering a master alone does not recover which correspondents were
// accepted, so that set must not overwrite the server's until rebuilt.
export function restoreMaster(newMaster, rotationIndex = 0) {
  bindMaster(newMaster, rotationIndex, false);
}

// Corrects the rotation index once a RESOLVE against the account's address
// reveals which one the server actually has bound (see restoreMaster()'s
// default of 0, the common case of a never-rotated identity).
export function setRotationIndex(n) {
  update(state => { state.rotations = n; });
}

// Whether the accepted-correspondent set held here may replace the server's
// on the next AUTH — false right after a restore from the master alone, whose
// empty set must not erase the server's (§4).
export function syncOk() {
  return load().syncOk ?? true;
}

export function setSyncOk(ok) {
  update(state => { state.syncOk = ok; });
}

// Rotation (§7): only the index advances; the superseded key stays derivable
// from the master, so nothing has to be archived.
export function advanceRotation() {
  const current = rotations();
  if (!master()) throw new Error("no identity to rotate");
  if (current >= MAX_CHAIN) throw new Error(`the rotation chain is full at ${MAX_CHAIN} links`);
  update(state => { state.rotations = current + 1; });
}

// --- account and server pins ---------------------------------------------------

export function account() {
  const a = load().account;
  return a ? { ...a } : null;
}

export function setAccount(addr) {
  update(state => {
    state.account = { user: addr.user, host: addr.host, port: addr.port };
  });
}

export function serverPin(host) {
  const raw = (load().servers || {})[host];
  return raw ? b32decode(raw) : null;
}

export function pinServer(host, keyBytes) {
  update(state => {
    state.servers = { ...(state.servers || {}), [host]: b32encode(keyBytes) };
  });
}

export function allPins() {
  return Object.entries(load().servers || {}).map(([host, key]) => ({ host, key }));
}

// §4/first contact: when true, resolving a never-before-seen recipient's key
// from a server that isn't pinned is refused instead of merely toasted after
// the envelope is already sealed. Persisted, not in-memory — an in-memory
// toggle fails open on every restart, worth little as a security setting.
export function requirePinnedServer() {
  return load().requirePinnedServer ?? false;
}

export function setRequirePinnedServer(value) {
  update(state => { state.requirePinnedServer = value; });
}

// --- FETCH behavior and cursor (§6.1) -------------------------------------------

// When true, fetch does not acknowledge (delete) what it retrieves — mail
// stays on the server until explicitly deleted. Defaults to the original
// behavior: fetched mail is acknowledged immediately.
export function leaveOnServer() {
  return load().leaveOnServer ?? false;
}

export function setLeaveOnServer(value) {
  update(state => { state.leaveOnServer = value; });
}

// Seconds a network call waits for a response before giving up. Applied per
// read, not per operation, so a slow-but-trickling exchange is not cut off
// just because it spans several reads.
export function timeoutSeconds() {
  return load().timeoutSeconds ?? 30;
}

export function setTimeoutSeconds(value) {
  const n = Math.trunc(Number(value));
  if (!Number.isFinite(n)) return;  // non-numeric input is ignored, not an error
  update(state => { state.timeoutSeconds = Math.min(600, Math.max(1, n)); });
}

// Undoes setIdentity()/restoreMaster(): onboarding's "back"/"start over",
// before anything is bound to an account. Pins and contacts are left alone —
// server trust is not identity-scoped.
export function discardIdentity() {
  if (account()) throw new Error("already bound to an account; log out instead");
  update(state => {
    delete state.master;
    delete state.rotations;
    delete state.syncOk;
    delete state.afterTime;
    delete state.afterId;
  });
}

export function cursor() {
  const state = load();
  const afterId = state.afterId;
  return [state.afterTime ?? 0, afterId ? unhex(afterId) : new Uint8Array(ID_LEN)];
}

export function setCursor(afterTime, afterId) {
  update(state => {
    state.afterTime = afterTime;
    state.afterId = hex(afterId);
  });
}

// --- contacts ------------------------------------------------------------------

export function contact(address) {
  const c = (load().contacts || {})[address];
  return c ? { key: b32decode(c.key), verified: c.verified, history: c.history ?? [] } : null;
}

// A key that displaces another is kept: it is the only local record that the
// contact rotated, and §8 turns on the user being able to notice such changes.
// Re-saving the same key is not a rotation and must not add an entry.
export function saveContact(address, keyBytes, verified) {
  const key = b32encode(keyBytes);
  update(state => {
    const contacts = { ...(state.contacts || {}) };
    const previous = contacts[address];
    const history = previous && previous.key !== key
      ? [...(previous.history ?? []), { key: previous.key, until: Date.now() }]
      : previous?.history ?? [];
    contacts[address] = { key, verified, seenAt: Date.now(), ...(history.length && { history }) };
    state.contacts = contacts;
  });
}

export function addressForKey(keyBytes) {
  return Object.entries(load().contacts || {})
    .find(([, c]) => c.key === b32encode(keyBytes))?.[0] ?? null;
}

export function allContacts() {
  return Object.entries(load().contacts || {})
    .map(([address, c]) => ({ address, key: c.key, verified: c.verified, history: c.history ?? [] }));
}

// --- accept tokens (§5.8) --------------------------------------------------------

export function accepted(address) {
  const a = (load().accepted || {})[address];
  return a ? { identity: b32decode(a.identity), active: a.active } : null;
}

// Admit a contact to the main tier. The identity is frozen at acceptance — a
// re-accept after a block must not change which key the token is derived
// from (§5.8).
export function accept(address, identityBytes) {
  update(state => {
    const table = { ...(state.accepted || {}) };
    const previous = table[address];
    table[address] = {
      identity: previous?.identity ?? b32encode(identityBytes),
      active: true,
      addedAt: previous?.addedAt ?? Date.now(),
    };
    state.accepted = table;
  });
}

// Withdraw a contact's accept token; their mail lands in the requests tier
// from their next message on. Throws if the contact was never accepted.
export function block(address) {
  if (!(load().accepted || {})[address]) throw new Error(`${address} was never accepted`);
  update(state => {
    const table = { ...(state.accepted || {}) };
    table[address] = { ...table[address], active: false };
    state.accepted = table;
  });
}

export function allAccepted() {
  return Object.entries(load().accepted || {})
    .map(([address, a]) => ({ address, identity: b32decode(a.identity), active: a.active }));
}

// §4: the tokens to push with AUTH, and whether to push at all. A client that
// cannot vouch for its own set — one restored from the master alone — must
// not replace the server's with an incomplete one.
export function tokenSet(masterBytes) {
  if (!syncOk()) return { sync: 0, tokens: [] };
  const active = allAccepted().filter(a => a.active).sort((a, b) => a.address.localeCompare(b.address));
  return { sync: 1, tokens: active.map(a => tokenFor(masterBytes, a.identity)) };
}

// A token received from a correspondent, filed under the address that issued
// it: an address outlives the keys behind it, so the token keeps working
// across the issuer's rotations (§5.8).
export function tokenFrom(address) {
  const raw = (load().tokens || {})[address];
  return raw ? b32decode(raw.token) : null;
}

export function learnToken(address, tokenBytes) {
  update(state => {
    const tokens = { ...(state.tokens || {}) };
    tokens[address] = { token: b32encode(tokenBytes), seenAt: Date.now() };
    state.tokens = tokens;
  });
}

// --- read markers ---------------------------------------------------------------

export function markRead(idHex) {
  update(state => {
    state.read = { ...(state.read || {}), [idHex]: true };
  });
}

export function isRead(idHex) {
  return Boolean((load().read || {})[idHex]);
}

// --- appearance ------------------------------------------------------------

// "light" or "dark" to pin a side; null to follow the system.
export function theme() {
  return load().theme ?? null;
}

export function setTheme(value) {
  update(state => {
    if (value) state.theme = value;
    else delete state.theme;
  });
}

// --- sealed mail (IndexedDB) ------------------------------------------------------

const DB_NAME = "gsmol", DB_VERSION = 1;

// Held open: reopening per read cost more than the reads themselves.
let dbPromise = null;

function openDb() {
  if (dbPromise) return dbPromise;
  dbPromise = new Promise((resolve, reject) => {
    const request = indexedDB.open(DB_NAME, DB_VERSION);
    request.onupgradeneeded = () => {
      const db = request.result;
      for (const box of ["inbox", "sent"]) {
        if (!db.objectStoreNames.contains(box)) db.createObjectStore(box, { keyPath: "id" });
      }
    };
    request.onsuccess = () => resolve(request.result);
    request.onerror = () => reject(request.error);
  }).catch(error => { dbPromise = null; throw error; });
  return dbPromise;
}

function tx(db, box, mode, fn) {
  return new Promise((resolve, reject) => {
    const request = fn(db.transaction(box, mode).objectStore(box));
    request.onsuccess = () => resolve(request.result);
    request.onerror = () => reject(request.error);
  });
}

// "requests" is a view over the same physical "inbox" records, filtered by
// tier (§5.8) — not a separate folder, so a message keeps one identity
// regardless of which tier it arrived in.
const physicalFolder = (folder) => folder === "requests" ? "inbox" : folder;

export async function storeMessage(box, record) {
  const db = await openDb();
  await tx(db, box, "readwrite", store => store.put(record));
}

// Returns null when the id already exists, so fetch can leave server state alone.
export async function storeIfNew(box, record) {
  const db = await openDb();
  const existing = await tx(db, box, "readonly", store => store.get(record.id));
  if (existing) return null;
  await tx(db, box, "readwrite", store => store.put(record));
  return record;
}

export async function listMessages(folder) {
  const physical = physicalFolder(folder);
  const db = await openDb();
  const rows = await tx(db, physical, "readonly", store => store.getAll());
  const wantTier = folder === "requests" ? TIER_REQUESTS : TIER_MAIN;
  const filtered = physical === "inbox" ? rows.filter(r => (r.tier ?? TIER_MAIN) === wantTier) : rows;
  return filtered.sort((a, b) => (b.receivedAt ?? b.sentAt) - (a.receivedAt ?? a.sentAt));  // newest first
}

export async function getMessage(folder, id) {
  const db = await openDb();
  return await tx(db, physicalFolder(folder), "readonly", store => store.get(id)) ?? null;
}

export async function removeMessage(folder, id) {
  const db = await openDb();
  await tx(db, physicalFolder(folder), "readwrite", store => store.delete(id));
}

// The unread badge must survive leaving the inbox, so it is counted from ids
// and read markers alone — no envelope is opened.
async function unreadIn(folder) {
  const rows = await listMessages(folder);
  const read = load().read || {};
  return rows.filter(row => !read[row.id]).length;
}

export const unreadCount = () => unreadIn("inbox");
export const requestsUnreadCount = () => unreadIn("requests");

// --- export / import: mail, contacts, pins — never the master --------------------

// btoa/atob work on JS's UTF-16 "binary string" convention (one code unit per
// byte); this just bridges that to Uint8Array. Browser-only, like the rest of
// this file — no Node test exercises store.js, unlike crypto.js/proto.js.
const toBase64 = (bytes) => btoa(String.fromCharCode(...bytes));
const fromBase64 = (text) => Uint8Array.from(atob(text), c => c.charCodeAt(0));

// A gsmol-only construction, not the wire protocol's own HKDF namespace
// ("smolmail/1 ..." in proto.js) — export/import has no counterpart in
// SPEC.md, so it gets its own label rather than borrowing the protocol's.
// Sealing to a key derived from the identity's own master means the file is
// opaque without it — no new passphrase to manage, and it can only ever be
// opened where that master is already restored, exactly the situation import
// already assumes.
const EXPORT_LABEL = utf8Bytes("gsmol/1 export");
const exportKey = (masterBytes) => hkdfSha256(masterBytes, new Uint8Array(0), EXPORT_LABEL, 32);

// Deliberately excludes the master: it already has its own reveal-and-copy
// flow in settings, meant for a password manager, not a downloadable file.
// Beyond that, everything else in here was worth hiding too — contacts and
// pins are a social graph and a list of which mail servers you use, not just
// the mail SPEC.md §5 makes irreplaceable once fetched — so the whole payload
// is sealed, not just the parts that were already ciphertext at rest.
export async function exportData() {
  const m = master();
  if (!m) throw new Error("no identity yet");
  const [inbox, requests, sent] = await Promise.all(
    [listMessages("inbox"), listMessages("requests"), listMessages("sent")]);
  const payload = {
    servers: Object.fromEntries(allPins().map(({ host, key }) => [host, key])),
    contacts: Object.fromEntries(allContacts().map(({ address, key, verified, history }) =>
      [address, { key, verified, history }])),
    inbox: [...inbox, ...requests].map(row => ({
      id: row.id, receivedAt: row.receivedAt, envelope: toBase64(row.envelope),
      tier: row.tier ?? TIER_MAIN, keptOnServer: row.keptOnServer ?? false,
    })),
    sent: sent.map(row =>
      ({ id: row.id, recipient: row.recipient, sentAt: row.sentAt, envelope: toBase64(row.envelope) })),
  };
  // Random per export: the key is the same every time (deterministic from the
  // master), so nonce reuse has to be ruled out the way it always is under a
  // fixed key — a fresh 96-bit nonce per seal, not a fixed one like seal()'s
  // in proto.js gets away with (there, a fresh ephemeral key each message
  // makes the derived key itself unique, so a zero nonce is safe).
  const nonce = randomBytes(12);
  const ciphertext = aeadEncrypt(exportKey(m), nonce, utf8Bytes(JSON.stringify(payload)), new Uint8Array(0));
  return {
    gsmolExport: 2,
    exportedAt: Date.now(),
    nonce: toBase64(nonce),
    ciphertext: toBase64(ciphertext),
  };
}

// Never overwrites a trust binding that already differs locally — the same
// rule refreshContact()/saveReplyAddress() apply elsewhere: an existing pin
// or contact key changes only by explicit user action, never silently. A
// malformed entry (hand-edited file, corruption) is skipped, not fatal — one
// bad record cannot abort the rest of the import, matching describe()'s
// per-message fail-open elsewhere in the app.
// A plain object, as JSON.parse would produce for `{...}`; Object.entries()
// on a string iterates its characters rather than failing, which is exactly
// the kind of malformed input this rejects as one unit instead of one per char.
const isRecord = (value) => typeof value === "object" && value !== null && !Array.isArray(value);

export async function importData(data) {
  let payload;
  if (data?.gsmolExport === 2) {
    const m = master();
    if (!m) throw new Error("no identity yet — restore it before importing");
    try {
      const plaintext = aeadDecrypt(
        exportKey(m), fromBase64(data.nonce), fromBase64(data.ciphertext), new Uint8Array(0));
      payload = JSON.parse(new TextDecoder().decode(plaintext));
    } catch {
      throw new Error("couldn't decrypt — exported by a different identity, or the file is corrupted");
    }
  } else if (data?.gsmolExport === 1) {
    payload = data;  // pre-encryption shape: the fields already sit at the top level
  } else {
    throw new Error("not a gsmol export file");
  }

  const summary = { pinsAdded: 0, pinsConflicted: 0, contactsAdded: 0, contactsConflicted: 0,
    mailAdded: 0, malformed: 0 };

  update(state => {
    const servers = { ...(state.servers || {}) };
    if (payload.servers !== undefined && !isRecord(payload.servers)) summary.malformed++;
    for (const [host, key] of Object.entries(isRecord(payload.servers) ? payload.servers : {})) {
      try {
        if (b32decode(key).length !== 32) throw new Error("bad length");
      } catch { summary.malformed++; continue; }
      if (!(host in servers)) { servers[host] = key; summary.pinsAdded++; }
      else if (servers[host] !== key) summary.pinsConflicted++;
    }
    state.servers = servers;

    const contacts = { ...(state.contacts || {}) };
    if (payload.contacts !== undefined && !isRecord(payload.contacts)) summary.malformed++;
    for (const [address, c] of Object.entries(isRecord(payload.contacts) ? payload.contacts : {})) {
      try {
        if (typeof c.key !== "string" || b32decode(c.key).length !== 32) throw new Error("bad key");
      } catch { summary.malformed++; continue; }
      if (!(address in contacts)) {
        contacts[address] = { key: c.key, verified: Boolean(c.verified), seenAt: Date.now(),
          ...(Array.isArray(c.history) && c.history.length && { history: c.history }) };
        summary.contactsAdded++;
      } else if (contacts[address].key !== c.key) {
        summary.contactsConflicted++;
      }
    }
    state.contacts = contacts;
  });

  for (const [box, rows] of [["inbox", payload.inbox], ["sent", payload.sent]]) {
    if (rows !== undefined && !Array.isArray(rows)) summary.malformed++;
    for (const row of Array.isArray(rows) ? rows : []) {
      try {
        const record = box === "inbox"
          ? { id: row.id, receivedAt: row.receivedAt, envelope: fromBase64(row.envelope),
              tier: row.tier ?? TIER_MAIN, keptOnServer: Boolean(row.keptOnServer) }
          : { id: row.id, recipient: row.recipient, sentAt: row.sentAt, envelope: fromBase64(row.envelope) };
        if (await storeIfNew(box, record)) summary.mailAdded++;
      } catch { summary.malformed++; }
    }
  }

  return summary;
}

// Logout: erases the master, pins, contacts and every cached message from
// this browser. Closes the held connection first so the delete isn't left
// "blocked" waiting for a handle that never closes on its own.
export async function clearAll() {
  localStorage.removeItem(KEY);
  cached = derived = null;
  if (dbPromise) {
    (await dbPromise).close();
    dbPromise = null;
  }
  await new Promise((resolve, reject) => {
    const request = indexedDB.deleteDatabase(DB_NAME);
    request.onsuccess = () => resolve();
    request.onerror = () => reject(request.error);
    request.onblocked = () => resolve();  // still completes once the reload closes every handle
  });
}