frontend/src/components/QualityMetricsTable.tsx

import React, { useState, useEffect, useRef } from 'react'
import { ArrowDownTrayIcon } from '@heroicons/react/24/outline'
import MetricInfoIcon from './MetricInfoIcon'
import Descriptions from '../Descriptions'

interface Row {
  metric: string
  [key: string]: string | number
}

interface QualityMetricsTableProps {
  qualityMetrics: string[]
  tableHeader: string[]
  selectedModels: Set<string>
  tableRows: Row[]
}

const QualityMetricsTable: React.FC<QualityMetricsTableProps> = ({
  qualityMetrics,
  tableHeader,
  selectedModels,
  tableRows,
}) => {
  // Sorting state
  const [rowSort, setRowSort] = useState<{ metric: string; direction: 'asc' | 'desc' } | null>(null)
  const [columnSort, setColumnSort] = useState<{ model: string; direction: 'asc' | 'desc' } | null>(
    null
  )
  const [descriptionsLoaded, setDescriptionsLoaded] = useState(false)
  const descriptions = useRef(Descriptions.getInstance())

  // Load descriptions
  useEffect(() => {
    descriptions.current.load().then(() => setDescriptionsLoaded(true))
  }, [])

  // Handle row sort (sort columns by this metric)
  const handleRowSort = (metric: string) => {
    setRowSort((prev) => {
      if (!prev || prev.metric !== metric) return { metric, direction: 'asc' }
      if (prev.direction === 'asc') return { metric, direction: 'desc' }
      return null // Remove sort
    })
  }

  // Handle column sort (sort rows by this model)
  const handleColumnSort = (model: string) => {
    setColumnSort((prev) => {
      if (!prev || prev.model !== model) return { model, direction: 'asc' }
      if (prev.direction === 'asc') return { model, direction: 'desc' }
      return null // Remove sort
    })
  }

  // Sort models (columns)
  let sortedModels = tableHeader.filter((model) => selectedModels.has(model))
  if (rowSort) {
    // Sort columns by the value in the selected metric row
    sortedModels = [...sortedModels].sort((a, b) => {
      const row = tableRows.find((r) => r.metric === rowSort.metric)
      if (!row) return 0
      const valA = Number(row[a])
      const valB = Number(row[b])
      if (isNaN(valA) && isNaN(valB)) return 0
      if (isNaN(valA)) return 1
      if (isNaN(valB)) return -1
      return rowSort.direction === 'asc' ? valA - valB : valB - valA
    })
  }

  // Sort metrics (rows)
  let sortedMetrics = [...qualityMetrics]
  if (columnSort) {
    // Sort rows by the value in the selected model column
    sortedMetrics = [...sortedMetrics].sort((a, b) => {
      const rowA = tableRows.find((r) => r.metric === a)
      const rowB = tableRows.find((r) => r.metric === b)
      if (!rowA || !rowB) return 0
      const valA = Number(rowA[columnSort.model])
      const valB = Number(rowB[columnSort.model])
      if (isNaN(valA) && isNaN(valB)) return 0
      if (isNaN(valA)) return 1
      if (isNaN(valB)) return -1
      return columnSort.direction === 'asc' ? valA - valB : valB - valA
    })
  }

  // CSV export logic
  function exportToCSV() {
    // Build header row with model full names when available
    const header = [
      'Quality Metric',
      ...sortedModels.map(model =>
        descriptionsLoaded && descriptions.current.getModelAlias(model) || model
      )
    ]
    // Build data rows
    const rows = sortedMetrics
      .map((metric) => {
        const row = tableRows.find((r) => r.metric === metric)
        if (!row) return null
        return [
          metric,
          ...sortedModels.map((col) => {
            const cell = row[col]
            // Format as displayed
            return !isNaN(Number(cell)) ? Number(Number(cell).toFixed(3)) : cell
          }),
        ]
      })
      .filter((row): row is (string | number)[] => !!row)
    // Combine
    const csv = [header, ...rows]
      .map((row) => row.map((cell) => `"${String(cell).replace(/"/g, '""')}"`).join(','))
      .join('\n')
    // Download
    const blob = new Blob([csv], { type: 'text/csv' })
    const url = URL.createObjectURL(blob)
    const a = document.createElement('a')
    a.href = url
    a.download = 'quality_metrics.csv'
    document.body.appendChild(a)
    a.click()
    document.body.removeChild(a)
    URL.revokeObjectURL(url)
  }

  if (qualityMetrics.length === 0) return null
  return (
    <div className="max-h-[80vh]">
      <div className="flex justify-end">
        <button className="btn btn-ghost btn-circle" title="Export CSV" onClick={exportToCSV}>
          <ArrowDownTrayIcon className="h-6 w-6" />
        </button>
      </div>
      <table className="table w-full min-w-max border-gray-700 border">
        <thead>
          <tr>
            <th className="sticky left-0 top-0 bg-base-100 z-20 border-gray-700 border">
              Quality Metric
            </th>
            {sortedModels.map((model) => {
              const isSorted = columnSort && columnSort.model === model
              const direction = isSorted ? columnSort.direction : undefined
              return (
                <th
                  key={`quality-${model}`}
                  className="sticky top-0 bg-base-100 z-10 text-center text-xs border-gray-700 border cursor-pointer select-none"
                  onClick={() => handleColumnSort(model)}
                  title={
                    isSorted
                      ? direction === 'asc'
                        ? 'Sort descending'
                        : 'Clear sort'
                      : 'Sort by this column'
                  }
                >
                  {descriptionsLoaded && descriptions.current.getModelAlias(model) || model}
                  <span className="ml-1">{isSorted ? (direction === 'asc' ? '↑' : '↓') : '⇅'}</span>
                </th>
              )
            })}
          </tr>
        </thead>
        <tbody>
          {sortedMetrics.map((metric) => {
            const row = tableRows.find((r) => r.metric === metric)
            if (!row) return null
            const isSorted = rowSort && rowSort.metric === metric
            const direction = isSorted ? rowSort.direction : undefined
            return (
              <tr key={`quality-${metric}`} className="hover:bg-base-100 hover:z-30 relative">
                <td
                  className="sticky left-0 bg-base-100 z-10 hover:z-40 border-gray-700 border cursor-pointer select-none pr-4"
                  onClick={() => handleRowSort(metric)}
                  title={
                    isSorted
                      ? direction === 'asc'
                        ? 'Sort descending'
                        : 'Clear sort'
                      : 'Sort by this row (sorts columns)'
                  }
                >
                  <div className="flex items-center">
                    <span className="inline-block">{metric}</span>
                    <MetricInfoIcon metricName={metric} />
                    <span className="ml-auto">
                      {isSorted ? (direction === 'asc' ? '←' : '→') : '⇆'}
                    </span>
                  </div>
                </td>
                {sortedModels.map((col) => {
                  const cell = row[col]
                  return (
                    <td
                      key={`quality-${metric}-${col}`}
                      className="text-center border-gray-700 border hover:z-40 relative"
                    >
                      {!isNaN(Number(cell)) ? Number(Number(cell).toFixed(3)) : cell}
                    </td>
                  )
                })}
              </tr>
            )
          })}
        </tbody>
      </table>
    </div>
  )
}

export default QualityMetricsTable