import { useMemo } from 'react'
import type { ApiGraphResponse, GraphNode3D, GraphEdge3D, TransformedGraph } from './types'
import { NODE_TYPE_COLORS, NODE_SIZE } from './constants'

const typeKeys = Object.keys(NODE_TYPE_COLORS)

function getColorIndex(type: string): number {
  const idx = typeKeys.indexOf(type?.toLowerCase())
  return idx >= 0 ? idx : typeKeys.length - 1
}

export function transformGraphData(raw: ApiGraphResponse | null): TransformedGraph | null {
  if (!raw?.data?.nodes) return null

  // nodes can be a dict {id: node} or an array — normalize to array
  const nodesRaw = raw.data.nodes
  const rawNodes: any[] = Array.isArray(nodesRaw)
    ? nodesRaw
    : Object.entries(nodesRaw).map(([k, v]: [string, any]) => ({ id: k, ...v }))
  const rawEdges: any[] = Array.isArray(raw.data.edges)
    ? raw.data.edges
    : Object.values(raw.data.edges || {})

  if (rawNodes.length === 0) return null

  // Build index map: id → array index
  const idToIndex = new Map<string, number>()
  rawNodes.forEach((n, i) => {
    idToIndex.set(String(n.id), i)
  })

  // Count degrees
  const degrees = new Float32Array(rawNodes.length)
  const edges: GraphEdge3D[] = []

  for (const e of rawEdges) {
    const srcId = String(e.from ?? e.source)
    const tgtId = String(e.to ?? e.target)
    const si = idToIndex.get(srcId)
    const ti = idToIndex.get(tgtId)
    if (si != null && ti != null) {
      degrees[si]++
      degrees[ti]++
      edges.push({ sourceIndex: si, targetIndex: ti, type: e.type || e.label || '' })
    }
  }

  // Transform nodes
  const maxDegree = Math.max(1, ...degrees)
  const nodes: GraphNode3D[] = rawNodes.map((n, i) => {
    const weight = n.weight ?? 1
    const degree = degrees[i]
    const radius = Math.min(
      NODE_SIZE.max,
      NODE_SIZE.min + weight * NODE_SIZE.weightScale + (degree / maxDegree) * NODE_SIZE.degreeScale * 10,
    )
    return {
      id: String(n.id),
      label: n.label || n.text || String(n.id),
      type: n.type?.toLowerCase() || 'default',
      weight,
      confidence: n.confidence ?? n.salience ?? 0.5,
      degree,
      radius,
      colorIndex: getColorIndex(n.type),
    }
  })

  return { nodes, edges }
}

export function useTransformedGraph(raw: ApiGraphResponse | null): TransformedGraph | null {
  return useMemo(() => transformGraphData(raw), [raw])
}

export function getTopNodes(graph: TransformedGraph, count: number): TransformedGraph {
  const scored = graph.nodes.map((n, i) => ({ node: n, index: i, score: n.weight * (1 + n.degree) }))
  scored.sort((a, b) => b.score - a.score)
  const topSet = new Set(scored.slice(0, count).map(s => s.index))

  const oldToNew = new Map<number, number>()
  const nodes: GraphNode3D[] = []
  for (const s of scored.slice(0, count)) {
    oldToNew.set(s.index, nodes.length)
    nodes.push(s.node)
  }

  const edges: GraphEdge3D[] = []
  for (const e of graph.edges) {
    if (topSet.has(e.sourceIndex) && topSet.has(e.targetIndex)) {
      edges.push({
        sourceIndex: oldToNew.get(e.sourceIndex)!,
        targetIndex: oldToNew.get(e.targetIndex)!,
        type: e.type,
      })
    }
  }

  return { nodes, edges }
}
