import { defineChart } from '@tanstack/charts'
import { polar, radialArc } from '@tanstack/charts/polar'
import { flare } from '@charts-poc/demo-data/flare'
import { partition } from 'd3-hierarchy'
import { arc } from 'd3-shape'
import { selectSunburstData } from './selection'
import { flareHierarchy } from './transform'
import { tanstackMount } from '../../shared/mount'
import type { FlareRow } from '@charts-poc/demo-data/flare'
import type { ConformanceInput } from '../../types'
import type { HierarchyRectangularNode } from 'd3-hierarchy'

interface SunburstArcDatum {
  name: string
  value: number
  depth: number
  startAngle: number
  endAngle: number
  category: string
}

function topLevelCategory(node: HierarchyRectangularNode<FlareRow>): string {
  return (
    node.ancestors().find((ancestor) => ancestor.depth === 1)?.data.name ??
    'Other'
  )
}

const definition = (input: ConformanceInput) => {
  const root = flareHierarchy(selectSunburstData(flare, input.revision)).sum(
    (node) => node.size ?? 0,
  )
  const layout = partition<FlareRow>().size([Math.PI * 2, root.height + 1])(
    root,
  )
  const data: SunburstArcDatum[] = layout
    .descendants()
    .filter((node) => node.depth > 0)
    .map((node) => ({
      name: node.data.name,
      value: node.value ?? 0,
      depth: node.depth,
      startAngle: Math.PI / 2 - node.x0,
      endAngle: Math.PI / 2 - node.x1,
      category: topLevelCategory(node),
    }))

  return defineChart({
    marks: [
      polar({
        radiusRatio: 0.88,
        marks: [
          radialArc(data, {
            className: 'ts-chart__sunburst',
            generator: ({ radius }) => {
              const innerRadius = radius * 0.14
              const treeDepth = root.height + 1
              const thickness = (radius - innerRadius) / treeDepth
              const ringPadding = 2

              return arc<unknown, SunburstArcDatum>()
                .startAngle((node) => node.startAngle)
                .endAngle((node) => node.endAngle)
                .innerRadius(
                  (node) =>
                    innerRadius + (node.depth - 1) * (thickness + ringPadding),
                )
                .outerRadius(
                  (node) =>
                    innerRadius +
                    (node.depth - 1) * (thickness + ringPadding) +
                    thickness,
                )
            },
            color: 'category',
            stroke: '#ffffff',
            strokeWidth: 2,
          }),
        ],
      }),
    ],
    color: {
      range: ['#7c3aed', '#0ea5e9', '#14b8a6'],
    },
    margin: 0,
  })
}

export const mount = tanstackMount(definition, 'Flare analytics sunburst', {
  format: ({ datum }) =>
    `${datum.name.replaceAll('.', ' › ')} · ${datum.value.toLocaleString('en-US')}`,
})

import type { FlareRow } from '@charts-poc/demo-data/flare'

export function selectSunburstData(rows: readonly FlareRow[], revision = 0) {
  return revision % 2 === 0
    ? rows.slice(1, 12)
    : [...rows.slice(1, 8), ...rows.slice(9, 13)]
}

import { stratify } from 'd3-hierarchy'
import type { FlareRow } from '@charts-poc/demo-data/flare'
import type { HierarchyNode } from 'd3-hierarchy'

export interface SunburstTreeNode {
  [key: string]: unknown
  name: string
  size: number | null
  value: number
  children?: SunburstTreeNode[]
}

export function flareHierarchy(rows: readonly FlareRow[]) {
  const rootName = rows[0]?.name
  if (rootName === undefined) {
    throw new TypeError('The Flare sunburst selection is empty')
  }

  return stratify<FlareRow>().path((row) =>
    row.name.slice(rootName.length).replaceAll('.', '/'),
  )(Array.from(rows))
}

export function sunburstTree(rows: readonly FlareRow[]): SunburstTreeNode {
  return treeNode(flareHierarchy(rows))
}

function treeNode(node: HierarchyNode<FlareRow>): SunburstTreeNode {
  const children = node.children?.map(treeNode)
  const value =
    children === undefined
      ? (node.data.size ?? 0)
      : children.reduce((total, child) => total + child.value, 0)

  return {
    name: node.data.name,
    size: node.data.size,
    value,
    ...(children === undefined ? {} : { children }),
  }
}