import type { AssistantReplyType, AssistantRequirement, RequirementCoverageReport, UnifiedRetrievalResult } from "../types/assistant";
import type { RouteHintSummary } from "../types/normalizer";
import type { InvestigationStateWithProblemUnits } from "../types/stage2ProblemUnits";
import { updateInvestigationState } from "./investigationState";

export interface BuildAssistantInvestigationStateSnapshotInput {
  featureEnabled: boolean;
  previousState: InvestigationStateWithProblemUnits | null | undefined;
  timestamp: string;
  questionId: string;
  userMessage: string;
  routeSummary: RouteHintSummary | null;
  requirements: AssistantRequirement[];
  coverageReport: RequirementCoverageReport;
  retrievalResults: UnifiedRetrievalResult[];
  replyType: AssistantReplyType;
  followupApplied: boolean;
}

export function buildAssistantInvestigationStateSnapshot(
  input: BuildAssistantInvestigationStateSnapshotInput
): InvestigationStateWithProblemUnits | null {
  if (!input.featureEnabled || !input.previousState) {
    return null;
  }
  return updateInvestigationState({
    previous: input.previousState,
    timestamp: input.timestamp,
    questionId: input.questionId,
    userMessage: input.userMessage,
    routeSummary: input.routeSummary,
    requirements: input.requirements,
    coverageReport: input.coverageReport,
    retrievalResults: input.retrievalResults,
    replyType: input.replyType,
    followupApplied: input.followupApplied
  });
}

export interface PersistAssistantInvestigationStateSnapshotInput {
  featureEnabled: boolean;
  sessionId: string;
  snapshot: InvestigationStateWithProblemUnits | null;
  persist: (sessionId: string, snapshot: InvestigationStateWithProblemUnits) => void;
}

export function persistAssistantInvestigationStateSnapshot(
  input: PersistAssistantInvestigationStateSnapshotInput
): boolean {
  if (!input.featureEnabled || !input.snapshot) {
    return false;
  }
  input.persist(input.sessionId, input.snapshot);
  return true;
}