/**
 * Realtime Store — PR 3 §13 follow-up
 *
 * Single shared WebSocket connection to the server's wsHub, exposed as a
 * typed event subscription API. Avoids one-socket-per-page churn: every
 * page that wants live updates just calls `subscribe('event:name', fn)`
 * and the store takes care of connecting, reconnecting, and routing.
 *
 * Auth: pulls the JWT from the same `auth-storage` localStorage key that
 * api.ts reads. Matches the pattern in pages/Messages.tsx.
 *
 * Server-side payloads (`services/wsHub.js`):
 *   { type: 'message:new',     data: {...} }
 *   { type: 'message:sent',    data: {...} }
 *   { type: 'announcement:new',data: {...} }
 *   { type: 'submission:graded', data: { kind, submission_id, ... } }
 *   { type: 'marked_script:returned', data: { kind, submission_id, marked_scripts } }
 *   { type: 'marked_script:uploaded', data: { submission_id, attachments } }
 */

import { create } from 'zustand';

export type RealtimeEvent =
  | 'message:new'
  | 'message:sent'
  | 'announcement:new'
  | 'submission:graded'
  | 'marked_script:returned'
  | 'marked_script:uploaded'
  | string; // allow other future event types without an enum churn

export interface RealtimePayload {
  type: RealtimeEvent;
  data: any;
}

type Handler = (payload: RealtimePayload) => void;

interface RealtimeState {
  connected: boolean;
  subscriberCount: number;
  // PR 3 §13 follow-up — derived counters / latest payload that pages can
  // read on every render. Pages don't have to subscribe themselves; the
  // store stays subscribed for the lifetime of the session and updates
  // these fields. Decoupling the "always on" WS from per-page mount/unmount
  // avoids React-StrictMode double-invoke losing the subscription.
  markedScriptReturnCount: number;
  latestMarkedScript: RealtimePayload | null;
  // Actions
  subscribe: (type: string, handler: Handler) => () => void;
  resetMarkedScriptCount: () => void;
  /** Internal — increments when first subscriber attaches; lets the
   *  Messages.tsx connection defer to ours so we share one socket. */
  _touch: () => void;
  _release: () => void;
}

const SUBSCRIBERS = new Map<string, Set<Handler>>();
let activeSocket: WebSocket | null = null;
let reconnectTimer: ReturnType<typeof setTimeout> | null = null;
let lastToken: string | null = null;

function readToken(): string | null {
  try {
    const raw = localStorage.getItem('auth-storage');
    if (!raw) return null;
    return JSON.parse(raw)?.state?.token ?? null;
  } catch {
    return null;
  }
}

function dispatch(payload: RealtimePayload) {
  if (!payload || typeof payload.type !== 'string') return;
  // PR 3 §13 follow-up — update derived counters in the store so pages
  // can read them via useRealtimeStore() without needing their own
  // subscription. Avoids the React-StrictMode double-mount dance that
  // makes a per-page useEffect unreliable in dev.
  if (payload.type === 'marked_script:returned') {
    try {
      const state = useRealtimeStore.getState();
      const items = (payload.data as any)?.marked_scripts;
      const next = items && Array.isArray(items) && items.length > 0
        ? payload
        : state.latestMarkedScript ?? payload; // preserve last non-empty payload
      useRealtimeStore.setState({
        markedScriptReturnCount: state.markedScriptReturnCount + 1,
        latestMarkedScript: next,
      });
      notifyMarkedScript(
        'Marked script returned',
        `A teacher has returned ${items?.length ?? 1} marked file${(items?.length ?? 1) === 1 ? '' : 's'}.`
      );
    } catch (err) {
      console.warn('[realtime] marked_script:returned store update failed:', err);
    }
  }

  const handlers = SUBSCRIBERS.get(payload.type);
  if (handlers) {
    for (const h of handlers) {
      try {
        h(payload);
      } catch (err) {
        console.warn('[realtime] handler error for', payload.type, err);
      }
    }
  }
  // Wildcard subscribers (type === '*') receive every event.
  const wildcards = SUBSCRIBERS.get('*');
  if (wildcards) {
    for (const h of wildcards) {
      try {
        h(payload);
      } catch (err) {
        console.warn('[realtime] wildcard handler error', err);
      }
    }
  }
}

function scheduleReconnect(token: string) {
  if (reconnectTimer) clearTimeout(reconnectTimer);
  reconnectTimer = setTimeout(() => connect(token), 5000);
}

function connect(token: string) {
  if (activeSocket && activeSocket.readyState === WebSocket.OPEN) return;
  const protocol = window.location.protocol === 'https:' ? 'wss:' : 'ws:';
  const wsUrl = `${protocol}//${window.location.host}/ws?token=${encodeURIComponent(token)}`;

  let socket: WebSocket;
  try {
    socket = new WebSocket(wsUrl);
  } catch {
    scheduleReconnect(token);
    return;
  }

  socket.onopen = () => {
    useRealtimeStore.setState({ connected: true });
  };

  socket.onmessage = (ev) => {
    try {
      const data = JSON.parse(ev.data);
      console.log('[realtime] onmessage raw:', typeof data, data && data.type);
      if (data && typeof data.type === 'string') {
        dispatch({ type: data.type, data: data.data ?? data });
      }
    } catch (err) {
      console.warn('[realtime] bad JSON frame', err);
    }
  };

  socket.onclose = () => {
    useRealtimeStore.setState({ connected: false });
    activeSocket = null;
    // Only reconnect if we still have subscribers and the token hasn't
    // been swapped (logout / re-login).
    if (lastToken === token && SUBSCRIBERS.size > 0) {
      scheduleReconnect(token);
    }
  };

  socket.onerror = () => {
    try {
      socket.close();
    } catch {
      /* noop */
    }
  };

  activeSocket = socket;
}

function ensureConnected() {
  const token = readToken();
  if (!token) {
    lastToken = null;
    if (activeSocket) {
      try { activeSocket.close(); } catch { /* noop */ }
      activeSocket = null;
    }
    useRealtimeStore.setState({ connected: false });
    return;
  }
  if (token !== lastToken || !activeSocket) {
    if (activeSocket) {
      try { activeSocket.close(); } catch { /* noop */ }
      activeSocket = null;
    }
    lastToken = token;
    connect(token);
  }
}

export const useRealtimeStore = create<RealtimeState>((set, get) => ({
  connected: false,
  subscriberCount: 0,
  markedScriptReturnCount: 0,
  latestMarkedScript: null,

  subscribe: (type, handler) => {
    let bucket = SUBSCRIBERS.get(type);
    if (!bucket) {
      bucket = new Set();
      SUBSCRIBERS.set(type, bucket);
    }
    bucket.add(handler);
    set({ subscriberCount: getSubscriberCount() });
    ensureConnected();
    return () => {
      bucket!.delete(handler);
      if (bucket!.size === 0) SUBSCRIBERS.delete(type);
      set({ subscriberCount: getSubscriberCount() });
    };
  },

  resetMarkedScriptCount: () => set({ markedScriptReturnCount: 0, latestMarkedScript: null }),

  _touch: () => {
    set({ subscriberCount: getSubscriberCount() });
    ensureConnected();
  },

  _release: () => {
    set({ subscriberCount: getSubscriberCount() });
  },
}));

if (typeof window !== 'undefined') {
  // Debug handle for browser-bridge E2E probes. Lets tests confirm the
  // dashboard's useEffect actually called subscribe().
  (window as any).__realtimeStore = useRealtimeStore;
  // Also expose the internal dispatcher so E2E tests can synthesize events
  // without needing the WS round-trip.
  (window as any).__realtimeDispatch = dispatch;
}

function getSubscriberCount(): number {
  let n = 0;
  for (const set of SUBSCRIBERS.values()) n += set.size;
  return n;
}

/** Optional desktop notification when a marked script arrives. Caller is
 *  responsible for asking the user for permission first. */
export function notifyMarkedScript(title: string, body: string): void {
  if (typeof window === 'undefined' || !('Notification' in window)) return;
  if (Notification.permission !== 'granted') return;
  try {
    new Notification(title, { body, icon: '/logo.svg' });
  } catch {
    /* noop — some browsers throw when called from background tabs */
  }
}