import { StateGraph, MemorySaver } from '@langchain/langgraph';
import { EvaluationAnnotation } from './state';
import { questionGeneratorNode } from './nodes/generator.node';
import { interviewerNode } from './nodes/interviewer.node';
import { graderNode } from './nodes/grader.node';
import { reportAnalyzerNode } from './nodes/analyzer.node';

/**
 * Conditional routing logic for the Grader node.
 */
export const routeAfterGrading = (state: typeof EvaluationAnnotation.State) => {
  const targetCount = state.questionCount || 5;
  const questionsLen = state.questions?.length || 0;
  const currentIndex = Math.max(0, state.currentQuestionIndex || 0);

  console.log('[Router] Evaluation Result:', {
    currentIndex,
    shouldFollowUp: state.shouldFollowUp,
    numQuestions: questionsLen,
    targetCount,
  });

  if (state.shouldFollowUp) {
    console.log('[Router] Routing to follow-up interviewer');
    return 'interviewer';
  }

  if (currentIndex < targetCount) {
    // If the next question isn't generated yet, go back to generator
    if (currentIndex >= questionsLen) {
      console.log('[Router] Index >= Questions, routing to generator');
      return 'generator';
    }
    // If it is generated, go to interviewer
    console.log('[Router] Index < Questions, routing to interviewer');
    return 'interviewer';
  }

  console.log('[Router] Assessment complete, routing to analyzer');
  return 'analyzer';
};

/**
 * Builds and compiles the Evaluation Graph.
 */
export const createEvaluationGraph = () => {
  const workflow = new StateGraph(EvaluationAnnotation)
    .addNode('generator', questionGeneratorNode)
    .addNode('interviewer', interviewerNode)
    .addNode('grader', graderNode)
    .addNode('analyzer', reportAnalyzerNode)

    // Flow definition
    .addEdge('__start__', 'generator')
    .addEdge('generator', 'interviewer')

    // After interviewer, the graph will naturally pause for user input
    // if we use it in a thread-safe way with interrupts or simple invocation.
    .addEdge('interviewer', 'grader')

    // After grading, decide where to go
    .addConditionalEdges('grader', routeAfterGrading, {
      interviewer: 'interviewer',
      generator: 'generator',
      analyzer: 'analyzer',
    })

    .addEdge('analyzer', '__end__');

  // Using MemorySaver for thread-based persistence
  const checkpointer = new MemorySaver();

  return workflow.compile({
    checkpointer,
    // We want the graph to stop after the interviewer presents the question
    interruptAfter: ['interviewer'],
  });
};