geointel/frontend/src/components/segmentation/SegmentationLab.tsx

import type {
  DatasetCreateResponse,
  SegmentationModelCapability,
  SegmentationQaResult,
  SegmentationRead,
  SegmentationRunRead,
  SegmentationRunResponse,
} from '../../types'

interface SegmentationLabProps {
  segmentationModels: SegmentationModelCapability[]
  loadingSegmentationModels: boolean
  segmentationModelError: string | null
  selectedSegmentationDatasetId: string
  selectedSegmentationModelId: string
  segmentationTileManifestPath: string
  segmentationConfidenceThreshold: number
  runningSegmentation: boolean
  segmentationRunResult: SegmentationRunResponse | null
  segmentationRunError: string | null
  segmentationRuns: SegmentationRunRead[]
  selectedSegmentationRunId: string
  segmentationItems: SegmentationRead[]
  segmentationClassFilter: string
  segmentationMinConfidenceFilter: number
  loadingSegmentationResults: boolean
  segmentationReferenceDatasetId: string
  segmentationQaResult: SegmentationQaResult | null
  segmentationQaError: string | null
  runningSegmentationQa: boolean
  selectedProjectId: string | null
  rasterDatasets: DatasetCreateResponse[]
  referenceDatasets: DatasetCreateResponse[]
  selectedSegmentationModelConfigured: boolean
  selectedSegmentationModelLimitation: string | null
  onLoadModels: () => void
  onSelectDataset: (datasetId: string) => void
  onSelectModel: (modelId: string) => void
  onSetTileManifestPath: (value: string) => void
  onSetConfidenceThreshold: (value: number) => void
  onRunSegmentation: () => void
  onLoadRuns: () => void
  onSelectRun: (runId: string) => void
  onSetClassFilter: (value: string) => void
  onSetMinConfidenceFilter: (value: number) => void
  onLoadResults: () => void
  onSelectReferenceDataset: (datasetId: string) => void
  onRunQa: () => void
}

export function SegmentationLab({
  segmentationModels,
  loadingSegmentationModels,
  segmentationModelError,
  selectedSegmentationDatasetId,
  selectedSegmentationModelId,
  segmentationTileManifestPath,
  segmentationConfidenceThreshold,
  runningSegmentation,
  segmentationRunResult,
  segmentationRunError,
  segmentationRuns,
  selectedSegmentationRunId,
  segmentationItems,
  segmentationClassFilter,
  segmentationMinConfidenceFilter,
  loadingSegmentationResults,
  segmentationReferenceDatasetId,
  segmentationQaResult,
  segmentationQaError,
  runningSegmentationQa,
  selectedProjectId,
  rasterDatasets,
  referenceDatasets,
  selectedSegmentationModelConfigured,
  selectedSegmentationModelLimitation,
  onLoadModels,
  onSelectDataset,
  onSelectModel,
  onSetTileManifestPath,
  onSetConfidenceThreshold,
  onRunSegmentation,
  onLoadRuns,
  onSelectRun,
  onSetClassFilter,
  onSetMinConfidenceFilter,
  onLoadResults,
  onSelectReferenceDataset,
  onRunQa,
}: SegmentationLabProps): JSX.Element {
  const segmentationHasDataset = selectedSegmentationDatasetId.length > 0
  const segmentationHasTileManifest = segmentationTileManifestPath.trim().length > 0
  const segmentationModelUiRunnable =
    selectedSegmentationModelConfigured && selectedSegmentationModelId !== 'fixture-segmenter'
  const segmentationRunReady =
    Boolean(selectedProjectId) && segmentationHasDataset && segmentationModelUiRunnable
  const segmentationRunBlockedReason = !selectedProjectId
    ? 'Select or create a project first'
    : !segmentationHasDataset
      ? 'Select a raster dataset'
      : selectedSegmentationModelId === 'fixture-segmenter'
        ? 'Fixture segmenter is explicit test/demo-only'
        : !selectedSegmentationModelConfigured
          ? selectedSegmentationModelLimitation ?? 'Selected segmentation model is not configured'
          : null

  return (
    <section className="workspace-panel ai-lab-shell segmentation-lab-shell">
      <div className="panel-title-row">
        <div>
          <p className="eyebrow">Polygon segmentation</p>
          <h2>Segmentation Lab</h2>
        </div>
        <button className="secondary-action" type="button" onClick={onLoadModels} disabled={loadingSegmentationModels}>
          Refresh models
        </button>
      </div>

      <div className="ai-lab-model-surface" aria-label="Segmentation model capabilities">
        <div className="ai-lab-section-header">
          <div>
            <h3>Model registry</h3>
            <p>Backend-reported segmenter states and limitations.</p>
          </div>
        </div>
        <div className="ai-lab-state-stack">
          {loadingSegmentationModels ? (
            <div className="result-state result-state-loading">
              <strong>Loading segmentation models.</strong>
              <p>Checking backend model registry availability.</p>
            </div>
          ) : null}
          {segmentationModelError ? (
            <div className="result-state result-state-error">
              <strong>Segmentation model registry unavailable.</strong>
              <p>{segmentationModelError}</p>
            </div>
          ) : null}
          {segmentationModels.length === 0 && !loadingSegmentationModels ? (
            <div className="result-state result-state-empty">
              <strong>No segmentation models reported by backend.</strong>
              <p>Refresh models after the backend is reachable.</p>
            </div>
          ) : null}
        </div>
        <ul className="model-list">
          {segmentationModels.map((model) => (
            <li className={model.configured ? 'model-card model-card-ready' : 'model-card'} key={model.model_id}>
              <strong>{model.display_name}</strong>
              <span className={model.configured ? 'status-badge status-badge-ready' : 'status-badge'}>{model.status}</span>
              <div className="entity-meta">
                <span>{model.model_id}</span>
                <span>{model.framework}</span>
                <span>{model.task_type}</span>
              </div>
              <p className="muted">classes: {model.supported_classes.join(', ')}</p>
              <p className="muted">{model.limitation_message}</p>
            </li>
          ))}
        </ul>
      </div>

      <div className="lab-block">
        <div className="ai-lab-run-surface" aria-label="Segmentation run controls">
          <h3>Run segmentation</h3>
          <div
            className={segmentationRunReady ? 'lab-readiness-panel lab-readiness-panel-ready' : 'lab-readiness-panel'}
            aria-label="Segmentation run readiness"
          >
            <div className="ai-lab-section-header">
              <div>
                <h3>Run readiness</h3>
                <p>Checks the selected raster and segmenter state before submitting a segmentation job.</p>
              </div>
              <span className={segmentationRunReady ? 'status-badge status-badge-ready' : 'status-badge'}>
                {segmentationRunReady ? 'Ready to submit' : 'Blocked'}
              </span>
            </div>
            <div className="lab-readiness-grid">
              <div className={segmentationHasDataset ? 'lab-readiness-item lab-readiness-item-ready' : 'lab-readiness-item'}>
                <span>Raster dataset</span>
                <strong>{segmentationHasDataset ? 'Selected' : 'Select a raster dataset'}</strong>
              </div>
              <div className={selectedSegmentationModelConfigured ? 'lab-readiness-item lab-readiness-item-ready' : 'lab-readiness-item'}>
                <span>Model availability</span>
                <strong>
                  {selectedSegmentationModelConfigured
                    ? 'Selected model is configured'
                    : selectedSegmentationModelLimitation ?? 'Select a configured segmentation model'}
                </strong>
              </div>
              <div className={segmentationHasTileManifest ? 'lab-readiness-item lab-readiness-item-ready' : 'lab-readiness-item'}>
                <span>Tile manifest</span>
                <strong>{segmentationHasTileManifest ? 'Provided for provenance' : 'Optional for the fixture segmenter'}</strong>
              </div>
            </div>
          </div>
          <div className={segmentationRunReady ? 'lab-action-guardrail lab-action-guardrail-ready' : 'lab-action-guardrail'}>
            <span>Run action</span>
            <strong>{segmentationRunReady ? 'Ready to submit a segmentation job' : segmentationRunBlockedReason}</strong>
          </div>
          {rasterDatasets.length === 0 ? (
            <div className="result-state result-state-empty">
              <strong>No raster datasets available for segmentation.</strong>
              <p>Upload or select a raster dataset in Data before running segmentation.</p>
            </div>
          ) : null}
          <div className="lab-form-grid">
            <label>
              Raster dataset
              <select value={selectedSegmentationDatasetId} onChange={(event) => onSelectDataset(event.target.value)}>
                <option value="">Select raster dataset</option>
                {rasterDatasets.map((dataset) => (
                  <option key={dataset.id} value={dataset.id}>
                    {dataset.name}
                  </option>
                ))}
              </select>
            </label>
            <label>
              Model
              <select value={selectedSegmentationModelId} onChange={(event) => onSelectModel(event.target.value)}>
                {segmentationModels.map((model) => (
                  <option key={model.model_id} value={model.model_id}>
                    {model.display_name}
                  </option>
                ))}
              </select>
            </label>
            <label>
              Min confidence
              <input
                type="number"
                min="0"
                max="1"
                step="0.05"
                value={segmentationConfidenceThreshold}
                onChange={(event) => onSetConfidenceThreshold(Number(event.target.value))}
              />
            </label>
          </div>
          <label>
            Tile manifest
            <input
              type="text"
              placeholder="Raster tile manifest path"
              value={segmentationTileManifestPath}
              onChange={(event) => onSetTileManifestPath(event.target.value)}
            />
          </label>
          <button
            className="primary-action"
            type="button"
            onClick={onRunSegmentation}
            disabled={runningSegmentation || !segmentationRunReady}
          >
            Run segmentation
          </button>
        </div>
      </div>

      <div className="ai-lab-state-stack">
        {!selectedSegmentationModelConfigured ? (
          <div className="result-state result-state-empty">
            <strong>Segmentation model is not ready.</strong>
            <p>{selectedSegmentationModelLimitation ?? 'Select a configured segmentation model'}</p>
          </div>
        ) : null}
        {segmentationRunError ? (
          <div className="result-state result-state-error">
            <strong>Segmentation run failed.</strong>
            <p>{segmentationRunError}</p>
          </div>
        ) : null}
        {segmentationRunResult ? (
          <div className="result-summary-card">
            <p>Status: {segmentationRunResult.status}</p>
            <p>Message: {segmentationRunResult.message}</p>
            <p>Analysis run: {segmentationRunResult.analysis_run_id}</p>
            <p>Job: {segmentationRunResult.job_id}</p>
            <p>Segmentations: {segmentationRunResult.segmentation_count}</p>
            {segmentationRunResult.error_code ? <p className="error">Code: {segmentationRunResult.error_code}</p> : null}
          </div>
        ) : null}
      </div>

      <div className="ai-lab-results-surface" aria-label="Segmentation results">
        <div className="panel-title-row">
          <div>
            <h3>Segmentation results</h3>
            <p className="muted">Load persisted segmentation polygons and filter by class or confidence.</p>
          </div>
          <button className="secondary-action" type="button" onClick={onLoadRuns} disabled={!selectedProjectId}>
            Refresh runs
          </button>
        </div>
        <div className="lab-form-grid">
          <label>
            Run
            <select value={selectedSegmentationRunId} onChange={(event) => onSelectRun(event.target.value)}>
              <option value="">Select segmentation run</option>
              {segmentationRuns.map((run) => (
                <option key={run.id} value={run.id}>
                  {run.model_name || 'segmentation'} - {run.status} - {run.id}
                </option>
              ))}
            </select>
          </label>
          <label>
            Class
            <input
              type="text"
              placeholder="Class filter"
              value={segmentationClassFilter}
              onChange={(event) => onSetClassFilter(event.target.value)}
            />
          </label>
          <label>
            Min confidence
            <input
              type="number"
              min="0"
              max="1"
              step="0.05"
              value={segmentationMinConfidenceFilter}
              onChange={(event) => onSetMinConfidenceFilter(Number(event.target.value))}
            />
          </label>
        </div>
        <button className="primary-action" type="button" onClick={onLoadResults} disabled={!selectedSegmentationRunId || loadingSegmentationResults}>
          Load segmentations
        </button>
        {loadingSegmentationResults ? (
          <div className="result-state result-state-loading">
            <strong>Loading segmentation results.</strong>
            <p>Retrieving persisted segmentation polygons for the selected run.</p>
          </div>
        ) : null}
        <div className="ai-lab-state-stack">
          <div className="result-state result-state-ready">
            <strong>Segmentations loaded: {segmentationItems.length}</strong>
            <p>{selectedSegmentationRunId ? 'Loaded from persisted segmentation records.' : 'Select a segmentation run before loading results.'}</p>
          </div>
        </div>
        {segmentationItems.length > 0 ? (
          <div className="table-scroll">
            <table>
              <thead>
                <tr>
                  <th>Class</th>
                  <th>Confidence</th>
                  <th>Area m2</th>
                  <th>Model</th>
                  <th>Tile</th>
                  <th>Mask path</th>
                </tr>
              </thead>
              <tbody>
                {segmentationItems.map((segmentation) => (
                  <tr key={segmentation.id}>
                    <td>{segmentation.class_name}</td>
                    <td>{segmentation.confidence?.toFixed(2) ?? 'n/a'}</td>
                    <td>{segmentation.area_m2?.toFixed(2) ?? 'n/a'}</td>
                    <td>{segmentation.model_name}</td>
                    <td>{segmentation.source_tile_path || (segmentation.tile_index ?? 'n/a')}</td>
                    <td>{segmentation.mask_path || 'n/a'}</td>
                  </tr>
                ))}
              </tbody>
            </table>
          </div>
        ) : null}
      </div>

      <div className="ai-lab-qa-surface" aria-label="Segmentation QA controls and results">
        <h3>Segmentation QA</h3>
        <label>
          Reference dataset
          <select value={segmentationReferenceDatasetId} onChange={(event) => onSelectReferenceDataset(event.target.value)}>
            <option value="">Select reference dataset</option>
            {referenceDatasets.map((dataset) => (
              <option key={dataset.id} value={dataset.id}>
                {dataset.name}
              </option>
            ))}
          </select>
        </label>
        <button className="primary-action" type="button" onClick={onRunQa} disabled={runningSegmentationQa || !selectedSegmentationRunId || !segmentationReferenceDatasetId}>
          Compare segmentations to reference
        </button>
        {segmentationQaError ? (
          <div className="result-state result-state-error">
            <strong>Segmentation QA failed.</strong>
            <p>{segmentationQaError}</p>
          </div>
        ) : null}
        {segmentationQaResult ? (
          <div className="result-summary-card">
            <p>Status: {segmentationQaResult.status}</p>
            <p>Quality check: {segmentationQaResult.quality_check_id}</p>
            <p>Precision: {segmentationQaResult.precision?.toFixed(3) ?? 'n/a'}</p>
            <p>Recall: {segmentationQaResult.recall?.toFixed(3) ?? 'n/a'}</p>
            <p>F1: {segmentationQaResult.f1_score?.toFixed(3) ?? 'n/a'}</p>
            <p>Mean IoU: {segmentationQaResult.mean_iou?.toFixed(3) ?? 'n/a'}</p>
            <p>False positives: {segmentationQaResult.false_positives}</p>
            <p>False negatives: {segmentationQaResult.false_negatives}</p>
          </div>
        ) : null}
      </div>
    </section>
  )
}