server/sessionStore.js

import crypto from 'crypto';
import fs from 'fs/promises';
import path from 'path';
import { loadEncryptedJson, saveEncryptedJson } from './cryptoUtils.js';

// Local encrypted chat storage.
// Chats are stored per user in /data/chat/users/<userId>/sessions/<sessionId>.json (encrypted),
// with a lightweight encrypted index for fast listing.

let _SUPABASE_URL;
let _SUPABASE_ANON_KEY;
export function initStoreConfig(url, key) {
  _SUPABASE_URL = url;
  _SUPABASE_ANON_KEY = key;
}

const TEMP_TTL_MS = 24 * 60 * 60 * 1000;
const TEMP_INACTIVITY = 12 * 60 * 60 * 1000;
const TEMP_MSG_LIMIT = 10;

const DATA_ROOT = '/data/chat';
const USERS_ROOT = path.join(DATA_ROOT, 'users');
const SHARES_FILE = path.join(DATA_ROOT, 'shares', 'index.json');
const TEMP_STORE_FILE = path.join(DATA_ROOT, 'temp_sessions.json');

const userCache = new Map(); // userId -> UserState
const tempStore = new Map(); // tempId -> TempData
const devSessions = new Map(); // token -> DeviceSession
const userWriteLocks = new Map(); // userId -> Promise

const shareState = {
  loaded: false,
  index: {
    shares: {},
  },
};

function nowIso() {
  return new Date().toISOString();
}

function safeFileId(value, fallback = 'unknown') {
  const normalized = String(value || '').trim();
  if (!normalized) return fallback;
  return normalized.replace(/[^a-zA-Z0-9_.-]+/g, '_').slice(0, 160) || fallback;
}

function cloneJson(value) {
  return JSON.parse(JSON.stringify(value));
}

function ensureMessageVersioningShape(message) {
  if (!message || typeof message !== 'object') return message;
  if (!Array.isArray(message.versions) || message.versions.length === 0) {
    message.versions = [{
      content: message.content ?? '',
      tail: [],
      timestamp: Number.isFinite(message.timestamp) ? message.timestamp : Date.now(),
    }];
    message.currentVersionIdx = 0;
  }

  const idx = Number.isInteger(message.currentVersionIdx)
    ? Math.max(0, Math.min(message.currentVersionIdx, message.versions.length - 1))
    : 0;
  message.currentVersionIdx = idx;

  const active = message.versions[idx] || {};
  if (!Array.isArray(active.tail)) active.tail = [];
  if (active.content === undefined || active.content === null) {
    active.content = message.content ?? '';
  }
  if (!Number.isFinite(active.timestamp)) {
    active.timestamp = Number.isFinite(message.timestamp) ? message.timestamp : Date.now();
  }
  message.versions[idx] = active;
  message.content = active.content;
  return message;
}

function getActiveVersionNode(message) {
  if (!message || typeof message !== 'object') return null;
  ensureMessageVersioningShape(message);
  return message.versions[message.currentVersionIdx];
}

function cloneAndRepairTree(rootMessage) {
  if (!rootMessage || typeof rootMessage !== 'object') return null;
  const cloned = cloneJson(rootMessage);
  const walk = (node) => {
    if (!node || typeof node !== 'object') return;
    ensureMessageVersioningShape(node);
    const active = getActiveVersionNode(node);
    for (const child of active?.tail || []) walk(child);
  };
  walk(cloned);
  return cloned;
}

function getActiveLeafMessage(rootMessage) {
  let current = rootMessage;
  while (current) {
    const active = getActiveVersionNode(current);
    const tail = Array.isArray(active?.tail) ? active.tail : [];
    if (!tail.length) return current;
    current = tail[tail.length - 1];
  }
  return rootMessage;
}

function appendEntriesToActiveLeaf(rootMessage, entries = []) {
  if (!rootMessage || !entries.length) return rootMessage;
  const leaf = getActiveLeafMessage(rootMessage);
  const active = getActiveVersionNode(leaf);
  active.tail = [...(active.tail || []), ...entries];
  return rootMessage;
}

function normalizeHistoryToTree(history) {
  const rawHistory = Array.isArray(history) ? history : [];
  if (!rawHistory.length) return [];

  // Already tree-shaped ([rootMessage]).
  if (
    rawHistory.length === 1 &&
    rawHistory[0] &&
    typeof rawHistory[0] === 'object' &&
    Array.isArray(rawHistory[0].versions)
  ) {
    const root = cloneAndRepairTree(rawHistory[0]);
    return root ? [root] : [];
  }

  // Legacy flat history; rebuild a linear active-branch tree.
  const nodes = rawHistory
    .filter((entry) => entry && typeof entry === 'object')
    .map((entry) => ensureMessageVersioningShape(cloneJson(entry)));
  if (!nodes.length) return [];
  const root = nodes[0];
  for (let i = 1; i < nodes.length; i++) {
    appendEntriesToActiveLeaf(root, [nodes[i]]);
  }
  return [root];
}

function ensureSessionShape(raw, fallbackId = null) {
  const created = Number.isFinite(raw?.created) ? raw.created : Date.now();
  let history = [];

  // Prefer desktop's legacy serialized tree when present.
  const legacyRootRaw = typeof raw?.__historyRootJson === 'string'
    ? raw.__historyRootJson.trim()
    : '';
  if (legacyRootRaw) {
    try {
      const parsedRoot = cloneAndRepairTree(JSON.parse(legacyRootRaw));
      if (parsedRoot) history = [parsedRoot];
    } catch {
      history = [];
    }
  }
  if (!history.length) {
    history = normalizeHistoryToTree(raw?.history);
  }

  return {
    id: raw?.id || fallbackId || crypto.randomUUID(),
    name: String(raw?.name || 'New Chat'),
    created,
    history,
    model: raw?.model || null,
  };
}

function buildSessionMeta(session, existingMeta = null) {
  const created = Number.isFinite(session?.created)
    ? session.created
    : (Number.isFinite(existingMeta?.created) ? existingMeta.created : Date.now());
  return {
    id: session.id,
    name: String(session.name || existingMeta?.name || 'New Chat'),
    created,
    model: session.model || existingMeta?.model || null,
    updatedAt: nowIso(),
  };
}

function ensureUserState(userId) {
  if (!userCache.has(userId)) {
    userCache.set(userId, {
      indexLoaded: false,
      sessionsMeta: new Map(),
      loadedSessions: new Map(),
      online: new Set(),
    });
  }
  return userCache.get(userId);
}

function userDir(userId) {
  return path.join(USERS_ROOT, safeFileId(userId));
}

function userIndexFile(userId) {
  return path.join(userDir(userId), 'index.json');
}

function userSessionFile(userId, sessionId) {
  return path.join(userDir(userId), 'sessions', `${safeFileId(sessionId)}.json`);
}

function userIndexAad(userId) {
  return `chat:user:${userId}:index`;
}

function userSessionAad(userId, sessionId) {
  return `chat:user:${userId}:session:${sessionId}`;
}

function toSerializableSessionMetaMap(map) {
  const sessions = {};
  for (const [id, meta] of map.entries()) {
    sessions[id] = {
      id,
      name: String(meta?.name || 'New Chat'),
      created: Number.isFinite(meta?.created) ? meta.created : Date.now(),
      model: meta?.model || null,
      updatedAt: meta?.updatedAt || nowIso(),
    };
  }
  return sessions;
}

async function ensureUserIndexLoaded(userId) {
  const state = ensureUserState(userId);
  if (state.indexLoaded) return state;

  const stored = await loadEncryptedJson(userIndexFile(userId), userIndexAad(userId));
  state.sessionsMeta.clear();

  const sessions = stored?.sessions || {};
  for (const [id, meta] of Object.entries(sessions)) {
    state.sessionsMeta.set(id, {
      id,
      name: String(meta?.name || 'New Chat'),
      created: Number.isFinite(meta?.created) ? meta.created : Date.now(),
      model: meta?.model || null,
      updatedAt: meta?.updatedAt || nowIso(),
    });
  }

  state.indexLoaded = true;
  return state;
}

async function saveUserIndex(userId) {
  const state = ensureUserState(userId);
  const payload = {
    sessions: toSerializableSessionMetaMap(state.sessionsMeta),
  };
  await saveEncryptedJson(userIndexFile(userId), payload, userIndexAad(userId));
}

async function loadUserSessionFromDisk(userId, sessionId) {
  const raw = await loadEncryptedJson(userSessionFile(userId, sessionId), userSessionAad(userId, sessionId));
  if (!raw) return null;
  return ensureSessionShape(raw, sessionId);
}

async function saveUserSessionToDisk(userId, session) {
  const shaped = ensureSessionShape(session);
  await saveEncryptedJson(userSessionFile(userId, shaped.id), shaped, userSessionAad(userId, shaped.id));
}

async function deleteUserSessionFromDisk(userId, sessionId) {
  await fs.rm(userSessionFile(userId, sessionId), { force: true }).catch(() => {});
}

async function withUserWriteLock(userId, fn) {
  const prior = userWriteLocks.get(userId) || Promise.resolve();
  const next = prior.catch(() => {}).then(fn);
  userWriteLocks.set(userId, next.finally(() => {
    if (userWriteLocks.get(userId) === next) userWriteLocks.delete(userId);
  }));
  return next;
}

function sessionForList(meta, loaded) {
  const source = loaded || meta;
  return {
    id: source.id,
    name: source.name || 'New Chat',
    created: Number.isFinite(source.created) ? source.created : Date.now(),
    history: loaded?.history || [],
    model: source.model || null,
  };
}

async function ensureShareIndexLoaded() {
  if (shareState.loaded) return;
  const stored = await loadEncryptedJson(SHARES_FILE, 'chat:shares:index');
  shareState.index = {
    shares: stored?.shares || {},
  };
  shareState.loaded = true;
}

async function saveShareIndex() {
  await saveEncryptedJson(SHARES_FILE, shareState.index, 'chat:shares:index');
}

async function loadTempStore() {
  const data = await loadEncryptedJson(TEMP_STORE_FILE, 'chat:temp:index');
  if (!data) return;
  for (const [id, d] of Object.entries(data)) {
    const restoredSessions = {};
    for (const [sessionId, sessionValue] of Object.entries(d.sessions || {})) {
      restoredSessions[sessionId] = ensureSessionShape(sessionValue, sessionId);
    }
    tempStore.set(id, {
      sessions: new Map(Object.entries(restoredSessions)),
      msgCount: d.msgCount || 0,
      created: d.created || Date.now(),
      lastActive: d.lastActive || Date.now(),
    });
  }
}

async function saveTempStore() {
  const data = {};
  for (const [id, d] of tempStore) {
    data[id] = {
      sessions: Object.fromEntries(d.sessions),
      msgCount: d.msgCount,
      created: d.created,
      lastActive: d.lastActive,
    };
  }
  await saveEncryptedJson(TEMP_STORE_FILE, data, 'chat:temp:index');
}

loadTempStore().catch((err) => console.error('Failed to load temp store:', err));

setInterval(async () => {
  const now = Date.now();
  for (const [id, d] of tempStore) {
    if (now - d.created > TEMP_TTL_MS || now - d.lastActive > TEMP_INACTIVITY) {
      tempStore.delete(id);
    }
  }
  await saveTempStore().catch((err) => console.error('Failed to save temp store:', err));
}, 30 * 60 * 1000);

export const sessionStore = {
  // TEMP
  initTemp(t) {
    if (!tempStore.has(t)) {
      tempStore.set(t, {
        sessions: new Map(),
        msgCount: 0,
        created: Date.now(),
        lastActive: Date.now(),
      });
    }
    return tempStore.get(t);
  },
  tempCanSend(t) {
    const d = tempStore.get(t);
    return d ? d.msgCount < TEMP_MSG_LIMIT : false;
  },
  tempBump(t) {
    const d = tempStore.get(t);
    if (d) {
      d.msgCount++;
      d.lastActive = Date.now();
    }
  },
  getTempSessions(t) {
    return [...(tempStore.get(t)?.sessions.values() || [])];
  },
  getTempSession(t, id) {
    return tempStore.get(t)?.sessions.get(id) || null;
  },
  createTempSession(t) {
    const d = this.initTemp(t);
    const s = ensureSessionShape({
      id: crypto.randomUUID(),
      name: 'New Chat',
      created: Date.now(),
      history: [],
      model: null,
    });
    d.sessions.set(s.id, s);
    d.lastActive = Date.now();
    saveTempStore().catch((err) => console.error('Failed to save temp store:', err));
    return s;
  },
  updateTempSession(t, id, patch) {
    const d = tempStore.get(t);
    if (!d) return null;
    const s = d.sessions.get(id);
    if (!s) return null;
    Object.assign(s, patch);
    const shaped = ensureSessionShape(s, id);
    d.sessions.set(id, shaped);
    d.lastActive = Date.now();
    saveTempStore().catch((err) => console.error('Failed to save temp store:', err));
    return shaped;
  },
  restoreTempSession(t, session) {
    const d = this.initTemp(t);
    const restored = ensureSessionShape(JSON.parse(JSON.stringify(session)));
    d.sessions.set(restored.id, restored);
    d.lastActive = Date.now();
    saveTempStore().catch((err) => console.error('Failed to save temp store:', err));
    return restored;
  },
  deleteTempSession(t, id) {
    tempStore.get(t)?.sessions.delete(id);
    saveTempStore().catch((err) => console.error('Failed to save temp store:', err));
  },
  deleteTempAll(t) {
    tempStore.get(t)?.sessions.clear();
    saveTempStore().catch((err) => console.error('Failed to save temp store:', err));
  },
  deleteTempSessionEverywhere(id) {
    let changed = false;
    for (const temp of tempStore.values()) {
      if (temp.sessions.delete(id)) changed = true;
    }
    if (changed) saveTempStore().catch((err) => console.error('Failed to save temp store:', err));
    return changed;
  },

  async transferTempToUser(tempId, userId, _accessToken) {
    const d = tempStore.get(tempId);
    if (!d || !d.sessions.size) return;

    await ensureUserIndexLoaded(userId);

    for (const s of d.sessions.values()) {
      if (!s.history || s.history.length === 0) continue;
      if (ensureUserState(userId).sessionsMeta.has(s.id)) continue;
      const copy = ensureSessionShape(JSON.parse(JSON.stringify(s)));
      await withUserWriteLock(userId, async () => {
        const state = ensureUserState(userId);
        state.loadedSessions.set(copy.id, copy);
        state.sessionsMeta.set(copy.id, buildSessionMeta(copy));
        await saveUserSessionToDisk(userId, copy);
        await saveUserIndex(userId);
      });
    }
  },

  // USERS
  _ensureUser(uid) {
    return ensureUserState(uid);
  },

  async loadUserSessions(userId, _accessToken) {
    const state = await ensureUserIndexLoaded(userId);
    return [...state.sessionsMeta.values()]
      .sort((a, b) => new Date(b.updatedAt || 0).getTime() - new Date(a.updatedAt || 0).getTime())
      .map((meta) => sessionForList(meta, state.loadedSessions.get(meta.id)));
  },

  getUserSessions(uid) {
    const state = userCache.get(uid);
    if (!state) return [];
    return [...state.sessionsMeta.values()].map((meta) => sessionForList(meta, state.loadedSessions.get(meta.id)));
  },

  getUserSession(uid, id) {
    return userCache.get(uid)?.loadedSessions.get(id) || null;
  },

  async getUserSessionResolved(uid, id) {
    const state = await ensureUserIndexLoaded(uid);
    if (state.loadedSessions.has(id)) return state.loadedSessions.get(id);
    const meta = state.sessionsMeta.get(id);
    if (!meta) return null;

    const loaded = await loadUserSessionFromDisk(uid, id);
    if (!loaded) return null;

    const merged = ensureSessionShape({
      ...loaded,
      id,
      name: loaded.name || meta.name,
      created: Number.isFinite(loaded.created) ? loaded.created : meta.created,
      model: loaded.model || meta.model,
    }, id);

    state.loadedSessions.set(id, merged);
    return merged;
  },

  async createUserSession(userId, _accessToken) {
    const s = ensureSessionShape({
      id: crypto.randomUUID(),
      name: 'New Chat',
      created: Date.now(),
      history: [],
      model: null,
    });

    await ensureUserIndexLoaded(userId);
    await withUserWriteLock(userId, async () => {
      const state = ensureUserState(userId);
      state.loadedSessions.set(s.id, s);
      state.sessionsMeta.set(s.id, buildSessionMeta(s));
      await saveUserSessionToDisk(userId, s);
      await saveUserIndex(userId);
    });
    return s;
  },

  async restoreUserSession(userId, _accessToken, session) {
    const restored = ensureSessionShape(JSON.parse(JSON.stringify(session)));
    await ensureUserIndexLoaded(userId);

    await withUserWriteLock(userId, async () => {
      const state = ensureUserState(userId);
      state.loadedSessions.set(restored.id, restored);
      state.sessionsMeta.set(restored.id, buildSessionMeta(restored, state.sessionsMeta.get(restored.id)));
      await saveUserSessionToDisk(userId, restored);
      await saveUserIndex(userId);
    });
    return restored;
  },

  async updateUserSession(userId, _accessToken, sessionId, patch) {
    await ensureUserIndexLoaded(userId);
    const current = await this.getUserSessionResolved(userId, sessionId);
    if (!current) return null;

    Object.assign(current, patch || {});
    const updated = ensureSessionShape(current, sessionId);

    await withUserWriteLock(userId, async () => {
      const state = ensureUserState(userId);
      state.loadedSessions.set(sessionId, updated);
      state.sessionsMeta.set(sessionId, buildSessionMeta(updated, state.sessionsMeta.get(sessionId)));
      await saveUserSessionToDisk(userId, updated);
      await saveUserIndex(userId);
    });

    return updated;
  },

  async deleteUserSession(userId, _accessToken, id) {
    await ensureUserIndexLoaded(userId);
    await withUserWriteLock(userId, async () => {
      const state = ensureUserState(userId);
      state.loadedSessions.delete(id);
      state.sessionsMeta.delete(id);
      await deleteUserSessionFromDisk(userId, id);
      await saveUserIndex(userId);
    });
  },

  async deleteAllUserSessions(userId, _accessToken) {
    await ensureUserIndexLoaded(userId);
    const state = ensureUserState(userId);
    const ids = [...state.sessionsMeta.keys()];

    await withUserWriteLock(userId, async () => {
      for (const id of ids) {
        await deleteUserSessionFromDisk(userId, id);
      }
      state.loadedSessions.clear();
      state.sessionsMeta.clear();
      await saveUserIndex(userId);
    });

    return true;
  },

  markOnline(uid, ws) {
    ensureUserState(uid).online.add(ws);
  },

  markOffline(uid, ws) {
    userCache.get(uid)?.online.delete(ws);
  },

  // SHARE
  async createShareToken(userId, _accessToken, sessionId) {
    const s = await this.getUserSessionResolved(userId, sessionId);
    if (!s) return null;

    const token = crypto.randomBytes(24).toString('base64url');
    await ensureShareIndexLoaded();
    shareState.index.shares[token] = {
      token,
      owner_id: userId,
      session_snapshot: JSON.parse(JSON.stringify(s)),
      created_at: nowIso(),
    };
    await saveShareIndex();
    return token;
  },

  async resolveShareToken(token) {
    await ensureShareIndexLoaded();
    return shareState.index.shares[String(token || '')] || null;
  },

  async importSharedSession(userId, accessToken, token) {
    const shared = await this.resolveShareToken(token);
    if (!shared) return null;
    const snap = ensureSessionShape(shared.session_snapshot);
    const newSession = {
      ...snap,
      id: crypto.randomUUID(),
      name: `${snap.name} (shared)`,
      created: Date.now(),
    };
    await this.restoreUserSession(userId, accessToken, newSession);
    return newSession;
  },
};

export const deviceSessionStore = {
  create(userId, ip, userAgent) {
    const token = crypto.randomBytes(32).toString('hex');
    devSessions.set(token, {
      token,
      userId,
      ip,
      userAgent,
      createdAt: nowIso(),
      lastSeen: nowIso(),
      active: true,
    });
    return token;
  },
  getForUser(uid) {
    return [...devSessions.values()].filter((s) => s.userId === uid && s.active);
  },
  revoke(token) {
    const s = devSessions.get(token);
    if (s) {
      s.active = false;
      return s;
    }
    return null;
  },
  revokeAllExcept(uid, except) {
    for (const [t, s] of devSessions) {
      if (s.userId === uid && t !== except) s.active = false;
    }
  },
  validate(token) {
    const s = devSessions.get(token);
    if (!s || !s.active) return null;
    s.lastSeen = nowIso();
    return s;
  },
};