import { Annotation, MessagesAnnotation } from '@langchain/langgraph';
import { BaseMessage } from '@langchain/core/messages';

/**
 * State representing the evaluation session using LangGraph Annotation.
 */
export const EvaluationAnnotation = Annotation.Root({
  /**
   * The message history of the conversation.
   * Inherits from MessagesAnnotation to handle message merging.
   */
  ...MessagesAnnotation.spec,

  /**
   * Historical evaluation feedback from the grader.
   * Separated from main messages to keep conversation context clean.
   */
  feedbackHistory: Annotation<BaseMessage[]>({
    reducer: (prev, next) => [...(prev || []), ...(next || [])],
    default: () => [],
  }),

  /**
   * The database ID of the current assessment session.
   */
  assessmentSessionId: Annotation<string>(),

  /**
   * The knowledge base ID used as ground truth for this evaluation.
   */
  knowledgeBaseId: Annotation<string>(),

  /**
   * List of questions generated for this session.
   */
  questions: Annotation<any[]>({
    reducer: (prev, next) => next ?? prev,
    default: () => [],
  }),

  /**
   * Index of the current question being discussed.
   */
  currentQuestionIndex: Annotation<number>({
    reducer: (prev, next) => next ?? prev,
    default: () => 0,
  }),

  /**
   * Flag indicating if the Grader believes a follow-up question is needed for clarity.
   */
  shouldFollowUp: Annotation<boolean>({
    reducer: (prev, next) => next ?? prev,
    default: () => false,
  }),

  /**
   * Map of scores for each question.
   */
  scores: Annotation<Record<string, number>>({
    reducer: (prev, next) => ({ ...prev, ...next }),
    default: () => ({}),
  }),

  /**
   * Final report generated by the ReportAnalyzer.
   */
  report: Annotation<string | undefined>({
    reducer: (prev, next) => next ?? prev,
  }),

  /**
   * Context chunks retrieved from the knowledge base for grounding.
   */
  context: Annotation<string[] | undefined>({
    reducer: (prev, next) => next ?? prev,
  }),

  /**
   * Preferred language for the assessment (zh, en, ja).
   */
  language: Annotation<string>({
    reducer: (prev, next) => next ?? prev,
    default: () => 'zh',
  }),

  /**
   * Number of times we have followed up on the current question.
   */
  followUpCount: Annotation<number>({
    reducer: (prev, next) => next ?? prev,
    default: () => 0,
  }),

  /**
   * Number of questions to generate.
   */
  questionCount: Annotation<number | undefined>({
    reducer: (prev, next) => next ?? prev,
  }),

  /**
   * Desired difficulty distribution.
   */
  difficultyDistribution: Annotation<any | undefined>({
    reducer: (prev, next) => next ?? prev,
  }),

  /**
   * Desired question style.
   */
  style: Annotation<string | undefined>({
    reducer: (prev, next) => next ?? prev,
  }),

  /**
   * Target keywords for question generation.
   */
  keywords: Annotation<string[] | undefined>({
    reducer: (prev, next) => next ?? prev,
  }),
});

export type EvaluationState = typeof EvaluationAnnotation.State;
