frontend/src/context/AgentGraphContext.tsx

import React, {
  createContext,
  useContext,
  useReducer,
  ReactNode,
  useMemo,
} from "react";
import { Trace, KnowledgeGraph, PipelineStage } from "@/types";
import { TemporalGraphData } from "@/types/temporal";

interface AgentGraphState {
  selectedTrace: Trace | null;
  selectedKnowledgeGraph: KnowledgeGraph | null;
  selectedTemporalData: TemporalGraphData | null;
  traces: Trace[];
  knowledgeGraphs: KnowledgeGraph[];
  pipelineStages: PipelineStage[];
  isLoading: boolean;
  error: string | null;
  sidebarCollapsed: boolean;
  graphComparison: {
    selectedGraphs: KnowledgeGraph[];
    comparisonResults: any | null;
    isComparing: boolean;
    availableGraphs: KnowledgeGraph[];
    similarityMatrix: any | null;
  };
  activeView:
    | "welcome"
    | "pipeline"
    | "trace"
    | "results"
    | "traces"
    | "trace-kg"
    | "kg-visualizer"
    | "advanced-processing"
    | "trace-editor"
    | "temporal-visualizer"
    | "graph-comparison"
    | "example-traces"
    | "connections"
    | "upload";
}

interface AgentGraphContextType {
  state: AgentGraphState;
  actions: {
    setSelectedTrace: (trace: Trace | null) => void;
    setSelectedKnowledgeGraph: (kg: KnowledgeGraph | null) => void;
    setSelectedTemporalData: (data: TemporalGraphData | null) => void;
    setTraces: (traces: Trace[]) => void;
    setKnowledgeGraphs: (kgs: KnowledgeGraph[]) => void;
    setPipelineStages: (stages: PipelineStage[]) => void;
    setLoading: (loading: boolean) => void;
    setError: (error: string | null) => void;
    setSidebarCollapsed: (collapsed: boolean) => void;
    setActiveView: (view: AgentGraphState["activeView"]) => void;
    updatePipelineStage: (
      stageId: number,
      updates: Partial<PipelineStage>
    ) => void;
    updateTrace: (traceId: string, updates: Partial<Trace>) => void;
    updateKnowledgeGraph: (
      kgId: string,
      updates: Partial<KnowledgeGraph>
    ) => void;
  };
}

type AgentGraphAction =
  | { type: "SET_SELECTED_TRACE"; payload: Trace | null }
  | { type: "SET_SELECTED_KG"; payload: KnowledgeGraph | null }
  | { type: "SET_SELECTED_TEMPORAL_DATA"; payload: TemporalGraphData | null }
  | { type: "SET_TRACES"; payload: Trace[] }
  | { type: "SET_KNOWLEDGE_GRAPHS"; payload: KnowledgeGraph[] }
  | { type: "SET_PIPELINE_STAGES"; payload: PipelineStage[] }
  | { type: "SET_LOADING"; payload: boolean }
  | { type: "SET_ERROR"; payload: string | null }
  | { type: "SET_SIDEBAR_COLLAPSED"; payload: boolean }
  | { type: "SET_ACTIVE_VIEW"; payload: AgentGraphState["activeView"] }
  | {
      type: "UPDATE_PIPELINE_STAGE";
      payload: { stageId: number; updates: Partial<PipelineStage> };
    }
  | {
      type: "UPDATE_TRACE";
      payload: { traceId: string; updates: Partial<Trace> };
    }
  | {
      type: "UPDATE_KNOWLEDGE_GRAPH";
      payload: { kgId: string; updates: Partial<KnowledgeGraph> };
    };

const initialState: AgentGraphState = {
  selectedTrace: null,
  selectedKnowledgeGraph: null,
  selectedTemporalData: null,
  traces: [],
  knowledgeGraphs: [],
  pipelineStages: [
    {
      id: 1,
      name: "Prompt Reconstruction",
      status: "pending",
      progress: 0,
      icon: "Code",
      description: "Reconstruct prompts from trace data",
    },
    {
      id: 2,
      name: "Perturbation Testing",
      status: "pending",
      progress: 0,
      icon: "Shield",
      description: "Test knowledge graph robustness",
    },
    {
      id: 3,
      name: "Causal Analysis",
      status: "pending",
      progress: 0,
      icon: "GitBranch",
      description: "Analyze causal relationships",
    },
  ],
  isLoading: false,
  error: null,
  sidebarCollapsed: false,
  activeView: "welcome",
  graphComparison: {
    selectedGraphs: [],
    comparisonResults: null,
    isComparing: false,
    availableGraphs: [],
    similarityMatrix: null,
  },
};

function agentGraphReducer(
  state: AgentGraphState,
  action: AgentGraphAction
): AgentGraphState {
  switch (action.type) {
    case "SET_SELECTED_TRACE":
      return {
        ...state,
        selectedTrace: action.payload,
        activeView: action.payload ? "pipeline" : "welcome",
      };
    case "SET_SELECTED_KG":
      return { ...state, selectedKnowledgeGraph: action.payload };
    case "SET_SELECTED_TEMPORAL_DATA":
      return { ...state, selectedTemporalData: action.payload };
    case "SET_TRACES":
      return { ...state, traces: action.payload };
    case "SET_KNOWLEDGE_GRAPHS":
      return { ...state, knowledgeGraphs: action.payload };
    case "SET_PIPELINE_STAGES":
      return { ...state, pipelineStages: action.payload };
    case "SET_LOADING":
      return { ...state, isLoading: action.payload };
    case "SET_ERROR":
      return { ...state, error: action.payload };
    case "SET_SIDEBAR_COLLAPSED":
      return { ...state, sidebarCollapsed: action.payload };
    case "SET_ACTIVE_VIEW":
      return { ...state, activeView: action.payload };
    case "UPDATE_PIPELINE_STAGE":
      return {
        ...state,
        pipelineStages: state.pipelineStages.map((stage) =>
          stage.id === action.payload.stageId
            ? { ...stage, ...action.payload.updates }
            : stage
        ),
      };
    case "UPDATE_TRACE":
      return {
        ...state,
        traces: state.traces.map((trace) =>
          trace.trace_id === action.payload.traceId
            ? { ...trace, ...action.payload.updates }
            : trace
        ),
        selectedTrace:
          state.selectedTrace?.trace_id === action.payload.traceId
            ? { ...state.selectedTrace, ...action.payload.updates }
            : state.selectedTrace,
      };
    case "UPDATE_KNOWLEDGE_GRAPH":
      return {
        ...state,
        knowledgeGraphs: state.knowledgeGraphs.map((kg) =>
          kg.id === action.payload.kgId
            ? { ...kg, ...action.payload.updates }
            : kg
        ),
        selectedKnowledgeGraph:
          state.selectedKnowledgeGraph?.id === action.payload.kgId
            ? { ...state.selectedKnowledgeGraph, ...action.payload.updates }
            : state.selectedKnowledgeGraph,
      };
    default:
      return state;
  }
}

const AgentGraphContext = createContext<AgentGraphContextType | undefined>(
  undefined
);

export function AgentGraphProvider({ children }: { children: ReactNode }) {
  const [state, dispatch] = useReducer(agentGraphReducer, initialState);

  const actions = useMemo(
    () => ({
      setSelectedTrace: (trace: Trace | null) =>
        dispatch({ type: "SET_SELECTED_TRACE", payload: trace }),
      setSelectedKnowledgeGraph: (kg: KnowledgeGraph | null) =>
        dispatch({ type: "SET_SELECTED_KG", payload: kg }),
      setSelectedTemporalData: (data: TemporalGraphData | null) =>
        dispatch({ type: "SET_SELECTED_TEMPORAL_DATA", payload: data }),
      setTraces: (traces: Trace[]) =>
        dispatch({ type: "SET_TRACES", payload: traces }),
      setKnowledgeGraphs: (kgs: KnowledgeGraph[]) =>
        dispatch({ type: "SET_KNOWLEDGE_GRAPHS", payload: kgs }),
      setPipelineStages: (stages: PipelineStage[]) =>
        dispatch({ type: "SET_PIPELINE_STAGES", payload: stages }),
      setLoading: (loading: boolean) =>
        dispatch({ type: "SET_LOADING", payload: loading }),
      setError: (error: string | null) =>
        dispatch({ type: "SET_ERROR", payload: error }),
      setSidebarCollapsed: (collapsed: boolean) =>
        dispatch({ type: "SET_SIDEBAR_COLLAPSED", payload: collapsed }),
      setActiveView: (view: AgentGraphState["activeView"]) =>
        dispatch({ type: "SET_ACTIVE_VIEW", payload: view }),
      updatePipelineStage: (stageId: number, updates: Partial<PipelineStage>) =>
        dispatch({
          type: "UPDATE_PIPELINE_STAGE",
          payload: { stageId, updates },
        }),
      updateTrace: (traceId: string, updates: Partial<Trace>) =>
        dispatch({ type: "UPDATE_TRACE", payload: { traceId, updates } }),
      updateKnowledgeGraph: (kgId: string, updates: Partial<KnowledgeGraph>) =>
        dispatch({
          type: "UPDATE_KNOWLEDGE_GRAPH",
          payload: { kgId, updates },
        }),
    }),
    [dispatch]
  );

  return (
    <AgentGraphContext.Provider value={{ state, actions }}>
      {children}
    </AgentGraphContext.Provider>
  );
}

export function useAgentGraph() {
  const context = useContext(AgentGraphContext);
  if (context === undefined) {
    throw new Error("useAgentGraph must be used within an AgentGraphProvider");
  }
  return context;
}