import type {
  AutoGroupResult,
  ClassSchema,
  ColumnInfo,
  MetadataRow,
  NumericBinning,
  SampleClass,
} from '../../types/auto-group'
import type { SampleGroup } from '../../types/components'
import type { TreeNode } from '../../types/componentLabTypes'
import { DEFAULT_COLORS, QC_OVERLAY_COLOR } from './colors'
import { classKey } from './classKey'

export interface ComposeInput {
  tokenizedSamples: string[][]
  sampleNames: string[]
  schemas: Record<string, ClassSchema>
  classes: SampleClass[]
  prefixClassLabels?: boolean
  missingGroupName?: string
}

/** Apply the operation pipeline to a single raw token value:
 *   numeric-parse → alias → exclude → bin
 *  Returns null if the row is excluded.
 */
function applyOps(col: ColumnInfo, raw: string): string | null {
  let value = raw

  // alias (operates on raw source value)
  if (col.ops?.alias && raw in col.ops.alias) {
    value = col.ops.alias[raw]
  }

  // exclude (post-alias)
  if (col.ops?.exclude?.includes(raw) || col.ops?.exclude?.includes(value)) {
    return null
  }

  // bin (only when role === 'numeric' and binning.mode === 'binned')
  if (col.role === 'numeric' && col.binning?.mode === 'binned') {
    const parsed = col.numeric?.byValue[raw]
    if (parsed && col.binning.edges) {
      const bin = pickBin(parsed.value, col.binning)
      if (bin !== null) value = bin
    }
  }

  return value
}

function pickBin(value: number, b: NumericBinning): string | null {
  if (!b.edges || b.edges.length === 0) return null
  for (let i = 0; i < b.edges.length; i++) {
    if (value < b.edges[i]) {
      return b.labels?.[i] ?? `< ${b.edges[i]}`
    }
  }
  const last = b.edges[b.edges.length - 1]
  return b.labels?.[b.edges.length] ?? `> ${last}`
}

export function composeGroups(input: ComposeInput): AutoGroupResult {
  const experimentalGroups: SampleGroup[] = []
  const qcGroups: SampleGroup[] = []
  const excludedSamples: string[] = []
  const metadata: MetadataRow[] = []
  let colorIdx = 0

  // Class-label prefixing exists ONLY to disambiguate group names when several
  // experimental classes could emit the same factor key (e.g. a Liver class and
  // a Kidney class both producing "5min / High"). With a single experimental
  // class there is nothing to disambiguate, so groups start directly at the
  // user's first layer and the sample selector's major-group slot is not
  // wasted on a constant root. When a prefix IS needed it should occupy ONE
  // hierarchy layer, so prefer the short distinguishing `subKind` ("Liver")
  // over the full label ("Biological / Liver") — falling back to the full
  // label when two classes share a subKind (e.g. token-length splits).
  const experimentalClasses = input.classes.filter(
    c => c.disposition === 'group' && input.schemas[classKey(c)],
  )
  const classPrefix = new Map<SampleClass, string>()
  if (input.prefixClassLabels !== false && experimentalClasses.length > 1) {
    const shortCounts = new Map<string, number>()
    for (const c of experimentalClasses) {
      const short = c.subKind ?? c.label
      shortCounts.set(short, (shortCounts.get(short) ?? 0) + 1)
    }
    for (const c of experimentalClasses) {
      const short = c.subKind ?? c.label
      classPrefix.set(c, (shortCounts.get(short) ?? 0) > 1 ? c.label : short)
    }
  }

  // Nested hierarchy (class → groupBy levels → sample leaves) built alongside
  // the flat groups. `treeIndex` lets repeated members reuse the same path node.
  // PATH_SEP is a control char so it never collides with a real token value.
  const PATH_SEP = String.fromCharCode(1)
  const groupTree: TreeNode[] = []
  const treeIndex = new Map<string, TreeNode>()
  let treeColorIdx = 0

  for (const cls of input.classes) {
    const schema = input.schemas[classKey(cls)]
    if (!schema) continue

    if (cls.disposition === 'exclude') {
      for (const m of cls.members) excludedSamples.push(input.sampleNames[m])
      continue
    }

    // Per-class group map. group key built from enabled (groupBy) columns,
    // joined with ' / ' for hierarchical display.
    const groupMap = new Map<string, { name: string; samples: string[] }>()
    const colByIdx = new Map(schema.columns.map(c => [c.index, c]))

    // Only experimental (group) classes form the nested tree; overlay/QC stay flat.
    const buildsTree = cls.disposition === 'group'
    const runCol = schema.columns.find(c => c.role === 'run-order')
    const classNodeId = classKey(cls)
    let classNode: TreeNode | null = null

    for (const m of cls.members) {
      const tokens = input.tokenizedSamples[m]
      const rowName = input.sampleNames[m]
      let excluded = false
      const fieldValues: Record<string, string> = {}

      for (const col of schema.columns) {
        const sourceValues = col.sourceIndices.map(idx => tokens[idx] ?? '')
        const raw = sourceValues.join('_') // merged columns join their source tokens
        const display = applyOps(col, raw)
        if (display === null) {
          excluded = true
          break
        }
        fieldValues[col.displayName ?? col.name] = display
      }

      if (excluded) {
        excludedSamples.push(rowName)
        continue
      }

      const keyParts: string[] =
        cls.disposition === 'group'
          ? schema.groupBy.map(idx => {
              const col = colByIdx.get(idx)!
              return fieldValues[col.displayName ?? col.name]
            })
          : [cls.label]

      // Prepend the class's short distinguishing label (see `classPrefix`
      // above) so group names from different classes don't collide when they
      // share a factor value (e.g. two tissue classes both emitting "5min").
      //
      // Only fires for disposition === 'group' classes (`classPrefix` holds no
      // other class): overlay/exclude classes already use cls.label AS their
      // single key part, so prefixing again would produce "IQC / IQC". A class
      // with one condition has no factor keyParts and is named by the prefix
      // alone ("Brain" next to "Liver / Ctrl"). The Unknown class opts out —
      // its label ("Unknown") adds no signal.
      const hasMissingGroupPart =
        input.missingGroupName &&
        cls.disposition === 'group' &&
        keyParts.some(part => part.trim().length === 0)
      const prefix = cls.kind !== 'unknown' ? classPrefix.get(cls) : undefined
      const groupKey = hasMissingGroupPart
        ? input.missingGroupName!
        : [...(prefix ? [prefix] : []), ...keyParts].join(' / ') || cls.label

      const bucket = groupMap.get(groupKey)
      if (bucket) {
        bucket.samples.push(rowName)
      } else {
        groupMap.set(groupKey, { name: groupKey, samples: [rowName] })
      }

      metadata.push({ sampleName: rowName, fields: fieldValues, group: groupKey })

      // Grow the nested tree: class root → one node per groupBy column → leaf.
      if (buildsTree) {
        if (!classNode) {
          classNode = {
            id: classNodeId,
            label: cls.label,
            badge: 0,
            children: [],
            metadata: { color: DEFAULT_COLORS[treeColorIdx++ % DEFAULT_COLORS.length] },
          }
          treeIndex.set(classNodeId, classNode)
          groupTree.push(classNode)
        }
        classNode.badge = (classNode.badge as number) + 1
        let cursor = classNode
        let pathId = classNodeId
        // Reuse the already-computed groupBy values (`keyParts`) so the tree
        // levels can never drift from the flat group key.
        for (const val of keyParts) {
          pathId += PATH_SEP + val
          let node = treeIndex.get(pathId)
          if (!node) {
            node = { id: pathId, label: val, badge: 0, children: [] }
            cursor.children!.push(node)
            treeIndex.set(pathId, node)
          }
          node.badge = (node.badge as number) + 1
          cursor = node
        }
        // Injection value goes through the column's op pipeline (alias/exclude)
        // via `fieldValues`, matching what every other surface shows.
        const injection = runCol ? fieldValues[runCol.displayName ?? runCol.name] : undefined
        const leaf: TreeNode = { id: pathId + PATH_SEP + rowName, label: rowName, type: 'sample' }
        if (injection) {
          leaf.badge = injection
          leaf.metadata = { injection }
        }
        cursor.children!.push(leaf)
      }
    }

    for (const { name, samples } of groupMap.values()) {
      const color =
        cls.disposition === 'overlay'
          ? QC_OVERLAY_COLOR
          : DEFAULT_COLORS[colorIdx++ % DEFAULT_COLORS.length]
      const g: SampleGroup = { name, color, samples }
      if (cls.disposition === 'overlay') qcGroups.push(g)
      else experimentalGroups.push(g)
    }
  }

  return {
    groups: [...experimentalGroups, ...qcGroups],
    experimentalGroups,
    qcGroups,
    metadata,
    excludedSamples,
    schemas: Object.values(input.schemas),
    groupTree,
  }
}