import type {
  AnswerGroundingCheck,
  AssistantRequirement,
  RequirementCoverageReport,
  UnifiedRetrievalResult
} from "../types/assistant";
import type { NormalizedPayload, RouteHintSummary } from "../types/normalizer";

export interface AssistantRequirementExtractionResult {
  requirements: AssistantRequirement[];
  byFragment: Map<string, string[]>;
}

export interface AssistantCoverageEvaluationResult {
  requirements: AssistantRequirement[];
  coverage: RequirementCoverageReport;
}

export interface AssistantCoverageGroundingPipelineInput {
  routeSummary: RouteHintSummary | null;
  normalized: NormalizedPayload | null | undefined;
  userMessage: string;
  retrievalResults: UnifiedRetrievalResult[];
  requirementExtraction?: AssistantRequirementExtractionResult;
  extractRequirements: (
    routeSummary: RouteHintSummary | null,
    normalized: NormalizedPayload | null | undefined,
    userMessage: string
  ) => AssistantRequirementExtractionResult;
  evaluateCoverage: (
    requirements: AssistantRequirement[],
    retrievalResults: UnifiedRetrievalResult[]
  ) => AssistantCoverageEvaluationResult;
  checkGrounding: (
    userMessage: string,
    requirements: AssistantRequirement[],
    coverage: RequirementCoverageReport,
    retrievalResults: UnifiedRetrievalResult[]
  ) => AnswerGroundingCheck;
}

export interface AssistantCoverageGroundingPipelineOutput {
  requirementExtraction: AssistantRequirementExtractionResult;
  coverageEvaluation: AssistantCoverageEvaluationResult;
  groundingCheckBase: AnswerGroundingCheck;
}

export function runAssistantCoverageGroundingPipeline(
  input: AssistantCoverageGroundingPipelineInput
): AssistantCoverageGroundingPipelineOutput {
  const requirementExtraction =
    input.requirementExtraction ?? input.extractRequirements(input.routeSummary, input.normalized, input.userMessage);
  const coverageEvaluation = input.evaluateCoverage(requirementExtraction.requirements, input.retrievalResults);
  const groundingCheckBase = input.checkGrounding(
    input.userMessage,
    coverageEvaluation.requirements,
    coverageEvaluation.coverage,
    input.retrievalResults
  );

  return {
    requirementExtraction,
    coverageEvaluation,
    groundingCheckBase
  };
}