import type { TraceNode } from '@/types/trace'

export type NodeGroupName = 'agent' | 'guard' | 'tool' | 'llm' | 'event' | 'span'

// theme-specific colors for each group (hex values for Cytoscape/SVG)
export const GROUP_THEME_COLORS: Record<
  NodeGroupName,
  { light: { bg: string; text: string }; dark: { bg: string; text: string } }
> = {
  agent: { light: { bg: '#eb78ff', text: '#fff' }, dark: { bg: '#eb78ff', text: '#fff' } },
  guard: { light: { bg: '#bc7aff', text: '#fff' }, dark: { bg: '#a080ff', text: '#fff' } },
  tool: { light: { bg: '#26d5a4', text: '#fff' }, dark: { bg: '#26d5a4', text: '#fff' } },
  llm: { light: { bg: '#2994ff', text: '#fff' }, dark: { bg: '#60a5fa', text: '#fff' } },
  event: { light: { bg: '#ffad42', text: '#fff' }, dark: { bg: '#fb923c', text: '#fff' } },
  span: { light: { bg: '#7e8895', text: '#fff' }, dark: { bg: '#94a3b8', text: '#fff' } },
}

// letter displayed in each node icon
export const GROUP_LETTERS: Record<NodeGroupName, string> = {
  agent: 'A',
  guard: 'G',
  tool: 'T',
  llm: 'L',
  event: 'E',
  span: 'S',
}

// generate SVG data URL for node icon
export function generateNodeSvgDataUrl(groupName: string, isDark: boolean): string {
  const name = groupName as NodeGroupName
  const colors = GROUP_THEME_COLORS[name] || GROUP_THEME_COLORS.span
  const letter = GROUP_LETTERS[name] || 'S'
  const theme = isDark ? colors.dark : colors.light

  // mix with white to desaturate colors
  const bgColor = `color-mix(in oklch, ${theme.bg} 85%, white)`
  const svg = `<svg xmlns="http://www.w3.org/2000/svg" width="100" height="100" viewBox="0 0 1024 1024">
    <rect width="100%" height="100%" fill="${bgColor}"/>
    <text x="50%" y="51.75%" fill="${theme.text}" font-size="170" font-family="Inter, Arial, sans-serif"
          text-anchor="middle" dominant-baseline="middle">${letter}</text>
  </svg>`

  return `data:image/svg+xml,${encodeURIComponent(svg)}`
}

const LLM_NODE_NAMES = [
  'dku_llm_mesh_call',
  'dku_llm_mesh_llm_call',
  'dku_llm_mesh_llm_call_streamed',
  'dku_llm_mesh_completion_query',
  'dku_llm_mesh_completion_query_streamed',
  'dku_llm_mesh_embedding_query',
  'dku_llm_mesh_image_generation_query',
  'dku_llm_mesh_reranking_query',
]

const AGENT_NODE_NAMES = [
  'dku_agent_call',
  'dku_agent_llm_call',
  'dku_agent_iteration',
  'dku_agent_tool_calls',
  'dku_tool_call',
]

export function classifyNodeGroup(node: TraceNode | null) {
  if (!node) return 'span'
  const name = (node.name || '').toLowerCase()
  const type = (node.type || '').toLowerCase()
  if (AGENT_NODE_NAMES.includes(name)) return 'agent'
  if (name.includes('enforcement')) return 'guard'
  if (type === 'event') return 'event'
  if (LLM_NODE_NAMES.includes(name)) return 'llm'
  if (name.includes('dku_managed_tool_call')) return 'tool'
  return 'span'
}

interface GroupConfig {
  shape: 'round-diamond' | 'ellipse' | 'round-rectangle'
  label: string
  icon: string
}

export const GROUPS_CONFIG: Record<NodeGroupName, GroupConfig> = {
  agent: {
    shape: 'round-diamond',
    label: 'Agent',
    icon: 'dku-icon-ai-agent',
  },
  guard: {
    shape: 'round-rectangle',
    label: 'Guard',
    icon: 'dku-icon-shield-check',
  },
  tool: {
    shape: 'round-diamond',
    label: 'Tool',
    icon: 'dku-icon-tool-wrench',
  },
  llm: {
    shape: 'ellipse',
    label: 'LLM',
    icon: 'dku-icon-llm',
  },
  event: {
    shape: 'round-rectangle',
    label: 'Event',
    icon: 'dku-icon-pulse',
  },
  span: {
    shape: 'ellipse',
    label: 'Span',
    icon: 'dku-icon-brackets-curly',
  },
}

export function getGroupColor(group: TraceNode | string, isDark: boolean = false) {
  const groupName = typeof group === 'string' ? group : classifyNodeGroup(group)
  const colors = GROUP_THEME_COLORS[groupName as NodeGroupName]
  if (!colors) return '#bbb6b6'
  return isDark ? colors.dark.bg : colors.light.bg
}

export function getGroupShape(group: string) {
  return GROUPS_CONFIG[group as NodeGroupName]?.shape || 'ellipse'
}

export function getGroupLabel(group: string) {
  return GROUPS_CONFIG[group as NodeGroupName]?.label || 'No Label'
}

export function getGroupIcon(groupName: string) {
  return GROUPS_CONFIG[groupName as NodeGroupName]?.icon || 'dku-icon-brackets-curly'
}

export function getGroupCytoscapeImage(groupName: string, isDark: boolean = false): string {
  return generateNodeSvgDataUrl(groupName, isDark)
}
