<script lang="ts">
  import { PLOT_COLORS } from '../colors'
  import { StatusMessage } from '../feedback'
  import { create_file_drop_handler, drag_over_handlers } from '../io'
  import { format_value } from '../labels'
  import type { Vec2 } from '../math'
  import { BarPlot } from '../plot'
  import type {
    AxisConfig,
    BarHandlerProps,
    BarSeries,
    Orientation,
  } from '../plot/core/types'
  import type { AnyStructure } from '../structure'
  import type { BondingStrategy } from '../structure/bonding'
  import { parse_any_structure } from '../structure/parse'
  import { is_crystal } from '../structure/validation'
  import type { ComponentProps } from 'svelte'
  import { SvelteMap, SvelteSet } from 'svelte/reactivity'
  import { calc_coordination_nums, type CoordinationData } from './calc-coordination'
  import type { SplitMode } from './index'
  import { to_error } from '../utils'

  type CoordinationMetadata = {
    element?: string
    structure_label?: string
  } & Record<string, unknown>

  interface StructureEntry {
    label: string
    structure: AnyStructure
    color?: string
    data?: CoordinationData
  }
  let {
    structures,
    strategy = `electroneg_ratio`,
    split_mode = `by_element`,
    mode = $bindable(`grouped`),
    orientation = `vertical` as Orientation,
    x_axis = {},
    y_axis = {},
    allow_file_drop = true,
    on_file_drop,
    loading = $bindable(false),
    error_msg = $bindable(),
    ...rest
  }: Omit<ComponentProps<typeof BarPlot>, `series`> & {
    structures:
      | AnyStructure
      | Record<string, AnyStructure | { structure: AnyStructure; color?: string }>
      | StructureEntry[]
    strategy?: BondingStrategy
    split_mode?: SplitMode
    x_axis?: AxisConfig
    y_axis?: AxisConfig
    allow_file_drop?: boolean
    on_file_drop?: (content: string | ArrayBuffer, filename: string) => void
    loading?: boolean
    error_msg?: string
  } = $props()

  let dragover = $state(false)
  let dropped_entries = $state<StructureEntry[]>([])
  let is_horizontal = $derived(orientation === `horizontal`)

  // Normalize input to consistent array of { label, structure, color }
  const structure_entries = $derived.by<StructureEntry[]>(() => {
    if (!structures) return []

    const base_entries = Array.isArray(structures)
      ? (structures as StructureEntry[])
      : (is_crystal(structures)
        ? [{ label: `Structure`, structure: structures as AnyStructure }]
        : Object.entries(
          structures as Record<
            string,
            AnyStructure | { structure: AnyStructure; color?: string }
          >,
        ).map(([label, value]) =>
          `structure` in value
            ? { label, ...value }
            : { label, structure: value as AnyStructure }
        ))

    // Merge user-provided structures with dropped structures
    return [...base_entries, ...dropped_entries]
  })

  // Compute coordination data for each structure
  const entries_with_data = $derived(structure_entries.map((entry) => ({
    ...entry,
    data: calc_coordination_nums(entry.structure, strategy),
  })))

  // Compute appropriate ranges
  const ranges = $derived.by(() => {
    // CN axis should always start at 0 and show at least [0,4]
    let max_cn = 4 // minimum max value
    for (const entry of entries_with_data) {
      for (const [cn] of entry.data.cn_histogram) max_cn = Math.max(max_cn, cn)
    }
    const cn_range: Vec2 = [-0.5, max_cn + 0.5]

    return { count: [0, null] as [number, null], cn: cn_range } // Count axis should always start at 0
  })

  // Derive integer CN ticks for axis labels
  const cn_ticks = $derived.by(() => {
    const all_cns = new SvelteSet<number>()
    // Always include minimum CN values 0-4
    for (let idx = 0; idx <= 4; idx++) all_cns.add(idx)
    // Add actual CN values from data
    for (const entry of entries_with_data) {
      for (const [cn] of entry.data.cn_histogram) all_cns.add(cn)
    }
    return Array.from(all_cns).sort((cn1, cn2) => cn1 - cn2)
  })

  // Build BarPlot series based on split_mode
  const bar_series = $derived.by<BarSeries<CoordinationMetadata>[]>(() => {
    if (split_mode === `by_element`) {
      // One series per element across all structures
      const element_series_map = new SvelteMap<string, Map<number, number>>()

      // Collect all unique CNs across all elements
      const all_cns = new SvelteSet<number>()

      for (const entry of entries_with_data) {
        for (const [element, cn_histogram] of entry.data.cn_histogram_by_element) {
          if (!element_series_map.has(element)) {
            element_series_map.set(element, new SvelteMap())
          }
          const element_map = element_series_map.get(element)
          if (!element_map) continue

          for (const [cn, count] of cn_histogram) {
            all_cns.add(cn)
            element_map.set(cn, (element_map.get(cn) ?? 0) + count)
          }
        }
      }

      // Sort CNs for consistent x-axis
      const sorted_cns = Array.from(all_cns).sort((a, b) => a - b)

      // Convert map to array and ensure all series have same x-values
      return Array.from(element_series_map.entries())
        .sort((a, b) => a[0].localeCompare(b[0]))
        .map(([element, cn_map], idx) => {
          return {
            x: sorted_cns,
            y: sorted_cns.map((cn) => cn_map.get(cn) ?? 0),
            label: element,
            color: PLOT_COLORS[idx % PLOT_COLORS.length],
            bar_width: 0.8,
            visible: true,
            metadata: { element },
          }
        })
    } else if (split_mode === `by_structure`) {
      // One series per structure
      // First collect all unique CNs
      const all_cns = new SvelteSet<number>()
      for (const entry of entries_with_data) {
        for (const [cn] of entry.data.cn_histogram) {
          all_cns.add(cn)
        }
      }
      const sorted_cns = Array.from(all_cns).sort((a, b) => a - b)

      return entries_with_data.map((entry, idx) => {
        return {
          x: sorted_cns,
          y: sorted_cns.map((cn) => entry.data.cn_histogram.get(cn) ?? 0),
          label: entry.label,
          color: entry.color ?? PLOT_COLORS[idx % PLOT_COLORS.length],
          bar_width: 0.8,
          visible: true,
          metadata: { structure_label: entry.label },
        }
      })
    }
    // split_mode === 'none': combine all into single series
    const combined_histogram = new SvelteMap<number, number>()

    for (const entry of entries_with_data) {
      for (const [cn, count] of entry.data.cn_histogram) {
        combined_histogram.set(cn, (combined_histogram.get(cn) ?? 0) + count)
      }
    }

    const x_vals = Array.from(combined_histogram.keys()).sort((a, b) => a - b)
    const y_vals = x_vals.map((cn) => combined_histogram.get(cn) ?? 0)

    return [
      {
        x: x_vals,
        y: y_vals,
        label: `All Sites`,
        color: PLOT_COLORS[0],
        bar_width: 0.8,
        visible: true,
        metadata: {},
      },
    ]
  })

  const compute_and_add = (content: string | ArrayBuffer, filename: string) => {
    try {
      const text_content = content instanceof ArrayBuffer
        ? new TextDecoder().decode(content)
        : content
      const parsed_structure = parse_any_structure(text_content, filename)
      if (is_crystal(parsed_structure)) {
        dropped_entries = [{
          label: filename || `Dropped structure`,
          structure: parsed_structure,
        }, ...dropped_entries]
      } else {
        error_msg = `Structure has no lattice or sites; cannot compute coordination`
      }
    } catch (exc) {
      error_msg = `Failed to process structure: ${
        to_error(exc).message
      }`
    }
  }

  const handle_file_drop = create_file_drop_handler({
    allow: () => allow_file_drop,
    on_drop: (content, filename) =>
      (on_file_drop || compute_and_add)(content, filename),
    on_error: (msg) => {
      error_msg = msg
    },
    set_loading: (val) => {
      loading = val
      if (val) [error_msg, dragover] = [undefined, false]
    },
  })

  let display = $state({ x_zero_line: false, y_zero_line: false })
  // Update display when orientation changes
  $effect(() => {
    display.x_zero_line = is_horizontal
    display.y_zero_line = !is_horizontal
  })

  const cn_axis = { label: `Coordination Number`, format: `d` }
  const count_axis = { label: `Count`, format: `d` }
</script>

<StatusMessage bind:message={error_msg} type="error" dismissible />

{#if bar_series.length === 0}
  <StatusMessage
    message={allow_file_drop
    ? `Drag and drop structure files here to compute coordination numbers`
    : `No coordination data to display`}
  />
{:else}
  {#snippet tooltip(info: BarHandlerProps<CoordinationMetadata>)}
    {@const cn = info.x}
    {@const count = info.y}
    {@const element = info.metadata?.element}
    {@const structure_label = info.metadata?.structure_label}
    {#if element}
      {element} —
    {/if}
    {#if structure_label}
      {structure_label} —
    {/if}
    CN: {format_value(cn, `.0f`)}
    <br />
    Sites: {format_value(count, `.0f`)}
  {/snippet}

  <BarPlot
    {...rest}
    series={bar_series}
    bind:orientation
    bind:mode
    x_axis={{
      ...(is_horizontal ? count_axis : cn_axis),
      range: is_horizontal ? ranges.count : ranges.cn,
      ticks: is_horizontal ? undefined : cn_ticks,
      ...x_axis,
      label_shift: { y: 20, ...(is_horizontal ? y_axis : x_axis).label_shift },
    }}
    y_axis={{
      ...(is_horizontal ? cn_axis : count_axis),
      range: is_horizontal ? ranges.cn : ranges.count,
      ticks: is_horizontal ? cn_ticks : undefined,
      ...y_axis,
      label_shift: { x: 2, ...(is_horizontal ? x_axis : y_axis).label_shift },
    }}
    bind:display
    {tooltip}
    ondrop={handle_file_drop}
    {...drag_over_handlers({ allow: () => allow_file_drop, set_dragover: (over) => dragover = over })}
    class={(rest.class ?? ``) + (dragover ? ` dragover` : ``)}
  />
{/if}
