src/app/chat/_components/StructureViewer.tsx

// StructureViewer.tsx — 3D protein structure viewer.
// Primary render path: dynamically imports 3Dmol.js from CDN at runtime (SSR-safe, zero npm dep).
// Fallback: if the CDN is unreachable or import fails, we render a lightweight CA-trace
// projection in SVG so the component still shows something meaningful.
"use client";

import { useEffect, useRef, useState } from "react";

export interface StructureViewerProps {
  url: string;
}

interface CAAtom {
  x: number;
  y: number;
  z: number;
  resName: string;
  resSeq: number;
}

function parsePdbCA(pdb: string): CAAtom[] {
  const out: CAAtom[] = [];
  for (const line of pdb.split(/\r?\n/)) {
    if (!line.startsWith("ATOM")) continue;
    const atomName = line.substring(12, 16).trim();
    if (atomName !== "CA") continue;
    const resName = line.substring(17, 20).trim();
    const resSeq = parseInt(line.substring(22, 26).trim(), 10) || 0;
    const x = parseFloat(line.substring(30, 38));
    const y = parseFloat(line.substring(38, 46));
    const z = parseFloat(line.substring(46, 54));
    if (Number.isFinite(x) && Number.isFinite(y) && Number.isFinite(z)) {
      out.push({ x, y, z, resName, resSeq });
    }
  }
  return out;
}

// Parse CA atoms from mmCIF _atom_site loop. Used only for the SVG fallback;
// 3Dmol handles CIF natively via addModel(data, "cif").
function parseCifCA(cif: string): CAAtom[] {
  const out: CAAtom[] = [];
  const lines = cif.split(/\r?\n/);

  // Find the _atom_site loop and extract column indices
  let inLoop = false;
  let inHeader = false;
  const columns: string[] = [];
  for (const line of lines) {
    const trimmed = line.trim();

    if (trimmed === "loop_") {
      inLoop = true;
      inHeader = true;
      columns.length = 0;
      continue;
    }

    if (inLoop && inHeader) {
      if (trimmed.startsWith("_atom_site.")) {
        columns.push(trimmed);
        continue;
      }
      // End of header — only proceed if this was an _atom_site loop
      if (columns.length === 0 || !columns[0].startsWith("_atom_site.")) {
        inLoop = false;
        inHeader = false;
        continue;
      }
      inHeader = false;
    }

    if (inLoop && !inHeader && columns.length > 0 && columns[0].startsWith("_atom_site.")) {
      // Data row or end of loop
      if (trimmed === "" || trimmed.startsWith("#") || trimmed.startsWith("_")) {
        inLoop = false;
        continue;
      }
      // A new loop_ line ends this loop but must fall through to the
      // loop_ detection above so _atom_site can be parsed if it follows
      // another loop without a '#' separator.
      if (trimmed === "loop_") {
        inLoop = false;
        inHeader = false;
        columns.length = 0;
        // Don't continue — fall through would skip the guard at line 52
        // because we already tested inLoop. Instead, re-enter directly.
        inLoop = true;
        inHeader = true;
        continue;
      }

      const tokens = trimmed.split(/\s+/);
      if (tokens.length < columns.length) continue;

      const get = (col: string): string | undefined => {
        const idx = columns.indexOf(col);
        return idx >= 0 ? tokens[idx] : undefined;
      };

      const group = get("_atom_site.group_PDB");
      const atomId = get("_atom_site.label_atom_id");
      if (group !== "ATOM" || atomId !== "CA") continue;

      const x = parseFloat(get("_atom_site.Cartn_x") ?? "");
      const y = parseFloat(get("_atom_site.Cartn_y") ?? "");
      const z = parseFloat(get("_atom_site.Cartn_z") ?? "");
      const resName = get("_atom_site.label_comp_id") ?? "";
      const resSeq = parseInt(get("_atom_site.label_seq_id") ?? "0", 10) || 0;

      if (Number.isFinite(x) && Number.isFinite(y) && Number.isFinite(z)) {
        out.push({ x, y, z, resName, resSeq });
      }
    }
  }
  return out;
}

// Detect structure format from URL extension and file content.
// Returns "cif" for mmCIF files, "pdb" otherwise.
function detectStructureFormat(url: string, content: string): "pdb" | "cif" {
  if (url.endsWith(".cif") || url.endsWith(".mmcif")) return "cif";
  if (content.trimStart().startsWith("data_")) return "cif";
  if (content.includes("_atom_site")) return "cif";
  return "pdb";
}

// Load 3Dmol.js from CDN once per page. Returns window.$3Dmol when resolved.
function load3Dmol(): Promise<unknown> {
  if (typeof window === "undefined") return Promise.reject(new Error("SSR"));
  const w = window as unknown as { $3Dmol?: unknown; __$3DmolLoader__?: Promise<unknown> };
  if (w.$3Dmol) return Promise.resolve(w.$3Dmol);
  if (w.__$3DmolLoader__) return w.__$3DmolLoader__;
  w.__$3DmolLoader__ = new Promise((resolve, reject) => {
    const existing = document.querySelector<HTMLScriptElement>(
      "script[data-3dmol-loader]",
    );
    const onLoad = () => {
      if (w.$3Dmol) resolve(w.$3Dmol);
      else reject(new Error("3Dmol loaded but global not found"));
    };
    if (existing) {
      existing.addEventListener("load", onLoad);
      existing.addEventListener("error", () => reject(new Error("3Dmol script error")));
      return;
    }
    const s = document.createElement("script");
    s.src = "https://cdn.jsdelivr.net/npm/3dmol@2.4.2/build/3Dmol-min.js";
    s.async = true;
    s.dataset.threedmolLoader = "true";
    s.setAttribute("data-3dmol-loader", "true");
    s.onload = onLoad;
    s.onerror = () => reject(new Error("Failed to load 3Dmol from CDN"));
    document.head.appendChild(s);
  });
  return w.__$3DmolLoader__;
}

export default function StructureViewer({ url }: StructureViewerProps) {
  const containerRef = useRef<HTMLDivElement | null>(null);
  const [mode, setMode] = useState<"loading" | "3dmol" | "fallback" | "error">("loading");
  const [error, setError] = useState<string | null>(null);
  const [atoms, setAtoms] = useState<CAAtom[]>([]);

  useEffect(() => {
    let cancelled = false;
    setMode("loading");
    setError(null);

    (async () => {
      let pdbText: string;
      try {
        const r = await fetch(url);
        if (!r.ok) throw new Error(`HTTP ${r.status}`);
        pdbText = await r.text();
      } catch (e) {
        if (!cancelled) {
          setError((e as Error).message);
          setMode("error");
        }
        return;
      }

      if (cancelled) return;

      const fmt = detectStructureFormat(url, pdbText);

      setAtoms(fmt === "cif" ? parseCifCA(pdbText) : parsePdbCA(pdbText));

      try {
        const mol = (await load3Dmol()) as {
          createViewer: (
            el: HTMLElement,
            opts: { backgroundColor: string },
          ) => {
            addModel: (data: string, fmt: string) => { setStyle: (sel: unknown, style: unknown) => void };
            setStyle: (sel: unknown, style: unknown) => void;
            zoomTo: () => void;
            render: () => void;
            resize: () => void;
          };
        };
        if (cancelled || !containerRef.current) return;
        const viewer = mol.createViewer(containerRef.current, { backgroundColor: "white" });
        viewer.addModel(pdbText, fmt);
        viewer.setStyle({}, { cartoon: { color: "spectrum" } });
        viewer.zoomTo();
        viewer.render();
        viewer.resize();
        if (!cancelled) setMode("3dmol");
      } catch (e) {
        // Fall back to SVG CA trace
        if (!cancelled) {
          setError((e as Error).message || "3Dmol unavailable");
          setMode("fallback");
        }
      }
    })();

    return () => {
      cancelled = true;
    };
  }, [url]);

  // SVG fallback: orthographic projection of CA trace onto XY plane.
  const svg = (() => {
    if (atoms.length === 0) return null;
    const xs = atoms.map((a) => a.x);
    const ys = atoms.map((a) => a.y);
    const minX = Math.min(...xs);
    const maxX = Math.max(...xs);
    const minY = Math.min(...ys);
    const maxY = Math.max(...ys);
    const w = 400;
    const h = 280;
    const pad = 16;
    const sx = (w - 2 * pad) / Math.max(1e-6, maxX - minX);
    const sy = (h - 2 * pad) / Math.max(1e-6, maxY - minY);
    const scale = Math.min(sx, sy);
    const pts = atoms
      .map((a) => {
        const px = pad + (a.x - minX) * scale;
        const py = h - pad - (a.y - minY) * scale;
        return `${px.toFixed(1)},${py.toFixed(1)}`;
      })
      .join(" ");
    return { w, h, pts };
  })();

  return (
    <div
      data-testid="structure-viewer"
      className="relative w-full max-h-96 overflow-hidden rounded-lg border border-border-subtle bg-white"
    >
      <div className="flex items-center justify-between px-3 py-2 border-b border-border-subtle">
        <span className="text-[10px] font-bold uppercase tracking-widest text-muted-fg">
          3D structure{atoms.length > 0 ? ` · ${atoms.length} residues` : ""}
          {mode === "fallback" && (
            <span className="ml-2 text-[9px] normal-case text-amber-600">
              (CA trace fallback)
            </span>
          )}
        </span>
        <a
          href={url}
          download
          className="text-[10px] text-muted-fg hover:text-accent underline underline-offset-2"
        >
          Download
        </a>
      </div>
      <div className="relative p-2">
        {mode === "loading" && (
          <div className="p-4 text-xs text-muted-fg">Loading structure…</div>
        )}
        {mode === "error" && (
          <div className="p-4 text-xs text-red-600">
            Failed to load structure: {error}
          </div>
        )}
        {/* 3Dmol container: always rendered while in 3dmol mode so the lib can attach */}
        <div
          ref={containerRef}
          style={{
            width: "100%",
            height: mode === "3dmol" ? 320 : 0,
            position: "relative",
          }}
        />
        {mode === "fallback" && svg && (
          <svg
            viewBox={`0 0 ${svg.w} ${svg.h}`}
            className="w-full h-auto max-h-80"
            role="img"
            aria-label="CA trace"
          >
            <polyline
              points={svg.pts}
              fill="none"
              stroke="#0ea5e9"
              strokeWidth={1.5}
              strokeLinejoin="round"
              strokeLinecap="round"
            />
          </svg>
        )}
      </div>
    </div>
  );
}

export { StructureViewer };