import { useEffect, useRef } from 'react'
import type { Landmark2D, Landmark3D } from '../gesture/types'
import { PoseLandmarker } from '@mediapipe/tasks-vision'

interface Props {
  landmarks: Landmark3D[]
  landmarks2D?: Landmark2D[]
  size?: number
}

type View = {
  title: string
  getX: (l: Landmark3D) => number
  getY: (l: Landmark3D) => number
}

const VIEWS: View[] = [
  { title: 'FRONT (X\u00b7Y)', getX: (l) => l.x, getY: (l) => l.y },
  { title: 'SIDE (Z\u00b7Y)', getX: (l) => l.z, getY: (l) => l.y },
  { title: 'TOP (X\u00b7Z)', getX: (l) => l.x, getY: (l) => l.z },
]

const LEFT_ARM = new Set([11, 13, 15, 17, 19, 21])
const RIGHT_ARM = new Set([12, 14, 16, 18, 20, 22])
const LEFT_LEG = new Set([23, 25, 27, 29, 31])
const RIGHT_LEG = new Set([24, 26, 28, 30, 32])
const HEAD = new Set([0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10])

type BodyPart = 'left-arm' | 'right-arm' | 'left-leg' | 'right-leg' | 'head' | 'torso'

const PART_COLORS: Record<BodyPart, { solid: string; faint: string }> = {
  'left-arm':  { solid: '#00ffcc', faint: 'rgba(0, 255, 204, 0.5)' },
  'right-arm': { solid: '#4488ff', faint: 'rgba(68, 136, 255, 0.5)' },
  'left-leg':  { solid: '#ff8800', faint: 'rgba(255, 136, 0, 0.5)' },
  'right-leg': { solid: '#ffcc00', faint: 'rgba(255, 204, 0, 0.5)' },
  'head':      { solid: '#aaaaaa', faint: 'rgba(170, 170, 170, 0.4)' },
  'torso':     { solid: '#888888', faint: 'rgba(136, 136, 136, 0.3)' },
}

function bodyPart(i: number): BodyPart {
  if (LEFT_ARM.has(i)) return 'left-arm'
  if (RIGHT_ARM.has(i)) return 'right-arm'
  if (LEFT_LEG.has(i)) return 'left-leg'
  if (RIGHT_LEG.has(i)) return 'right-leg'
  if (HEAD.has(i)) return 'head'
  return 'torso'
}

function connectionColor(a: number, b: number): string {
  const pa = bodyPart(a), pb = bodyPart(b)
  if (pa === pb) return PART_COLORS[pa].faint
  // mixed connections (e.g. shoulder-hip) — use the more colorful end
  if (pa === 'torso') return PART_COLORS[pb].faint
  return PART_COLORS[pa].faint
}

function landmarkColor(i: number): string {
  return PART_COLORS[bodyPart(i)].solid
}

function drawView(
  ctx: CanvasRenderingContext2D,
  w: number,
  h: number,
  landmarks: Landmark3D[],
  view: View,
) {
  ctx.fillStyle = 'rgba(0, 0, 0, 0.85)'
  ctx.fillRect(0, 0, w, h)

  if (landmarks.length < 33) {
    ctx.fillStyle = '#666'
    ctx.font = '14px monospace'
    ctx.fillText(view.title + ' (no data)', 8, 20)
    return
  }

  const VIS_THRESHOLD = 0.3
  const visible = landmarks.map((l) => (l.visibility ?? 1) >= VIS_THRESHOLD)
  const projected = landmarks.map((l) => ({ x: view.getX(l), y: view.getY(l) }))

  // Auto-scale: fit only visible landmarks into canvas with padding
  const xs = projected.filter((_, i) => visible[i]).map((p) => p.x)
  const ys = projected.filter((_, i) => visible[i]).map((p) => p.y)
  const minX = Math.min(...xs), maxX = Math.max(...xs)
  const minY = Math.min(...ys), maxY = Math.max(...ys)
  const range = Math.max(maxX - minX, maxY - minY) || 0.001
  const pad = 30
  const scale = (Math.min(w, h) - pad * 2) / range
  const cx = w / 2, cy = h / 2
  const midX = (minX + maxX) / 2, midY = (minY + maxY) / 2

  const toScreen = (p: { x: number; y: number }) => ({
    x: cx + (p.x - midX) * scale,
    y: cy + (p.y - midY) * scale,
  })

  // Crosshair through center of mass
  ctx.strokeStyle = '#333'
  ctx.lineWidth = 1
  ctx.setLineDash([4, 4])
  ctx.beginPath()
  ctx.moveTo(cx, pad); ctx.lineTo(cx, h - pad)
  ctx.moveTo(pad, cy); ctx.lineTo(w - pad, cy)
  ctx.stroke()
  ctx.setLineDash([])

  // Connections
  ctx.lineWidth = 1.5
  for (const conn of PoseLandmarker.POSE_CONNECTIONS) {
    if (!visible[conn.start] || !visible[conn.end]) continue
    const a = toScreen(projected[conn.start])
    const b = toScreen(projected[conn.end])
    ctx.strokeStyle = connectionColor(conn.start, conn.end)
    ctx.beginPath()
    ctx.moveTo(a.x, a.y)
    ctx.lineTo(b.x, b.y)
    ctx.stroke()
  }

  // Landmarks
  for (let i = 0; i < projected.length; i++) {
    if (!visible[i]) continue
    const p = toScreen(projected[i])
    ctx.fillStyle = landmarkColor(i)
    ctx.beginPath()
    ctx.arc(p.x, p.y, 3, 0, Math.PI * 2)
    ctx.fill()
  }

  // Title
  ctx.font = 'bold 14px monospace'
  ctx.fillStyle = '#fff'
  ctx.fillText(view.title, 8, 20)
}

const VIEW_2D: View = { title: '2D (X\u00b7Y)', getX: (l) => l.x, getY: (l) => l.y }

export function Pose3DView({ landmarks, landmarks2D, size = 220 }: Props) {
  const ref0 = useRef<HTMLCanvasElement>(null)
  const ref1 = useRef<HTMLCanvasElement>(null)
  const ref2 = useRef<HTMLCanvasElement>(null)
  const ref3 = useRef<HTMLCanvasElement>(null)

  useEffect(() => {
    const refs = [ref0, ref1, ref2]
    for (let i = 0; i < VIEWS.length; i++) {
      const canvas = refs[i].current
      if (!canvas) continue
      const ctx = canvas.getContext('2d')!
      drawView(ctx, size, size, landmarks, VIEWS[i])
    }

    if (ref3.current && landmarks2D) {
      const lm3d = landmarks2D.map((l) => ({ x: l.x, y: l.y, z: 0, visibility: l.visibility }))
      drawView(ref3.current.getContext('2d')!, size, size, lm3d, VIEW_2D)
    }
  }, [landmarks, landmarks2D, size])

  return (
    <div style={{
      position: 'fixed',
      right: 10,
      bottom: 10,
      zIndex: 100,
      display: 'flex',
      gap: 4,
    }}>
      <canvas ref={ref0} width={size} height={size} />
      <canvas ref={ref1} width={size} height={size} />
      <canvas ref={ref2} width={size} height={size} />
      <canvas ref={ref3} width={size} height={size} />
    </div>
  )
}
