import { safeSend, broadcastToUser } from './helpers.js';
import { LIGHTNING_BASE, PUBLIC_URL } from './config.js';
import { sessionStore, deviceSessionStore } from './sessionStore.js';
import { rateLimiter } from './rateLimiter.js';
import {
  verifySupabaseToken, getUserSettings, saveUserSettings,
  getUserProfile, setUsername, getSubscriptionInfo,
  getTierConfig, getUsageInfo,
} from './auth.js';
import { streamChat } from './chatStream.js';
import crypto from 'crypto';

const activeStreams = new Map();

export async function handleWsMessage(ws, msg, wsClients) {
  const client = wsClients.get(ws); if (!client) return;
  const h = handlers[msg.type];
  if (h) return h(ws, msg, client, wsClients);
  safeSend(ws, { type: 'error', message: `Unknown: ${msg.type}` });
}

// wrap broadcastToUser to pass wsClients automatically
function bcast(wsClients, userId, data, excludeWs) {
  broadcastToUser(wsClients, userId, data, excludeWs);
}

const handlers = {
  'auth:login': async (ws, msg, client, wsClients) => {
    const { accessToken, tempId: clientTempId } = msg;
    if (!accessToken) return safeSend(ws, { type: 'auth:error', message: 'Missing token' });
    const user = await verifySupabaseToken(accessToken);
    if (!user) return safeSend(ws, { type: 'auth:error', message: 'Invalid token' });

    client.userId = user.id; client.accessToken = accessToken; client.authenticated = true;
    client.deviceToken = deviceSessionStore.create(user.id, client.ip, client.userAgent);
    sessionStore.markOnline(user.id, ws);

    const tId = clientTempId || client.tempId;
    await sessionStore.transferTempToUser(tId, user.id, accessToken);

    const [sessions, settings, profile] = await Promise.all([
      sessionStore.loadUserSessions(user.id, accessToken),
      getUserSettings(user.id, accessToken),
      getUserProfile(user.id, accessToken),
    ]);

    safeSend(ws, { type: 'auth:ok', userId: user.id, email: user.email,
      deviceToken: client.deviceToken, sessions: sessions.map(ser), settings, profile });

    bcast(wsClients, user.id, { type: 'auth:newLogin', message: 'New login on your account.',
      ip: client.ip, userAgent: client.userAgent, timestamp: new Date().toISOString() }, ws);
  },

  'auth:logout': (ws, msg, client) => {
    if (client.deviceToken) deviceSessionStore.revoke(client.deviceToken);
    Object.assign(client, { userId: null, authenticated: false, accessToken: null, deviceToken: null });
    safeSend(ws, { type: 'auth:loggedOut' });
  },

  'auth:guest': (ws, msg, client) => {
    const t = msg.tempId || client.tempId;
    sessionStore.initTemp(t);
    safeSend(ws, { type: 'auth:guestOk', tempId: t, sessions: sessionStore.getTempSessions(t).map(ser) });
  },

  'sessions:list': (ws, msg, client) => {
    const list = client.userId
      ? sessionStore.getUserSessions(client.userId)
      : sessionStore.getTempSessions(client.tempId);
    list.sort((a, b) => b.created - a.created);
    safeSend(ws, { type: 'sessions:list', sessions: list.map(ser) });
  },

  'sessions:create': async (ws, msg, client) => {
    const s = client.userId
      ? await sessionStore.createUserSession(client.userId, client.accessToken)
      : sessionStore.createTempSession(client.tempId);
    safeSend(ws, { type: 'sessions:created', session: ser(s) });
  },

  'sessions:delete': async (ws, msg, client) => {
    if (client.userId) await sessionStore.deleteUserSession(client.userId, client.accessToken, msg.sessionId);
    else sessionStore.deleteTempSession(client.tempId, msg.sessionId);
    safeSend(ws, { type: 'sessions:deleted', sessionId: msg.sessionId });
  },

  'sessions:deleteAll': async (ws, msg, client) => {
    if (client.userId) await sessionStore.deleteAllUserSessions(client.userId, client.accessToken);
    else sessionStore.deleteTempAll(client.tempId);
    safeSend(ws, { type: 'sessions:deletedAll' });
  },

  'sessions:rename': async (ws, msg, client) => {
    const name = (msg.name || '').trim(); if (!name) return;
    if (client.userId)
      await sessionStore.updateUserSession(client.userId, client.accessToken, msg.sessionId, { name });
    else sessionStore.updateTempSession(client.tempId, msg.sessionId, { name });
    safeSend(ws, { type: 'sessions:renamed', sessionId: msg.sessionId, name });
  },

  'sessions:get': (ws, msg, client) => {
    const s = client.userId
      ? sessionStore.getUserSession(client.userId, msg.sessionId)
      : sessionStore.getTempSession(client.tempId, msg.sessionId);
    if (!s) return safeSend(ws, { type: 'error', message: 'Session not found' });
    safeSend(ws, { type: 'sessions:data', session: ser(s) });
  },

  'sessions:share': async (ws, msg, client) => {
    if (!client.userId) return safeSend(ws, { type: 'error', message: 'Sign in to share' });
    const token = await sessionStore.createShareToken(client.userId, client.accessToken, msg.sessionId);
    if (!token) return safeSend(ws, { type: 'error', message: 'Share failed' });
    safeSend(ws, { type: 'sessions:shareUrl', url: `${PUBLIC_URL}/?share=${token}`, sessionId: msg.sessionId });
  },

  'sessions:import': async (ws, msg, client) => {
    if (!client.userId) return safeSend(ws, { type: 'error', message: 'Sign in to import' });
    const s = await sessionStore.importSharedSession(client.userId, client.accessToken, msg.token);
    if (!s) return safeSend(ws, { type: 'error', message: 'Invalid share link' });
    safeSend(ws, { type: 'sessions:imported', session: ser(s) });
  },

  'chat:send': async (ws, msg, client) => {
    const { sessionId, content, tools } = msg;
    if (!client.userId) {
      if (!sessionStore.tempCanSend(client.tempId)) return safeSend(ws, { type: 'chat:limitReached' });
      sessionStore.tempBump(client.tempId);
    }
    const session = client.userId
      ? sessionStore.getUserSession(client.userId, sessionId)
      : sessionStore.getTempSession(client.tempId, sessionId);
    if (!session) return safeSend(ws, { type: 'error', message: 'Session not found' });

    if (activeStreams.has(ws)) activeStreams.get(ws).abort();
    const abort = new AbortController();
    activeStreams.set(ws, abort);
    safeSend(ws, { type: 'chat:start', sessionId });

    let fullText = '';
    const assetsCollected = [], toolCallsCollected = [];

    await streamChat(ws, {
      history: session.history || [], userMessage: content, tools: tools || {},
      accessToken: client.accessToken, clientId: msg.clientId, abortSignal: abort.signal,
      onToken(t) { fullText += t; safeSend(ws, { type: 'chat:token', token: t, sessionId }); },
      onToolCall(call) {
        safeSend(ws, { type: 'chat:toolCall', call, sessionId });
        if (call.state === 'resolved' || call.state === 'canceled') toolCallsCollected.push(call);
      },
      onNewAsset(asset) { safeSend(ws, { type: 'chat:asset', asset, sessionId }); assetsCollected.push(asset); },
      async onDone(text, toolCalls, aborted) {
        activeStreams.delete(ws);
        const finalText = text || fullText;
        const userEntry = buildEntry('user', content);
        const asstEntry = buildEntry('assistant', finalText, toolCalls || []);
        const newHistory = [...(session.history || []), userEntry, asstEntry];
        for (const a of assetsCollected) newHistory.push({ role: a.role, content: a.content, timestamp: Date.now() });
        let newName = session.name;
        if (!(session.history || []).length || session.name === 'New Chat')
          newName = await autoName(content, client.accessToken);
        if (client.userId)
          await sessionStore.updateUserSession(client.userId, client.accessToken, sessionId, { history: newHistory, name: newName });
        else sessionStore.updateTempSession(client.tempId, sessionId, { history: newHistory, name: newName });
        safeSend(ws, { type: aborted ? 'chat:aborted' : 'chat:done', sessionId, name: newName, history: newHistory });
      },
      onError(err) { activeStreams.delete(ws); safeSend(ws, { type: 'chat:error', error: String(err), sessionId }); },
    });
  },

  'chat:stop': (ws) => { if (activeStreams.has(ws)) { activeStreams.get(ws).abort(); activeStreams.delete(ws); } },

  'chat:editMessage': async (ws, msg, client) => {
    const { sessionId, messageIndex, newContent } = msg;
    const session = client.userId
      ? sessionStore.getUserSession(client.userId, sessionId)
      : sessionStore.getTempSession(client.tempId, sessionId);
    if (!session) return;
    const history = session.history || [];
    const m = history[messageIndex]; if (!m) return;
    if (!m.versions) m.versions = [{ content: m.content, tail: history.slice(messageIndex + 1), timestamp: m.timestamp || Date.now() }];
    m.versions.push({ content: newContent, tail: history.slice(messageIndex + 1), timestamp: Date.now() });
    m.currentVersionIdx = m.versions.length - 1;
    m.content = newContent;
    const newHistory = history.slice(0, messageIndex + 1);
    if (client.userId) await sessionStore.updateUserSession(client.userId, client.accessToken, sessionId, { history: newHistory });
    else sessionStore.updateTempSession(client.tempId, sessionId, { history: newHistory });
    safeSend(ws, { type: 'chat:messageEdited', sessionId, messageIndex, message: m, history: newHistory });
  },

  'chat:selectVersion': async (ws, msg, client) => {
    const { sessionId, messageIndex, versionIdx } = msg;
    const session = client.userId
      ? sessionStore.getUserSession(client.userId, sessionId)
      : sessionStore.getTempSession(client.tempId, sessionId);
    if (!session) return;
    const history = session.history || [];
    const m = history[messageIndex];
    if (!m?.versions || versionIdx < 0 || versionIdx >= m.versions.length) return;
    const v = m.versions[versionIdx];
    m.currentVersionIdx = versionIdx; m.content = v.content;
    const newHistory = [...history.slice(0, messageIndex + 1), ...(v.tail || [])];
    if (client.userId) await sessionStore.updateUserSession(client.userId, client.accessToken, sessionId, { history: newHistory });
    else sessionStore.updateTempSession(client.tempId, sessionId, { history: newHistory });
    safeSend(ws, { type: 'chat:versionSelected', sessionId, history: newHistory });
  },

  'settings:get': async (ws, msg, client) => {
    const s = client.userId
      ? await getUserSettings(client.userId, client.accessToken)
      : { theme: 'dark', webSearch: true, imageGen: true, videoGen: true, audioGen: true };
    safeSend(ws, { type: 'settings:data', settings: s });
  },
  'settings:save': async (ws, msg, client, wsClients) => {
    if (!client.userId) return;
    await saveUserSettings(client.userId, client.accessToken, msg.settings);
    safeSend(ws, { type: 'settings:saved' });
    bcast(wsClients, client.userId, { type: 'settings:updated', settings: msg.settings }, ws);
  },

  'account:getProfile':      async (ws, msg, c) => { if (!c.userId) return; safeSend(ws, { type: 'account:profile', profile: await getUserProfile(c.userId, c.accessToken) }); },
  'account:setUsername':     async (ws, msg, c) => { if (!c.userId) return; safeSend(ws, { type: 'account:usernameResult', ...await setUsername(c.userId, c.accessToken, msg.username) }); },
  'account:getSubscription': async (ws, msg, c) => { if (!c.userId) return; safeSend(ws, { type: 'account:subscription', info: await getSubscriptionInfo(c.accessToken) }); },
  'account:getUsage':        async (ws, msg, c) => { safeSend(ws, { type: 'account:usage', usage: await getUsageInfo(c.accessToken) }); },
  'account:getTierConfig':   async (ws)          => { safeSend(ws, { type: 'account:tierConfig', config: await getTierConfig() }); },
  'account:getSessions':     (ws, msg, c)         => { if (!c.userId) return; safeSend(ws, { type: 'account:deviceSessions', sessions: deviceSessionStore.getForUser(c.userId), currentToken: c.deviceToken }); },
  'account:revokeSession':   (ws, msg, c)         => { if (!c.userId) return; deviceSessionStore.revoke(msg.token); safeSend(ws, { type: 'account:sessionRevoked', token: msg.token }); },
  'account:revokeAllOthers': (ws, msg, c, wsClients) => {
    if (!c.userId) return;
    deviceSessionStore.revokeAllExcept(c.userId, c.deviceToken);
    for (const [ows, oc] of wsClients)
      if (oc.userId === c.userId && ows !== ws) safeSend(ows, { type: 'auth:forcedLogout', reason: 'Session revoked by another device' });
    safeSend(ws, { type: 'account:allOthersRevoked' });
  },
};

function ser(s) { return { id: s.id, name: s.name, created: s.created, history: s.history || [], model: s.model }; }

function buildEntry(role, content, toolCalls = []) {
  return { role, content, timestamp: Date.now(),
    versions: [{ content, tail: [], timestamp: Date.now() }], currentVersionIdx: 0,
    ...(toolCalls.length ? { toolCalls } : {}) };
}

async function autoName(prompt, accessToken) {
  try {
    const text = typeof prompt === 'string' ? prompt
      : (Array.isArray(prompt) ? (prompt.find(p => p.type === 'text')?.text || '') : String(prompt));
    const h = { 'Content-Type': 'application/json' };
    if (accessToken) h.Authorization = `Bearer ${accessToken}`;
    const r = await fetch(`${LIGHTNING_BASE}/gen/v1/chat/completions`, {
      method: 'POST', headers: h,
      body: JSON.stringify({ model: 'lightning',
        messages: [
          { role: 'system', content: 'Generate a 2-4 word title. No quotes. Just the title.' },
          { role: 'user',   content: text.slice(0, 200) }],
        max_tokens: 20, stream: false }),
    });
    if (!r.ok) return text.slice(0, 40) || 'New Chat';
    const d = await r.json();
    return d.choices?.[0]?.message?.content?.trim() || 'New Chat';
  } catch { return 'New Chat'; }
}