import { useMemo, useState } from 'react'
import {
  Box, Button, Checkbox, Chip, FormControl, FormControlLabel, InputLabel,
  LinearProgress, MenuItem, Select, Stack, Typography,
} from '@mui/material'
import { jsMatmul } from '../methods/jsMatmul'
import { webgpuMatmul } from '../methods/webgpuMatmul'
import { wasmNaiveMatmul, wasmBlockedMatmul, wasmBlockedMtMatmul } from '../methods/wasmMatmul'
import { blisStMatmul, blisMtMatmul } from '../methods/blisMatmul'
import { runInWorker } from '../methods/workerClient'
import { runInThreadedWorker } from '../methods/threadedClient'
import { generateMatrix } from '../methods/random'
import { gflops, type MatmulMethod } from '../methods/types'
import { ResultsTable, type CellState } from './ResultsTable'

const METHODS: MatmulMethod[] = [
  jsMatmul, webgpuMatmul, wasmNaiveMatmul, wasmBlockedMatmul, wasmBlockedMtMatmul,
  blisStMatmul, blisMtMatmul,
]
const ALL_SIZES = [128, 256, 512, 1024, 2048]

const HW = typeof navigator !== 'undefined' ? navigator.hardwareConcurrency || 4 : 4
const THREAD_OPTIONS = [...new Set([1, 2, 4, 8, HW])].filter((t) => t <= HW).sort((a, b) => a - b)
const isolated = typeof globalThis !== 'undefined' && globalThis.crossOriginIsolated === true

const yield_ = () => new Promise((r) => setTimeout(r, 0))

// Deterministic per-size seeds so every method (in this run) multiplies
// identical A and B — the `sample` (C[0]) values are then directly comparable.
function seedsFor(n: number) {
  return { seedA: n * 7 + 1, seedB: n * 13 + 2 }
}

type Results = Record<string, Record<number, CellState>>

export function BenchmarkRunner() {
  const [sizes, setSizes] = useState<Set<number>>(new Set(ALL_SIZES.filter((n) => n <= 1024)))
  const [threads, setThreads] = useState<number>(Math.min(8, HW))
  const [results, setResults] = useState<Results>({})
  const [running, setRunning] = useState(false)
  const [progress, setProgress] = useState(0)

  // Threaded column headers reflect the currently selected thread count:
  // the word "threaded" in the label becomes "N threads".
  const displayMethods = useMemo<MatmulMethod[]>(
    () => METHODS.map((m) => (
      m.label.includes('threaded')
        ? { ...m, label: m.label.replace('threaded', `${threads} threads`) }
        : m
    )),
    [threads],
  )

  const setCell = (methodId: string, n: number, state: CellState) => {
    setResults((prev) => ({
      ...prev,
      [methodId]: { ...prev[methodId], [n]: state },
    }))
  }

  const runOne = async (method: MatmulMethod, n: number) => {
    if (!method.available()) {
      setCell(method.id, n, 'unavailable')
      return
    }
    setCell(method.id, n, 'pending')
    await yield_()
    try {
      const { seedA, seedB } = seedsFor(n)
      let result
      if (method.threadedKind) {
        result = await runInThreadedWorker(method.threadedKind, n, seedA, seedB, threads)
      } else if (method.worker) {
        result = await runInWorker(method.id, n, seedA, seedB)
      } else {
        result = await method.run!(n, generateMatrix(n, seedA), generateMatrix(n, seedB))
      }
      setCell(method.id, n, { gflops: gflops(n, result.ms), ms: result.ms, sample: result.sample })
    } catch {
      setCell(method.id, n, 'error')
    }
  }

  const activeSizes = ALL_SIZES.filter((n) => sizes.has(n))

  const runAll = async () => {
    setRunning(true)
    const total = activeSizes.length * METHODS.length
    let done = 0
    for (const n of activeSizes) {
      for (const method of METHODS) {
        await runOne(method, n)
        done++
        setProgress((done / total) * 100)
      }
    }
    setRunning(false)
  }

  const runMethod = async (method: MatmulMethod) => {
    setRunning(true)
    const total = activeSizes.length
    let done = 0
    for (const n of activeSizes) {
      await runOne(method, n)
      done++
      setProgress((done / total) * 100)
    }
    setRunning(false)
  }

  const toggleSize = (n: number) => {
    setSizes((prev) => {
      const next = new Set(prev)
      if (next.has(n)) next.delete(n)
      else next.add(n)
      return next
    })
  }

  const hasResults = Object.values(results).some(
    (byN) => byN && Object.values(byN).some((c) => c && typeof c === 'object'),
  )

  // Tidy (long) CSV of every measured cell. Method labels carry the thread
  // count (from displayMethods), so no separate threads column is needed.
  const downloadCsv = () => {
    const esc = (s: string) => (/[",\n]/.test(s) ? `"${s.replace(/"/g, '""')}"` : s)
    const lines = ['n,method,precision,gflops,ms,checksum']
    for (const n of ALL_SIZES) {
      for (const m of displayMethods) {
        const r = results[m.id]?.[n]
        if (r && typeof r === 'object') {
          lines.push([n, esc(m.label), m.precision, r.gflops.toFixed(4), r.ms.toFixed(3), r.sample].join(','))
        }
      }
    }
    const blob = new Blob([lines.join('\n')], { type: 'text/csv' })
    const a = document.createElement('a')
    a.href = URL.createObjectURL(blob)
    a.download = 'matmul-bench-results.csv'
    a.click()
    URL.revokeObjectURL(a.href)
  }

  return (
    <Stack spacing={2}>
      <Typography variant="body2" color="text.secondary" sx={{ maxWidth: 900 }}>
        Multiplies two random n×n matrices with each method and reports GFLOP/s
        (2n³ ÷ time). All methods use identical inputs per size — see the
        cross-check below. <strong>WebGPU runs in single precision (f32)</strong>;
        every other in-browser method uses double precision (f64).
      </Typography>

      <Stack direction="row" spacing={1} sx={{ flexWrap: 'wrap', alignItems: 'center' }}>
        <Chip
          size="small"
          color={isolated ? 'success' : 'default'}
          label={isolated ? 'cross-origin isolated ✓' : 'not cross-origin isolated'}
        />
        <Typography variant="caption" color="text.secondary">
          {isolated
            ? 'WASM threads (SharedArrayBuffer) available — threaded methods enabled.'
            : 'Threaded methods need SharedArrayBuffer; they will show as n/a here.'}
        </Typography>
      </Stack>

      <Stack direction="row" spacing={1} sx={{ flexWrap: 'wrap', alignItems: 'center' }}>
        {ALL_SIZES.map((n) => (
          <FormControlLabel
            key={n}
            control={<Checkbox size="small" checked={sizes.has(n)} onChange={() => toggleSize(n)} disabled={running} />}
            label={`n=${n}`}
          />
        ))}
        <FormControl size="small" sx={{ minWidth: 180 }} disabled={running}>
          <InputLabel id="threads-label">threads (mt methods)</InputLabel>
          <Select
            labelId="threads-label"
            label="threads (mt methods)"
            value={threads}
            onChange={(e) => {
              // The measured mt numbers no longer match the new thread count
              // (and the column headers would relabel), so drop stale results.
              setThreads(Number(e.target.value))
              setResults({})
              setProgress(0)
            }}
          >
            {THREAD_OPTIONS.map((t) => (
              <MenuItem key={t} value={t}>{t} thread{t > 1 ? 's' : ''}</MenuItem>
            ))}
          </Select>
        </FormControl>
      </Stack>

      <Stack direction="row" spacing={1} sx={{ flexWrap: 'wrap', alignItems: 'center' }}>
        <Button variant="contained" onClick={runAll} disabled={running || activeSizes.length === 0}>
          {running ? 'Running…' : 'Run all methods'}
        </Button>
        {displayMethods.map((m) => (
          <Button key={m.id} size="small" variant="outlined" onClick={() => runMethod(m)} disabled={running || activeSizes.length === 0}>
            Run {m.label}
          </Button>
        ))}
      </Stack>
      {running && <LinearProgress variant="determinate" value={progress} />}

      <ResultsTable
        methods={displayMethods}
        sizes={activeSizes}
        cell={(methodId, n) => results[methodId]?.[n]}
      />

      <Box>
        <Button size="small" variant="outlined" onClick={downloadCsv} disabled={!hasResults}>
          Download CSV
        </Button>
      </Box>

      <Box>
        <Typography variant="subtitle2">Cross-check (C[0] for identical inputs)</Typography>
        <Typography variant="caption" color="text.secondary" component="div" sx={{ mb: 1 }}>
          The f64 methods (JS, the WASM kernels, both BLIS builds) should agree
          to ~1e-9; WebGPU (f32) will be close but not identical.
        </Typography>
        <Stack spacing={0.5}>
          {activeSizes.map((n) => (
            <Typography key={n} variant="caption" component="div">
              n={n}: {displayMethods.map((m) => {
                const r = results[m.id]?.[n]
                const v = r && typeof r === 'object' ? r.sample.toFixed(6) : '—'
                return `${m.label}=${v}`
              }).join('  ·  ')}
            </Typography>
          ))}
        </Stack>
      </Box>
    </Stack>
  )
}