bioai-platform/frontend/src/components/DockingViewer.tsx

'use client';

import { useCallback, useEffect, useRef, useState } from 'react';
import type { DockingInteraction } from '@/lib/api';
import { HudPanel, HudLegend } from '@/components/ui';
import { useTheme } from '@/contexts/theme';
import { ArrowCounterClockwise as RotateCcw, ArrowClockwise as RotateCw, Camera, FlipHorizontal, Palette, Cube as Box, Stack as Layers, Minus, Circle, Hexagon } from '@phosphor-icons/react';

interface DockingViewerProps {
  pdbId: string;
  ligandPdb: string;
  pdbUrl?: string;
  pdbUrlFormat?: 'pdb' | 'cif';
  interactions?: DockingInteraction;
  height?: number | string;
  backgroundColor?: string;
  chains?: Array<{ id: string; residue_count: number }>;
  ligands?: Array<{ id: string; chain: string; residue_count: number }>;
  highlightRange?: { start: number; end: number };
}

type RepresentationType = 'cartoon' | 'ball-and-stick' | 'spacefill' | 'gaussian-surface' | 'molecular-surface' | 'putty' | 'ribbon';
type ColorScheme = 'spectrum' | 'chain' | 'secondary-structure' | 'residue-type' | 'bfactor' | 'uniform';
type InteractionLayer = 'hbonds' | 'hydrophobic' | 'pi_stacking' | 'salt_bridges';

const INTERACTION_LAYERS: { key: InteractionLayer; label: string; color: string }[] = [
  { key: 'hbonds', label: 'H-bonds', color: '#4488ff' },
  { key: 'hydrophobic', label: 'Hydrophobic', color: '#ff8844' },
  { key: 'pi_stacking', label: 'Pi-stacking', color: '#ff44ff' },
  { key: 'salt_bridges', label: 'Salt bridges', color: '#00cccc' },
];

const REP_OPTIONS: { value: RepresentationType; label: string; icon: typeof Box }[] = [
  { value: 'cartoon', label: 'Cartoon', icon: Layers },
  { value: 'ribbon', label: 'Ribbon', icon: Minus },
  { value: 'ball-and-stick', label: 'Ball+Stick', icon: Box },
  { value: 'spacefill', label: 'Spacefill', icon: Circle },
  { value: 'gaussian-surface', label: 'Surface', icon: Hexagon },
  { value: 'putty', label: 'Putty', icon: Minus },
];

const COLOR_MAP: Record<ColorScheme, Record<string, unknown>> = {
  'spectrum': { name: 'spectrum' },
  'chain': { name: 'chain-id' },
  'secondary-structure': { name: 'secondary-structure' },
  'residue-type': { name: 'residue-type' },
  'bfactor': { name: 'bfactor' },
  'uniform': { name: 'uniform', value: 0x66ccff },
};

interface PDBeElement extends HTMLElement {
  viewerInstance?: {
    plugin?: {
      canvas3d?: {
        requestDraw: () => void;
        camera: {
          spin: (on: boolean | string, speed?: number) => void;
          reset: () => void;
        };
      };
      primitives?: {
        clear: () => void;
        add: (primitives: unknown[]) => void;
      };
    };
    visual: {
      update: (options: Record<string, unknown>, fullLoad?: boolean) => Promise<boolean>;
      select: (params: Record<string, unknown>) => Promise<void>;
      reset: (params: Record<string, unknown>) => Promise<void>;
      toggleSpin: (spin?: boolean) => Promise<void>;
      focus: (selection: unknown[]) => Promise<void>;
    };
    clear: () => Promise<void>;
    load: (params: Record<string, unknown>, fullLoad?: boolean) => Promise<void>;
    deleteStructure: (id?: string | number) => Promise<void>;
    events: {
      loadComplete: { subscribe: (cb: (success: boolean) => void) => void };
    };
  };
}

export function DockingViewer({
  pdbId,
  ligandPdb,
  pdbUrl,
  pdbUrlFormat,
  interactions,
  height = 480,
  backgroundColor,
  chains = [],
  ligands = [],
  highlightRange,
}: DockingViewerProps) {
  const { theme } = useTheme();
  const sceneBg = backgroundColor ?? (theme === 'light' ? '#FFFFFF' : '#0B0C14');
  const containerRef = useRef<HTMLDivElement>(null);
  const viewerRef = useRef<PDBeElement | null>(null);
  const initRef = useRef(false);
  const [status, setStatus] = useState<'loading' | 'ready' | 'error'>('loading');
  const [error, setError] = useState<string | null>(null);
  const [representation, setRepresentation] = useState<RepresentationType>('cartoon');
  const [colorScheme, setColorScheme] = useState<ColorScheme>('spectrum');
  const [spinning, setSpinning] = useState(false);
  const [selectedChain, setSelectedChain] = useState('');
  const [visibleInteractions, setVisibleInteractions] = useState<Record<InteractionLayer, boolean>>({
    hbonds: true,
    hydrophobic: true,
    pi_stacking: true,
    salt_bridges: true,
  });
  const representationRef = useRef<RepresentationType>('cartoon');
  const colorSchemeRef = useRef<ColorScheme>('spectrum');

  const applyVisualization = useCallback(async (rep: RepresentationType, color: ColorScheme) => {
    const viewer = viewerRef.current?.viewerInstance;
    if (!viewer) return;

    try {
      const options: Record<string, unknown> = {
        visualStyle: {
          polymer: { type: rep, color: { name: COLOR_MAP[color].name } },
          het: 'ball-and-stick',
          water: 'spacefill',
        },
      };
      if (pdbUrl) {
        options.customDataUrl = pdbUrl;
        options.customDataFormat = pdbUrlFormat || 'pdb';
      } else {
        options.moleculeId = pdbId.toLowerCase();
      }
      await viewer.visual.update(options, true);
    } catch {
      // Best-effort
    }
  }, [pdbId, pdbUrl, pdbUrlFormat]);

  const loadLigand = useCallback(async () => {
    const viewer = viewerRef.current?.viewerInstance;
    if (!viewer || !ligandPdb) return;

    try {
      const blob = new Blob([ligandPdb], { type: 'text/plain' });
      const url = URL.createObjectURL(blob);

      await viewer.load({
        url,
        format: 'pdb',
        id: 'ligand',
      }, false);

      URL.revokeObjectURL(url);
    } catch {
      // Best-effort
    }
  }, [ligandPdb]);

  const drawInteractions = useCallback((inter: DockingInteraction, visible: Record<InteractionLayer, boolean>) => {
    const plugin = viewerRef.current?.viewerInstance?.plugin;
    if (!plugin?.primitives) return;

    try {
      plugin.primitives.clear();
      const primitives: unknown[] = [];

      for (const hb of visible.hbonds ? inter.hbonds || [] : []) {
        if (hb.ligand_coords && hb.protein_coords) {
          primitives.push({
            kind: 'cylinder',
            start: { x: hb.ligand_coords[0], y: hb.ligand_coords[1], z: hb.ligand_coords[2] },
            end: { x: hb.protein_coords[0], y: hb.protein_coords[1], z: hb.protein_coords[2] },
            color: 0x4488ff,
            radius: 0.15,
          });
        }
      }

      for (const hp of visible.hydrophobic ? inter.hydrophobic || [] : []) {
        if (hp.ligand_coords && hp.protein_coords) {
          primitives.push({
            kind: 'cylinder',
            start: { x: hp.ligand_coords[0], y: hp.ligand_coords[1], z: hp.ligand_coords[2] },
            end: { x: hp.protein_coords[0], y: hp.protein_coords[1], z: hp.protein_coords[2] },
            color: 0xff8844,
            radius: 0.12,
          });
        }
      }

      for (const ps of visible.pi_stacking ? inter.pi_stacking || [] : []) {
        if (ps.ring_centroid && ps.ligand_centroid) {
          primitives.push({
            kind: 'cylinder',
            start: { x: ps.ring_centroid[0], y: ps.ring_centroid[1], z: ps.ring_centroid[2] },
            end: { x: ps.ligand_centroid[0], y: ps.ligand_centroid[1], z: ps.ligand_centroid[2] },
            color: 0xff44ff,
            radius: 0.18,
          });
        }
      }

      for (const sb of visible.salt_bridges ? inter.salt_bridges || [] : []) {
        if (sb.ligand_coords && sb.protein_coords) {
          primitives.push({
            kind: 'cylinder',
            start: { x: sb.ligand_coords[0], y: sb.ligand_coords[1], z: sb.ligand_coords[2] },
            end: { x: sb.protein_coords[0], y: sb.protein_coords[1], z: sb.protein_coords[2] },
            color: 0x00cccc,
            radius: 0.2,
          });
        }
      }

      if (primitives.length > 0) {
        plugin.primitives.add(primitives);
      }
      plugin.canvas3d?.requestDraw();
    } catch {
      // Best-effort
    }
  }, []);

  const selectChain = useCallback(async (chain: string) => {
    const viewer = viewerRef.current?.viewerInstance;
    if (!viewer) return;
    try {
      if (!chain) {
        await viewer.visual.reset({ camera: false, theme: false });
        return;
      }
      await viewer.visual.select({
        data: [{ auth_asym_id: chain, color: '#5eead4' }],
        nonSelectedColor: '#334155',
        keepRepresentations: true,
      });
      await viewer.visual.focus([{ auth_asym_id: chain }]);
    } catch {
      // The PDBe API can differ between deposited structures; preserve the canvas if a selection is unsupported.
    }
  }, []);

  const focusContacts = useCallback(async () => {
    const viewer = viewerRef.current?.viewerInstance;
    if (!viewer || !interactions) return;
    const residues = [
      ...interactions.hbonds,
      ...interactions.hydrophobic,
      ...interactions.pi_stacking,
      ...interactions.salt_bridges,
    ].map(contact => ({ auth_asym_id: contact.protein_chain, auth_seq_id: contact.protein_residue_seq }));
    if (!residues.length) return;
    try {
      await viewer.visual.select({ data: residues, keepRepresentations: true });
      await viewer.visual.focus(residues);
    } catch {
      // Best-effort viewer enhancement; the tabular fingerprint remains available.
    }
  }, [interactions]);

  const handleScreenshot = useCallback(() => {
    const canvas = viewerRef.current?.querySelector('canvas');
    if (!canvas) return;
    const link = document.createElement('a');
    link.download = `docking-${pdbId}-${Date.now()}.png`;
    link.href = canvas.toDataURL('image/png');
    link.click();
  }, [pdbId]);

  const handleResetCamera = useCallback(() => {
    viewerRef.current?.viewerInstance?.visual?.reset({ camera: true });
  }, []);

  const handleFlipView = useCallback(() => {
    const cam = viewerRef.current?.viewerInstance?.plugin?.canvas3d?.camera;
    if (!cam) return;
    cam.spin('x');
    setTimeout(() => cam.spin(false), 1500);
  }, []);

  // Initialize viewer ONCE per pdbId / pdbUrl
  useEffect(() => {
    if (!pdbId && !pdbUrl) {
      setStatus('error');
      setError('Missing structure identifier');
      return;
    }

    const container = containerRef.current;
    if (!container) return;
    if (initRef.current) return;
    initRef.current = true;

    let cancelled = false;
    let checkInterval: ReturnType<typeof setInterval> | null = null;

    container.innerHTML = '';

    const el = document.createElement('pdbe-molstar') as PDBeElement;
    if (pdbUrl) {
      el.setAttribute('custom-data-url', pdbUrl);
      el.setAttribute('custom-data-format', pdbUrlFormat || 'pdb');
    } else {
      el.setAttribute('molecule-id', pdbId.toLowerCase());
    }
    el.setAttribute('hide-controls', '');
    el.setAttribute('background-color', sceneBg);
    el.id = `pdbe-docking-${(pdbId || pdbUrl || 'structure').toLowerCase().replace(/[^a-z0-9]/gi, '')}`;
    container.appendChild(el);
    viewerRef.current = el;

    let cancelled_after_event = false;

    const onModelsLoaded = async () => {
      if (cancelled || cancelled_after_event) return;
      cancelled_after_event = true;
      if (checkInterval) clearInterval(checkInterval);

      const viewer = el.viewerInstance;
      if (!viewer) {
        if (!cancelled) setStatus('ready');
        return;
      }

      // Load ligand as separate structure
      await loadLigand();

      // Draw interaction lines
      if (interactions) drawInteractions(interactions, visibleInteractions);

      if (!cancelled) setStatus('ready');
    };

    el.addEventListener('molstar-models-loaded', onModelsLoaded);

    let attempts = 0;
    checkInterval = setInterval(() => {
      if (cancelled || cancelled_after_event) {
        if (checkInterval) clearInterval(checkInterval);
        return;
      }
      const hasCanvas = el.querySelector('canvas');
      if (hasCanvas) {
        onModelsLoaded();
        return;
      }
      attempts++;
      if (attempts > 100) {
        if (checkInterval) clearInterval(checkInterval);
        if (!cancelled) setStatus('ready');
      }
    }, 300);

    return () => {
      cancelled = true;
      cancelled_after_event = true;
      if (checkInterval) clearInterval(checkInterval);
      el.removeEventListener('molstar-models-loaded', onModelsLoaded);
    };
  }, [pdbId, pdbUrl, ligandPdb, sceneBg]); // eslint-disable-line react-hooks/exhaustive-deps

  // Highlight a residue range (deep-linked from motif scan on the structure page).
  // Uses PDBe Molstar's visual.select range format so the residues actually get
  // colored and focused on the canvas.
  const applyHighlight = useCallback(async () => {
    const viewer = viewerRef.current?.viewerInstance;
    if (!viewer || !highlightRange) return;
    const { start, end } = highlightRange;
    if (!Number.isFinite(start) || !Number.isFinite(end) || end < start) return;
    try {
      await viewer.visual.select({
        data: [{ start_residue_number: start, end_residue_number: end, color: '#fbbf24', focus: true }],
        nonSelectedColor: '#1e2436',
        keepRepresentations: true,
      });
    } catch {
      // Best-effort viewer enhancement; the tabular residue data remains available.
    }
  }, [highlightRange]);

  // Representation + color changes
  useEffect(() => {
    if (status === 'ready' && (representation !== representationRef.current || colorScheme !== colorSchemeRef.current)) {
      representationRef.current = representation;
      colorSchemeRef.current = colorScheme;
      void applyVisualization(representation, colorScheme).then(() => applyHighlight());
    }
  }, [representation, colorScheme, status, applyVisualization, applyHighlight]);

  // Interaction updates
  useEffect(() => {
    if (status === 'ready' && interactions) {
      drawInteractions(interactions, visibleInteractions);
    }
  }, [interactions, status, drawInteractions, visibleInteractions]);

  useEffect(() => {
    if (status === 'ready') {
      void applyHighlight();
    }
  }, [status, applyHighlight]);

  // Spin
  useEffect(() => {
    viewerRef.current?.viewerInstance?.visual?.toggleSpin(spinning);
  }, [spinning, status]);

  return (
    <div className="relative z-20 w-full overflow-hidden rounded-xl border border-glass-border bg-viewer">
      {/* HUD toolbar — spatial chrome for the viewer */}
      <HudPanel className="flex flex-wrap items-center gap-2 border-b border-glass-border px-3 py-2">
        <span className="mr-2 font-mono text-xs text-text-muted">{pdbId}</span>

        {highlightRange && (
          <span className="flex items-center gap-1.5 rounded border border-amber-400/40 bg-amber-400/10 px-2 py-1 text-[10px] text-amber-300">
            <span className="inline-block w-2 h-2 rounded-full bg-amber-400" />
            Highlight {highlightRange.start}–{highlightRange.end}
          </span>
        )}

        {/* Representation */}
        <div className="flex items-center gap-1">
          {REP_OPTIONS.map(opt => {
            const Icon = opt.icon;
            return (
              <button
                key={opt.value}
                type="button"
                title={opt.label}
                onClick={() => setRepresentation(opt.value)}
                disabled={status !== 'ready'}
                className={`p-1.5 rounded text-[10px] transition disabled:opacity-40 ${
                  representation === opt.value
                    ? 'bg-accent-cyan/20 text-accent-cyan'
                    : 'text-hud-text/50 hover:text-hud-text/80 hover:bg-hud/5'
                }`}
              >
                <Icon className="w-3.5 h-3.5" />
              </button>
            );
          })}
        </div>

        <div className="w-px h-5 bg-hud/10" />

        {chains.length > 0 && (
          <select
            value={selectedChain}
            onChange={(event) => { const chain = event.target.value; setSelectedChain(chain); selectChain(chain); }}
            disabled={status !== 'ready'}
            title="Focus protein chain"
            className="rounded border border-hud-line/10 bg-hud/30 px-2 py-1 text-[10px] text-hud-text/80 disabled:opacity-40"
          >
            <option value="">All chains</option>
            {chains.map(chain => <option key={chain.id} value={chain.id}>Chain {chain.id} · {chain.residue_count} residues</option>)}
          </select>
        )}

        {interactions && (
          <button type="button" onClick={focusContacts} disabled={status !== 'ready'} className="rounded border border-hud-line/10 px-2 py-1 text-[10px] text-hud-text/60 hover:bg-hud/5 hover:text-hud-text disabled:opacity-40">
            Focus pocket
          </button>
        )}

        {/* Color scheme */}
        <div className="relative">
          <select
            value={colorScheme}
            onChange={(e) => setColorScheme(e.target.value as ColorScheme)}
            disabled={status !== 'ready'}
            title="Color scheme"
            className="appearance-none rounded border border-hud-line/10 bg-hud/30 pl-6 pr-2 py-1 text-[10px] text-hud-text/80 disabled:opacity-40 cursor-pointer"
          >
            <option value="spectrum">Spectrum</option>
            <option value="chain">Chain</option>
            <option value="secondary-structure">Sec. Structure</option>
            <option value="residue-type">Residue Type</option>
            <option value="bfactor">B-factor</option>
            <option value="uniform">Uniform</option>
          </select>
          <Palette className="absolute left-1.5 top-1/2 -translate-y-1/2 w-3 h-3 text-hud-text/40 pointer-events-none" />
        </div>

        <div className="w-px h-5 bg-hud/10" />

        {/* Action buttons */}
        <button
          type="button"
          title="Toggle spin"
          onClick={() => setSpinning(s => !s)}
          disabled={status !== 'ready'}
          className={`p-1.5 rounded transition disabled:opacity-40 ${
            spinning ? 'text-accent-cyan bg-accent-cyan/10' : 'text-hud-text/50 hover:text-hud-text/80 hover:bg-hud/5'
          }`}
        >
          <RotateCcw className="w-3.5 h-3.5" />
        </button>

        <button
          type="button"
          title="Reset camera"
          onClick={handleResetCamera}
          disabled={status !== 'ready'}
          className="p-1.5 rounded text-hud-text/50 hover:text-hud-text/80 hover:bg-hud/5 transition disabled:opacity-40"
        >
          <RotateCw className="w-3.5 h-3.5" />
        </button>

        <button
          type="button"
          title="Flip view (X-axis)"
          onClick={handleFlipView}
          disabled={status !== 'ready'}
          className="p-1.5 rounded text-hud-text/50 hover:text-hud-text/80 hover:bg-hud/5 transition disabled:opacity-40"
        >
          <FlipHorizontal className="w-3.5 h-3.5" />
        </button>

        <button
          type="button"
          title="Screenshot"
          onClick={handleScreenshot}
          disabled={status !== 'ready'}
          className="p-1.5 rounded text-hud-text/50 hover:text-hud-text/80 hover:bg-hud/5 transition disabled:opacity-40"
        >
          <Camera className="w-3.5 h-3.5" />
        </button>
      </HudPanel>

      {/* Interaction legend */}
      {status === 'ready' && interactions && (
        <HudLegend className="flex flex-wrap items-center gap-2 border-b border-glass-border">
          <span className="mr-1 text-[10px] text-text-muted">Contacts</span>
          {INTERACTION_LAYERS.map(({ key, label, color }) => {
            const count = interactions[key]?.length || 0;
            if (!count) return null;
            const active = visibleInteractions[key];
            return (
              <button
                key={key}
                type="button"
                aria-pressed={active}
                onClick={() => setVisibleInteractions(current => ({ ...current, [key]: !current[key] }))}
                className={`flex items-center gap-1 rounded px-1.5 py-1 transition ${active ? 'bg-hud/10 text-hud-text/90' : 'text-hud-text/35 line-through'}`}
              >
                <span className="inline-block w-3 h-0.5" style={{ backgroundColor: color }} />
                {label} ({count})
              </button>
            );
          })}
        </HudLegend>
      )}

      {(chains.length > 0 || ligands.length > 0) && (
        <div className="flex flex-wrap gap-x-4 gap-y-1 border-b border-glass-border bg-viewer/60 px-3 py-2 text-[10px] text-text-muted">
          {chains.length > 0 && <span>{chains.length} protein chain{chains.length === 1 ? '' : 's'}</span>}
          {ligands.length > 0 && <span>Deposited ligands: {ligands.map(ligand => `${ligand.id}:${ligand.chain}`).join(', ')}</span>}
          {ligandPdb && <span className="text-accent-cyan">Docked ligand loaded</span>}
        </div>
      )}

      {/* Canvas */}
      <div ref={containerRef} style={{ height, minHeight: height, position: 'relative' }} />
      {status === 'loading' && (
        <div className="flex items-center justify-center py-4 text-sm text-hud-text/60">
          Loading structure...
        </div>
      )}
      {status === 'error' && (
        <div className="flex items-center justify-center py-4 px-6 text-center text-sm text-error">
          {error}
        </div>
      )}
    </div>
  );
}