import type { TraceNode } from '@/common/interface/trace'
import agentBackground from '@/assets/images/agent.svg'
import guardBackground from '@/assets/images/guard.svg'
import toolBackground from '@/assets/images/tool.svg'
import llmBackground from '@/assets/images/llm.svg'
import eventBackground from '@/assets/images/event.svg'
import spanBackground from '@/assets/images/span.svg'

import emptyAgentBackground from '@/assets/images/agent-empty.svg'
import emptyGuardBackground from '@/assets/images/guard-empty.svg'
import emptyToolBackground from '@/assets/images/tool-empty.svg'
import emptyLlmBackground from '@/assets/images/llm-empty.svg'
import emptyEventBackground from '@/assets/images/event-empty.svg'
import emptySpanBackground from '@/assets/images/span-empty.svg'

export function classifyNodeGroup(node: TraceNode | null) {
  if (!node) return 'span'
  const name = (node.name || '').toLowerCase()
  const type = (node.type || '').toLowerCase()
  if (name.includes('python_llm_call') || name.includes('dku_agent_call')) return 'agent'
  if (name.includes('enforcement')) return 'guard'
  if (type === 'event') return 'event'
  if (node.attributes && node.attributes.completionQuery) return 'llm'
  if (name.includes('dku_llm_mesh_embedding_query')) return 'llm'
  if (name.includes('dku_managed_tool_call')) return 'tool'
  return 'span'
}

interface GroupConfig {
  [key: string]: {
    color: string
    background: string
    emptyBackground: string
    shape: string
    label: string
  }
}

export const GROUPS_CONFIG: GroupConfig = {
  agent: {
    color: '#fda5d5',
    background: `url(${agentBackground})`,
    emptyBackground: `url(${emptyAgentBackground})`,
    shape: 'diamond',
    label: 'Agent'
  },
  guard: {
    color: '#74d4ff',
    background: `url(${guardBackground})`,
    emptyBackground: `url(${emptyGuardBackground})`,
    shape: 'ellipse',
    label: 'Guard'
  },
  tool: {
    color: '#00d492',
    background: `url(${toolBackground})`,
    emptyBackground: `url(${emptyToolBackground})`,
    shape: 'ellipse',
    label: 'Tool'
  },
  llm: {
    color: '#51a2ff',
    background: `url(${llmBackground})`,
    emptyBackground: `url(${emptyLlmBackground})`,
    shape: 'ellipse',
    label: 'LLM'
  },
  event: {
    color: '#ff8904',
    background: `url(${eventBackground})`,
    emptyBackground: `url(${emptyEventBackground})`,
    shape: 'ellipse',
    label: 'Event'
  },
  span: {
    color: '#d6d3d1',
    background: `url(${spanBackground})`,
    emptyBackground: `url(${emptySpanBackground})`,
    shape: 'ellipse',
    label: 'Span'
  }
}

export function getGroupColor(group: string) {
  return GROUPS_CONFIG[group]?.color || '#bbb6b6'
}

export function getGroupBackground(group: string, empty: boolean) {
  if (empty) {
    return GROUPS_CONFIG[group]?.emptyBackground || null
  }
  return GROUPS_CONFIG[group]?.background || null
}

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

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