import { useQuery } from '@tanstack/react-query';
import { fetchCodeGraphSampled } from '@/api/code';
import { NODE_COLORS } from './graphConfig';

export interface GraphNode {
  id: string;
  label: string;
  size?: number;
  color?: string;
  x?: number;
  y?: number;
}

export interface GraphEdge {
  source: string;
  target: string;
  weight?: number;
  color?: string;
}

interface RawNode {
  name?: string;
  id?: string;
  type?: string;
  size?: number;
}

interface RawEdge {
  source: string;
  target: string;
  weight?: number;
}

interface RawGraphResponse {
  nodes?: RawNode[];
  edges?: RawEdge[];
}

export function useGraphData(limit?: number) {
  return useQuery({
    queryKey: ['code', 'graph', 'sampled', { limit }],
    queryFn: () => fetchCodeGraphSampled(limit),
    staleTime: 30_000,
    select(raw: unknown): { nodes: GraphNode[]; edges: GraphEdge[] } {
      const data = raw as RawGraphResponse;
      const rawNodes = data.nodes ?? [];
      const rawEdges = data.edges ?? [];

      const nodeIds = new Set(rawNodes.map((n) => n.id ?? n.name ?? ''));

      const nodes: GraphNode[] = rawNodes.map((n) => {
        const id = n.id ?? n.name ?? '';
        const nodeType = n.type ?? 'default';
        return {
          id,
          label: String(n.name ?? id ?? ''),
          size: n.size ?? 5,
          color: NODE_COLORS[nodeType] ?? NODE_COLORS['default'],
        };
      });

      const edges: GraphEdge[] = rawEdges
        .filter((e) => nodeIds.has(e.source) && nodeIds.has(e.target))
        .map((e) => ({
          source: e.source,
          target: e.target,
          weight: e.weight,
        }));

      return { nodes, edges };
    },
  });
}
