import { useState } from 'react'
import { useQuery, useMutation, useQueryClient } from '@tanstack/react-query'
import { toast } from 'sonner'
import { RefreshCw, ChevronDown, ChevronRight, Cpu, Layers, Minus, Plus } from 'lucide-react'
import {
  getCeleryWorkers,
  getQueueStatus,
  scaleWorkers,
  type CeleryWorker,
  type ScaleRequest,
} from '../api/worker'

// ---------------------------------------------------------------------------
// Worker card
// ---------------------------------------------------------------------------

function WorkerCard({ worker }: { worker: CeleryWorker }) {
  const [expanded, setExpanded] = useState(false)
  return (
    <div className="rounded-xl border border-border-default p-4">
      <div className="flex items-start justify-between gap-3">
        <div className="flex items-center gap-2 min-w-0">
          <Cpu size={16} className="text-accent shrink-0" />
          <span className="text-sm font-medium text-content truncate">{worker.name}</span>
        </div>
        <div className="flex items-center gap-2 shrink-0">
          <span
            className={`text-xs font-semibold px-2 py-0.5 rounded-full ${
              worker.active_task_count > 0
                ? 'bg-blue-500/20 text-blue-400'
                : 'bg-green-500/20 text-green-400'
            }`}
          >
            {worker.active_task_count > 0 ? `${worker.active_task_count} active` : 'idle'}
          </span>
          {worker.active_tasks.length > 0 && (
            <button
              onClick={() => setExpanded((e) => !e)}
              className="text-content-muted hover:text-content transition-colors"
            >
              {expanded ? <ChevronDown size={14} /> : <ChevronRight size={14} />}
            </button>
          )}
        </div>
      </div>

      {/* Queues */}
      <div className="mt-2 flex flex-wrap gap-1">
        {worker.queues.map((q) => (
          <span
            key={q}
            className="text-xs px-2 py-0.5 rounded bg-surface-muted text-content-muted"
          >
            {q}
          </span>
        ))}
      </div>

      {/* Active tasks */}
      {expanded && worker.active_tasks.length > 0 && (
        <div className="mt-3 space-y-1">
          {worker.active_tasks.map((t) => (
            <div key={t.id} className="text-xs text-content-muted font-mono truncate">
              {t.name}
            </div>
          ))}
        </div>
      )}
    </div>
  )
}

// ---------------------------------------------------------------------------
// Scale controls
// ---------------------------------------------------------------------------

type ScalableService = ScaleRequest['service']

const SCALABLE_SERVICES: { service: ScalableService; label: string; description: string }[] = [
  { service: 'render-worker', label: 'Render Worker', description: 'Blender renders — concurrency=1' },
  { service: 'worker', label: 'Step Worker', description: 'STEP processing — concurrency=8' },
  { service: 'worker-thumbnail', label: 'Thumbnail Worker', description: 'Thumbnail rendering' },
]

function ScaleControl({
  service,
  label,
  description,
}: {
  service: ScalableService
  label: string
  description: string
}) {
  const [count, setCount] = useState(1)
  const scaleMut = useMutation({
    mutationFn: () => scaleWorkers({ service, count }),
    onSuccess: (data) => toast.success(`${data.service} → ${data.count} instance(s)`),
    onError: (e: unknown) => {
      const detail = (e as { response?: { data?: { detail?: string } } })?.response?.data?.detail
      toast.error(detail ?? `Failed to scale ${service}`)
    },
  })

  return (
    <div className="rounded-xl border border-border-default p-4 flex items-center justify-between gap-4">
      <div>
        <p className="text-sm font-medium text-content">{label}</p>
        <p className="text-xs text-content-muted mt-0.5">{description}</p>
      </div>
      <div className="flex items-center gap-2 shrink-0">
        <button
          onClick={() => setCount((c) => Math.max(0, c - 1))}
          className="p-1 rounded-md bg-surface-muted hover:bg-surface-hover text-content transition-colors"
        >
          <Minus size={14} />
        </button>
        <span className="w-6 text-center text-sm font-semibold text-content">{count}</span>
        <button
          onClick={() => setCount((c) => Math.min(20, c + 1))}
          className="p-1 rounded-md bg-surface-muted hover:bg-surface-hover text-content transition-colors"
        >
          <Plus size={14} />
        </button>
        <button
          onClick={() => scaleMut.mutate()}
          disabled={scaleMut.isPending}
          className="btn-primary text-xs px-3 py-1.5 ml-2"
        >
          {scaleMut.isPending ? 'Scaling…' : 'Scale'}
        </button>
      </div>
    </div>
  )
}

// ---------------------------------------------------------------------------
// Queue depth bar
// ---------------------------------------------------------------------------

function QueueDepthRow({ queue, depth }: { queue: string; depth: number }) {
  return (
    <div className="flex items-center gap-3">
      <span className="text-sm text-content w-44 truncate font-mono">{queue}</span>
      <div className="flex-1 h-2 rounded-full bg-surface-muted overflow-hidden">
        <div
          className="h-full rounded-full transition-all"
          style={{
            width: `${Math.min(100, depth * 5)}%`,
            backgroundColor: depth > 10 ? 'var(--color-red-500)' : 'var(--color-accent)',
          }}
        />
      </div>
      <span
        className={`text-xs font-semibold w-8 text-right ${
          depth > 10 ? 'text-red-400' : 'text-content-muted'
        }`}
      >
        {depth}
      </span>
    </div>
  )
}

// ---------------------------------------------------------------------------
// Main page
// ---------------------------------------------------------------------------

export default function WorkerManagement() {
  const qc = useQueryClient()

  const { data: workerData, isLoading: workersLoading } = useQuery({
    queryKey: ['celery-workers'],
    queryFn: getCeleryWorkers,
    refetchInterval: 10_000,
  })

  const { data: queueData, isLoading: queuesLoading } = useQuery({
    queryKey: ['queue-status'],
    queryFn: getQueueStatus,
    refetchInterval: 5_000,
  })

  function refresh() {
    qc.invalidateQueries({ queryKey: ['celery-workers'] })
    qc.invalidateQueries({ queryKey: ['queue-status'] })
  }

  const workers = workerData?.workers ?? []
  const queueDepths = queueData?.queue_depths ?? {}

  return (
    <div className="p-8 space-y-8">
      {/* Header */}
      <div className="flex items-center justify-between">
        <div>
          <h1 className="text-2xl font-bold text-content">Worker Management</h1>
          <p className="text-sm text-content-muted mt-1">
            Monitor active Celery workers and scale services up or down.
          </p>
        </div>
        <button onClick={refresh} className="btn-secondary flex items-center gap-2 text-sm">
          <RefreshCw size={14} />
          Refresh
        </button>
      </div>

      {/* Queue depths */}
      <section>
        <div className="flex items-center gap-2 mb-3">
          <Layers size={16} className="text-accent" />
          <h2 className="text-base font-semibold text-content">Queue Depths</h2>
        </div>
        {queuesLoading ? (
          <div className="space-y-2">
            {[0, 1, 2].map((i) => (
              <div key={i} className="h-6 rounded bg-surface-muted animate-pulse" />
            ))}
          </div>
        ) : Object.keys(queueDepths).length === 0 ? (
          <p className="text-sm text-content-muted">No queue data available.</p>
        ) : (
          <div className="rounded-xl border border-border-default p-4 space-y-3">
            {Object.entries(queueDepths).map(([queue, depth]) => (
              <QueueDepthRow key={queue} queue={queue} depth={depth} />
            ))}
          </div>
        )}
      </section>

      {/* Active workers */}
      <section>
        <div className="flex items-center gap-2 mb-3">
          <Cpu size={16} className="text-accent" />
          <h2 className="text-base font-semibold text-content">
            Active Workers
            {workers.length > 0 && (
              <span className="ml-2 text-xs font-normal text-content-muted">
                ({workers.length})
              </span>
            )}
          </h2>
        </div>
        {workersLoading ? (
          <div className="grid grid-cols-2 gap-3">
            {[0, 1].map((i) => (
              <div key={i} className="h-20 rounded-xl bg-surface-muted animate-pulse" />
            ))}
          </div>
        ) : workerData?.error ? (
          <div className="rounded-xl border border-border-default p-4 text-sm text-red-400">
            Failed to fetch workers: {workerData.error}
          </div>
        ) : workers.length === 0 ? (
          <div className="rounded-xl border border-border-default p-4 text-sm text-content-muted">
            No active workers detected. Make sure Celery workers are running.
          </div>
        ) : (
          <div className="grid grid-cols-1 md:grid-cols-2 gap-3">
            {workers.map((w) => (
              <WorkerCard key={w.name} worker={w} />
            ))}
          </div>
        )}
      </section>

      {/* Scale controls */}
      <section>
        <h2 className="text-base font-semibold text-content mb-3">Scale Services</h2>
        <p className="text-xs text-content-muted mb-4">
          Adjust the number of container instances for each service via Docker Compose.
          Changes take effect immediately but are not persisted across deployments.
        </p>
        <div className="space-y-2">
          {SCALABLE_SERVICES.map((s) => (
            <ScaleControl key={s.service} {...s} />
          ))}
        </div>
      </section>
    </div>
  )
}
