import type { Assessment, CanonicalAssessmentModel } from '../../domain/canonical-model';
import type { ImportedDataset } from '../../domain/dataset';
import type { UnifiedResultQuery, UnifiedResultQueryResult } from '../../domain/unified-result';
import { queryAllUnifiedResults } from '../../exports/unified-result-jsonl';
import { canonicalModelsToDataset } from './canonical-model-to-dataset';
import { appendPolicyObservationsToDataset } from './policy-observations-to-dataset';

export interface CanonicalProjectReader {
  listProjectAssessments(projectId: string): Promise<Assessment[]>;
  loadAssessmentModel(assessmentId: string): Promise<CanonicalAssessmentModel>;
  queryUnifiedResults?(query?: UnifiedResultQuery): Promise<UnifiedResultQueryResult>;
}

export interface CanonicalProjectWorkspace {
  models: CanonicalAssessmentModel[];
  dataset: ImportedDataset | null;
  assessmentIds: string[];
}

export function canonicalModelForScanRuns(
  model: CanonicalAssessmentModel,
  scanRunIds: string[]
): CanonicalAssessmentModel | null {
  const selectedIds = new Set(scanRunIds);
  const scanRuns = model.scanRuns.filter(run => selectedIds.has(run.id));
  if (!scanRuns.length) return null;
  const importFiles = model.importFiles.filter(file => selectedIds.has(file.scanRunId));
  const fileIds = new Set(importFiles.map(file => file.id));
  const sourceTables = model.sourceTables.filter(table => fileIds.has(table.importFileId));
  const tableIds = new Set(sourceTables.map(table => table.id));
  const findingOccurrences = model.findingOccurrences.filter(occurrence => selectedIds.has(occurrence.scanRunId));
  const occurrenceIds = new Set(findingOccurrences.map(occurrence => occurrence.id));
  const definitionIds = new Set(findingOccurrences.map(occurrence => occurrence.definitionId));
  return {
    ...model,
    scanRuns,
    importFiles,
    sourceTables,
    findingDefinitions: model.findingDefinitions.filter(definition => definitionIds.has(definition.id)),
    findingOccurrences,
    evidence: model.evidence.filter(record => occurrenceIds.has(record.occurrenceId)),
    reviews: model.reviews.filter(review => review.target.kind === 'definition'
      ? definitionIds.has(review.target.definitionId)
      : occurrenceIds.has(review.target.occurrenceId)),
    diagnostics: model.diagnostics.filter(diagnostic =>
      fileIds.has(diagnostic.importFileId)
      && (diagnostic.sourceTableId === undefined || tableIds.has(diagnostic.sourceTableId)))
  };
}

export function canonicalModelForScanRun(model: CanonicalAssessmentModel, scanRunId: string): CanonicalAssessmentModel | null {
  return canonicalModelForScanRuns(model, [scanRunId]);
}
export async function loadCanonicalProjectWorkspace(
  reader: CanonicalProjectReader,
  projectId: string,
  scanRunIds?: string[]
): Promise<CanonicalProjectWorkspace> {
  const assessments = await reader.listProjectAssessments(projectId);
  const loadedModels = await Promise.all(assessments.map(assessment => reader.loadAssessmentModel(assessment.id)));
  const models = scanRunIds?.length
    ? loadedModels.map(model => canonicalModelForScanRuns(model, scanRunIds))
      .filter((model): model is CanonicalAssessmentModel => model !== null)
    : loadedModels;
  const assessmentIds = [...new Set(models.map(model => model.assessment.id))];
  const policyObservations = reader.queryUnifiedResults && assessmentIds.length
    ? await queryAllUnifiedResults(
      query => reader.queryUnifiedResults!(query),
      { projectId, assessmentIds, sourceKind: ['policy-observation'] }
    )
    : [];
  const scannerDataset = canonicalModelsToDataset(models);
  const projectName = models[0]?.project.name ?? assessments.find(item => assessmentIds.includes(item.id))?.name ?? projectId;
  return {
    models,
    assessmentIds,
    dataset: appendPolicyObservationsToDataset(scannerDataset, projectName, policyObservations)
  };
}