frontend/src/components/StructureViewer.tsx

import { useEffect, useRef, useState } from 'react';
import { Box } from 'lucide-react';

/* PDBe Mol* is ~6 MB, so it is loaded from the CDN only when a viewer is opened. */
const MOLSTAR_VERSION = '3.12.0';
const BASE = `https://cdn.jsdelivr.net/npm/pdbe-molstar@${MOLSTAR_VERSION}/build`;

interface MolstarInstance {
  render: (el: HTMLElement, options: Record<string, unknown>) => Promise<void> | void;
  visual: {
    select: (params: Record<string, unknown>) => Promise<void> | void;
    tooltips?: (params: Record<string, unknown>) => Promise<void> | void;
  };
  events: { loadComplete: { subscribe: (cb: (ok: boolean) => void) => { unsubscribe?: () => void } } };
  plugin?: {
    dispose?: () => void;
    canvas3d?: { setProps: (props: Record<string, unknown>) => void };
  };
}

declare global {
  interface Window {
    PDBeMolstarPlugin?: new () => MolstarInstance;
  }
}

let loader: Promise<void> | null = null;
function loadMolstar(): Promise<void> {
  if (window.PDBeMolstarPlugin) return Promise.resolve();
  if (!loader) {
    loader = new Promise<void>((resolve, reject) => {
      const css = document.createElement('link');
      css.rel = 'stylesheet';
      css.href = `${BASE}/pdbe-molstar-light.css`;
      document.head.appendChild(css);
      const s = document.createElement('script');
      s.src = `${BASE}/pdbe-molstar-plugin.js`;
      s.async = true;
      s.onload = () => resolve();
      s.onerror = () => {
        loader = null;
        reject(new Error('Could not load the Mol* viewer'));
      };
      document.head.appendChild(s);
    });
  }
  return loader;
}

function hexToRgb(hex: string) {
  const n = parseInt(hex.replace('#', ''), 16);
  return { r: (n >> 16) & 255, g: (n >> 8) & 255, b: n & 255 };
}

function isDark() {
  const t = document.documentElement.dataset.theme;
  if (t) return t === 'dark';
  return window.matchMedia?.('(prefers-color-scheme: dark)').matches ?? false;
}

export function StructureViewer({
  pdbId,
  chain,
  color = '#2a78d6',
  height = 380,
  label,
}: {
  pdbId: string;
  chain: string;
  color?: string;
  height?: number;
  label?: string;
}) {
  const ref = useRef<HTMLDivElement>(null);
  const [status, setStatus] = useState<'loading' | 'ready' | 'error'>('loading');

  useEffect(() => {
    let alive = true;
    let sub: { unsubscribe?: () => void } | undefined;
    let instance: MolstarInstance | undefined;
    const el = ref.current;
    setStatus('loading');
    // Mol* reports nothing when WebGL is unavailable; don't spin forever
    const timeout = setTimeout(() => alive && setStatus((st) => (st === 'loading' ? 'error' : st)), 90000);
    loadMolstar()
      .then(() => {
        if (!alive || !el || !window.PDBeMolstarPlugin) return;
        instance = new window.PDBeMolstarPlugin();
        sub = instance.events.loadComplete.subscribe((ok) => {
          if (!alive) return;
          if (ok === false) {
            setStatus('error');
            return;
          }
          setStatus('ready');
          instance?.visual.select({
            data: [{ auth_asym_id: chain, color: hexToRgb(color), focus: true }],
            nonSelectedColor: isDark() ? { r: 70, g: 78, b: 92 } : { r: 214, g: 219, b: 226 },
          });
        });
        const bg = isDark() ? { r: 18, g: 24, b: 33 } : { r: 255, g: 255, b: 255 };
        instance.render(el, {
          moleculeId: pdbId.toLowerCase(),
          bgColor: bg,
          hideControls: true,
          sequencePanel: false,
          landscape: true,
          visualStyle: 'cartoon',
          hideCanvasControls: ['expand', 'selection', 'animation', 'controlToggle', 'controlInfo'],
          loadingOverlay: true,
        });
      })
      .catch(() => alive && setStatus('error'));
    return () => {
      alive = false;
      clearTimeout(timeout);
      sub?.unsubscribe?.();
      instance?.plugin?.dispose?.();
      if (el) el.innerHTML = '';
    };
  }, [pdbId, chain, color]);

  return (
    <figure className="viewer" style={{ height }}>
      <div ref={ref} className="viewer__canvas" />
      {status !== 'ready' && (
        <div className="viewer__overlay">
          {status === 'loading' ? (
            <>
              <span className="spinner" aria-hidden /> Loading {pdbId.toUpperCase()}…
            </>
          ) : (
            <>
              <Box aria-hidden /> 3D view unavailable for {pdbId.toUpperCase()}
            </>
          )}
        </div>
      )}
      {label && <figcaption className="viewer__label">{label}</figcaption>}
    </figure>
  );
}

export interface ResidueColorRun {
  chain: string;
  start: number;
  end: number;
  color: string;
  /** 0–1; below 1 draws the run semi-transparent */
  opacity?: number;
}

export interface ResidueTooltip {
  chain: string;
  pos: number;
  text: string;
}

/**
 * Mol* view of a custom coordinate file (e.g. two superposed chains), coloured by
 * residue runs. Residue numbers are the file's auth_seq_id.
 */
export function CustomStructureViewer({
  url,
  runs,
  height = 460,
  label,
  fallbackColor = '#c3c9d2',
  tooltips = [],
}: {
  url: string;
  runs: ResidueColorRun[];
  height?: number;
  label?: string;
  fallbackColor?: string;
  /** per-residue hover text, shown in place of the generic Mol* label */
  tooltips?: ResidueTooltip[];
}) {
  const ref = useRef<HTMLDivElement>(null);
  const [status, setStatus] = useState<'loading' | 'ready' | 'error'>('loading');
  const runsKey = JSON.stringify(runs);
  const tipsKey = JSON.stringify(tooltips);

  useEffect(() => {
    let alive = true;
    let sub: { unsubscribe?: () => void } | undefined;
    let instance: MolstarInstance | undefined;
    const el = ref.current;
    setStatus('loading');
    const timeout = setTimeout(() => alive && setStatus((st) => (st === 'loading' ? 'error' : st)), 90000);
    const data = (JSON.parse(runsKey) as ResidueColorRun[]).map((r) => ({
      auth_asym_id: r.chain,
      beg_auth_seq_id: r.start,
      end_auth_seq_id: r.end,
      color: hexToRgb(r.color),
      ...(r.opacity != null && r.opacity < 1 ? { opacity: r.opacity } : {}),
    }));
    const tips = (JSON.parse(tipsKey) as ResidueTooltip[]).map((t) => ({
      auth_asym_id: t.chain,
      beg_auth_seq_id: t.pos,
      end_auth_seq_id: t.pos,
      tooltip: t.text,
    }));
    loadMolstar()
      .then(() => {
        if (!alive || !el || !window.PDBeMolstarPlugin) return;
        instance = new window.PDBeMolstarPlugin();
        sub = instance.events.loadComplete.subscribe((ok) => {
          if (!alive) return;
          if (ok === false) {
            setStatus('error');
            return;
          }
          setStatus('ready');
          // semi-transparent runs stay hoverable: Mol* only picks geometry at or
          // above this opacity (default 0.5)
          instance?.plugin?.canvas3d?.setProps({ renderer: { pickingAlphaThreshold: 0.05 } });
          instance?.visual.select({ data, nonSelectedColor: hexToRgb(fallbackColor) });
          if (tips.length) instance?.visual.tooltips?.({ data: tips });
        });
        const bg = isDark() ? { r: 18, g: 24, b: 33 } : { r: 255, g: 255, b: 255 };
        instance.render(el, {
          customData: { url: new URL(url, window.location.href).toString(), format: 'pdb', binary: false },
          bgColor: bg,
          hideControls: true,
          sequencePanel: false,
          landscape: true,
          visualStyle: 'cartoon',
          hideCanvasControls: ['expand', 'selection', 'animation', 'controlToggle', 'controlInfo'],
          loadingOverlay: true,
        });
      })
      .catch(() => alive && setStatus('error'));
    return () => {
      alive = false;
      clearTimeout(timeout);
      sub?.unsubscribe?.();
      instance?.plugin?.dispose?.();
      if (el) el.innerHTML = '';
    };
  }, [url, runsKey, tipsKey, fallbackColor]);

  return (
    <figure className="viewer" style={{ height }}>
      <div ref={ref} className="viewer__canvas" />
      {status !== 'ready' && (
        <div className="viewer__overlay">
          {status === 'loading' ? (
            <>
              <span className="spinner" aria-hidden /> Loading structures…
            </>
          ) : (
            <>
              <Box aria-hidden /> 3D view unavailable
            </>
          )}
        </div>
      )}
      {label && <figcaption className="viewer__label">{label}</figcaption>}
    </figure>
  );
}