HartOMat/frontend/src/components/cad/ThreeDViewer.tsx

import {
  Suspense,
  useRef,
  useCallback,
  useState,
  useEffect,
  Component,
  type ErrorInfo,
  type ReactNode,
} from 'react'
import { useQuery } from '@tanstack/react-query'
import { Canvas, useThree, useFrame } from '@react-three/fiber'
import { OrbitControls, useGLTF, Environment } from '@react-three/drei'
import { toast } from 'sonner'
import {
  X, Camera, Loader2, AlertTriangle, Box, Cpu, Download, ChevronDown,
} from 'lucide-react'
import api from '../../api/client'

// ---------------------------------------------------------------------------
// Types
// ---------------------------------------------------------------------------

export interface ThreeDViewerProps {
  cadFileId: string
  onClose: () => void
  /** URL for the geometry-only GLB (from OCC export) */
  geometryGltfUrl?: string
  /** URL for the production-quality GLB (Blender + PBR materials) */
  productionGltfUrl?: string
  /** Whether a geometry GLB exists (for hint display) */
  hasGeometryGlb?: boolean
  /** Whether a production GLB exists (for hint display) */
  hasProductionGlb?: boolean
  /** Called when the user clicks "Generate Geometry GLB" from the hint banner */
  onGenerateGeometry?: () => void
  /** Whether a geometry GLB generation is in progress */
  isGeneratingGeometry?: boolean
  /** Download URLs for assets */
  downloadUrls?: {
    glb?: string
    production?: string
    blend?: string
  }
}

type ViewMode = 'geometry' | 'production'

const ENV_PRESETS = ['city', 'sunset', 'dawn', 'night', 'warehouse', 'forest', 'apartment', 'studio', 'park', 'lobby'] as const
type EnvPreset = typeof ENV_PRESETS[number]

// ---------------------------------------------------------------------------
// Inner model loader – separated so Suspense can catch it
// ---------------------------------------------------------------------------

function GltfModel({ url, wireframe }: { url: string; wireframe: boolean }) {
  const { scene } = useGLTF(url)

  useEffect(() => {
    scene.traverse((child: any) => {
      if (child.isMesh) {
        child.material = child.material.clone()
        child.material.wireframe = wireframe
      }
    })
  }, [scene, wireframe])

  return <primitive object={scene} />
}

// ---------------------------------------------------------------------------
// Screenshot helper – lives inside Canvas so it can access gl / useThree
// ---------------------------------------------------------------------------

interface ScreenshotCaptureProps {
  enabled: boolean
  cadFileId: string
  onDone: () => void
}

function ScreenshotCapture({ enabled, cadFileId, onDone }: ScreenshotCaptureProps) {
  const { gl } = useThree()
  const didCapture = useRef(false)

  useFrame(() => {
    if (!enabled || didCapture.current) return
    didCapture.current = true

    const dataUrl = gl.domElement.toDataURL('image/png')
    const [header, base64Data] = dataUrl.split(',')
    const mimeMatch = header.match(/:(.*?);/)
    const mimeType = mimeMatch ? mimeMatch[1] : 'image/png'
    const byteCharacters = atob(base64Data)
    const byteArray = new Uint8Array(byteCharacters.length)
    for (let i = 0; i < byteCharacters.length; i++) {
      byteArray[i] = byteCharacters.charCodeAt(i)
    }
    const blob = new Blob([byteArray], { type: mimeType })
    const formData = new FormData()
    formData.append('thumbnail', blob, 'thumbnail.png')

    api
      .post(`/cad/${cadFileId}/regenerate-thumbnail`, formData, {
        headers: { 'Content-Type': 'multipart/form-data' },
      })
      .then(() => toast.success('Thumbnail captured and saved'))
      .catch(() => toast.error('Failed to save thumbnail'))
      .finally(() => {
        didCapture.current = false
        onDone()
      })
  })

  return null
}

// ---------------------------------------------------------------------------
// Error boundary
// ---------------------------------------------------------------------------

class GltfErrorBoundary extends Component<
  { children: ReactNode; onError: (msg: string) => void },
  { hasError: boolean }
> {
  constructor(props: { children: ReactNode; onError: (msg: string) => void }) {
    super(props)
    this.state = { hasError: false }
  }
  static getDerivedStateFromError(): { hasError: boolean } {
    return { hasError: true }
  }
  componentDidCatch(error: Error, _info: ErrorInfo): void {
    this.props.onError(error.message || 'Failed to parse GLTF')
  }
  render(): ReactNode {
    if (this.state.hasError) return null
    return this.props.children
  }
}

// ---------------------------------------------------------------------------
// Loading overlay
// ---------------------------------------------------------------------------

function LoadingOverlay() {
  return (
    <div className="absolute inset-0 flex flex-col items-center justify-center bg-gray-900 text-white gap-3 pointer-events-none z-10">
      <Loader2 size={40} className="animate-spin text-accent" />
      <p className="text-sm text-gray-300">Loading 3D model…</p>
    </div>
  )
}

// ---------------------------------------------------------------------------
// Model loader with ready tracking
// ---------------------------------------------------------------------------

interface ModelWithReadyProps {
  url: string
  wireframe: boolean
  onReady: () => void
}

function ModelWithReady({ url, wireframe, onReady }: ModelWithReadyProps) {
  const { scene } = useGLTF(url)

  useEffect(() => {
    scene.traverse((child: any) => {
      if (child.isMesh) {
        child.material = child.material.clone()
        child.material.wireframe = wireframe
      }
    })
  }, [scene, wireframe])

  useEffect(() => { onReady() }, [onReady])
  return <primitive object={scene} />
}

// ---------------------------------------------------------------------------
// Env preset dropdown
// ---------------------------------------------------------------------------

function EnvDropdown({
  value,
  onChange,
}: {
  value: EnvPreset
  onChange: (v: EnvPreset) => void
}) {
  const [open, setOpen] = useState(false)
  return (
    <div className="relative">
      <button
        onClick={() => setOpen((o) => !o)}
        className="flex items-center gap-1.5 px-3 py-1.5 rounded-md bg-gray-700 hover:bg-gray-600 text-white text-xs font-medium transition-colors"
      >
        {value}
        <ChevronDown size={12} />
      </button>
      {open && (
        <div className="absolute right-0 top-full mt-1 z-50 bg-gray-800 border border-gray-700 rounded-md shadow-xl min-w-[130px]">
          {ENV_PRESETS.map((p) => (
            <button
              key={p}
              onClick={() => { onChange(p); setOpen(false) }}
              className={`w-full text-left px-3 py-1.5 text-xs hover:bg-gray-700 transition-colors ${
                p === value ? 'text-accent font-semibold' : 'text-gray-300'
              }`}
            >
              {p}
            </button>
          ))}
        </div>
      )}
    </div>
  )
}

// ---------------------------------------------------------------------------
// Main exported component
// ---------------------------------------------------------------------------

export default function ThreeDViewer({
  cadFileId,
  onClose,
  geometryGltfUrl,
  productionGltfUrl,
  hasGeometryGlb,
  hasProductionGlb,
  onGenerateGeometry,
  isGeneratingGeometry,
  downloadUrls,
}: ThreeDViewerProps) {
  // Default to production mode if only production GLB is available
  const initialMode: ViewMode = productionGltfUrl && !geometryGltfUrl ? 'production' : 'geometry'
  const [mode, setMode] = useState<ViewMode>(initialMode)
  const [wireframe, setWireframe] = useState(false)
  const [envPreset, setEnvPreset] = useState<EnvPreset>('city')
  const [capturing, setCapturing] = useState(false)
  const [loadError, setLoadError] = useState<string | null>(null)
  const [modelReady, setModelReady] = useState(false)

  const { data: settings3d } = useQuery({
    queryKey: ['admin-settings'],
    queryFn: () => api.get('/admin/settings').then(r => r.data),
    staleTime: 60_000,
  })

  // Resolve the active model URL: prefer selected mode, fall back to whichever URL exists
  const activeUrl =
    mode === 'production' && productionGltfUrl
      ? productionGltfUrl
      : geometryGltfUrl ?? productionGltfUrl

  const handleModelReady = useCallback(() => setModelReady(true), [])
  const handleError = useCallback((msg: string) => setLoadError(msg), [])
  const handleCaptureDone = useCallback(() => setCapturing(false), [])

  // Reset ready state when URL changes
  useEffect(() => {
    setModelReady(false)
    setLoadError(null)
  }, [activeUrl])

  function handleDownload(url: string, filename: string) {
    const a = document.createElement('a')
    a.href = url
    a.download = filename
    document.body.appendChild(a)
    a.click()
    document.body.removeChild(a)
  }

  const hasBothModes = !!(geometryGltfUrl && productionGltfUrl)

  return (
    <div className="fixed inset-0 z-50 flex flex-col bg-gray-950">
      {/* Toolbar */}
      <div className="flex items-center justify-between px-5 py-3 bg-gray-900 border-b border-gray-800 shrink-0 gap-3 flex-wrap">
        <span className="text-white font-semibold tracking-wide">3D Viewer</span>

        <div className="flex items-center gap-2 flex-wrap">
          {/* Mode toggle */}
          {hasBothModes && (
            <div className="flex rounded-md overflow-hidden border border-gray-700">
              <button
                onClick={() => setMode('geometry')}
                className={`flex items-center gap-1.5 px-3 py-1.5 text-xs font-medium transition-colors ${
                  mode === 'geometry'
                    ? 'bg-accent text-white'
                    : 'bg-gray-800 text-gray-300 hover:bg-gray-700'
                }`}
              >
                <Box size={12} />
                Geometry
              </button>
              <button
                onClick={() => setMode('production')}
                className={`flex items-center gap-1.5 px-3 py-1.5 text-xs font-medium transition-colors ${
                  mode === 'production'
                    ? 'bg-accent text-white'
                    : 'bg-gray-800 text-gray-300 hover:bg-gray-700'
                }`}
              >
                <Cpu size={12} />
                Production
              </button>
            </div>
          )}

          {/* Wireframe toggle */}
          <button
            onClick={() => setWireframe((w) => !w)}
            className={`px-3 py-1.5 rounded-md text-xs font-medium transition-colors border ${
              wireframe
                ? 'bg-accent border-accent text-white'
                : 'bg-gray-800 border-gray-700 text-gray-300 hover:bg-gray-700'
            }`}
          >
            Wireframe
          </button>

          {/* Environment preset */}
          <EnvDropdown value={envPreset} onChange={setEnvPreset} />

          {/* Download buttons */}
          {downloadUrls?.glb && (
            <button
              onClick={() => handleDownload(downloadUrls.glb!, `${cadFileId}_geometry.glb`)}
              className="flex items-center gap-1.5 px-3 py-1.5 rounded-md bg-gray-700 hover:bg-gray-600 text-white text-xs font-medium transition-colors"
            >
              <Download size={12} />
              Geometry GLB
            </button>
          )}
          {downloadUrls?.production && (
            <button
              onClick={() => handleDownload(downloadUrls.production!, `${cadFileId}_production.glb`)}
              className="flex items-center gap-1.5 px-3 py-1.5 rounded-md bg-gray-700 hover:bg-gray-600 text-white text-xs font-medium transition-colors"
            >
              <Download size={12} />
              Production GLB
            </button>
          )}
          {downloadUrls?.blend && (
            <button
              onClick={() => handleDownload(downloadUrls.blend!, `${cadFileId}.blend`)}
              className="flex items-center gap-1.5 px-3 py-1.5 rounded-md bg-gray-700 hover:bg-gray-600 text-white text-xs font-medium transition-colors"
            >
              <Download size={12} />
              .blend
            </button>
          )}

          {/* Capture button */}
          <button
            onClick={() => setCapturing(true)}
            disabled={capturing || !modelReady || loadError !== null}
            className="flex items-center gap-2 px-4 py-1.5 rounded-md bg-accent hover:bg-accent-hover disabled:opacity-40 disabled:cursor-not-allowed text-white text-sm font-medium transition-colors"
          >
            {capturing ? <Loader2 size={15} className="animate-spin" /> : <Camera size={15} />}
            {capturing ? 'Capturing…' : 'Capture Angle'}
          </button>

          {/* Close */}
          <button
            onClick={onClose}
            className="p-1.5 rounded-md text-gray-400 hover:text-white hover:bg-gray-700 transition-colors"
            aria-label="Close viewer"
          >
            <X size={20} />
          </button>
        </div>
      </div>

      {/* Hint banners */}
      {!hasProductionGlb && (
        <div className="bg-amber-900/60 border-b border-amber-700/50 px-4 py-2 flex items-center gap-2 text-amber-200 text-xs shrink-0">
          <Cpu size={13} className="shrink-0" />
          <span>
            <strong>No Production GLB yet.</strong> Go to the product page and click "Generate Production GLB" to create a high-quality version with PBR materials and proper mesh smoothing.
          </span>
        </div>
      )}
      {!hasGeometryGlb && hasProductionGlb && onGenerateGeometry && (
        <div className="bg-blue-900/50 border-b border-blue-700/50 px-4 py-2 flex items-center gap-3 text-blue-200 text-xs shrink-0">
          <Box size={13} className="shrink-0" />
          <span>
            <strong>Showing Production GLB.</strong> Generate a Geometry GLB to enable the mode toggle and compare geometry vs. production quality.
          </span>
          {isGeneratingGeometry ? (
            <span className="flex items-center gap-1 text-blue-300 ml-auto shrink-0">
              <Loader2 size={11} className="animate-spin" />
              Generating…
            </span>
          ) : (
            <button
              onClick={onGenerateGeometry}
              className="ml-auto shrink-0 px-3 py-1 rounded bg-blue-700 hover:bg-blue-600 text-white text-xs font-medium transition-colors"
            >
              Generate Geometry GLB
            </button>
          )}
        </div>
      )}

      {/* Viewport */}
      <div className="relative flex-1">
        {loadError && (
          <div className="absolute inset-0 flex flex-col items-center justify-center bg-gray-900 text-white gap-4 z-20">
            <AlertTriangle size={48} className="text-red-400" />
            <p className="text-lg font-semibold">Failed to load 3D model</p>
            <p className="text-sm text-gray-400 max-w-sm text-center">{loadError}</p>
            <button
              onClick={onClose}
              className="mt-2 px-4 py-2 rounded-md bg-gray-700 hover:bg-gray-600 text-sm transition-colors"
            >
              Close
            </button>
          </div>
        )}

        {!modelReady && !loadError && <LoadingOverlay />}

        <Canvas
          camera={{ position: [0, 0.1, 0.3], fov: 45 }}
          gl={{ preserveDrawingBuffer: true }}
          style={{ width: '100%', height: '100%', background: '#111827' }}
        >
          <ambientLight intensity={0.5} />
          <directionalLight position={[5, 10, 7]} intensity={1.0} castShadow />
          <directionalLight position={[-5, -5, -5]} intensity={0.25} />

          {activeUrl && (
            <GltfErrorBoundary onError={handleError}>
              <Suspense fallback={null}>
                <ModelWithReady
                  key={activeUrl}
                  url={activeUrl}
                  wireframe={wireframe}
                  onReady={handleModelReady}
                />
              </Suspense>
            </GltfErrorBoundary>
          )}

          <OrbitControls
            enablePan
            enableZoom
            enableRotate
            minDistance={settings3d?.viewer_min_distance ?? 0.001}
            maxDistance={settings3d?.viewer_max_distance ?? 50}
          />
          <Environment preset={envPreset} />

          {capturing && (
            <ScreenshotCapture
              enabled={capturing}
              cadFileId={cadFileId}
              onDone={handleCaptureDone}
            />
          )}
        </Canvas>
      </div>
    </div>
  )
}