import type { TrainingConfigSchema } from '../../config/training-defaults.d.ts';
import type { Tensor } from '../../gpu/tensor.js';
import type { DynamicLossScaler } from './loss-scaling.js';
import type { TrainingBatch, TrainingOptimizer, ClipMetrics } from './trainer.js';
import type { DataLoader } from './dataloader.js';
import type { TrainingObjective } from './objectives/base.js';
import type { UlArtifactFinalizeResult, DistillArtifactFinalizeResult } from './artifacts.js';

export interface TrainingStepMetricsEntry {
  schemaVersion: number;
  step: number;
  epoch: number;
  batch: number;
  total_loss: number;
  step_time_ms: number;
  forward_ms?: number;
  backward_ms?: number;
  optimizer_ms?: number;
  effective_lr?: number | null;
  lr?: number | null;
  seed?: number;
  model_id?: string;
  runtime_profile?: string | null;
  kernel_path?: string | null;
  environment_metadata?: Record<string, unknown>;
  memory_stats?: Record<string, unknown> | null;
  build_provenance?: Record<string, unknown> | null;
  scheduler_index?: number | null;
  scheduler_phase?: string | null;
  gradient_norm_unclipped?: number;
  gradient_norm_clipped?: number;
  clipped_event_count?: number;
  total_param_count?: number;
  trainable_param_count?: number | null;
  trainable_groups?: string[];
  frozen_groups?: string[];
  nan_count?: number;
  inf_count?: number;
  saturation_count?: number;
  telemetry_mode?: 'step' | 'window' | 'epoch';
  telemetry_window_size?: number;
  telemetry_alerts?: string[];
  window_loss_avg?: number | null;
  window_step_time_ms_avg?: number | null;
  ul_stage?: string | null;
  distill_stage?: string | null;
  lambda?: number | null;
  progress_shard_index?: number | null;
  progress_shard_count?: number | null;
  progress_step_in_shard?: number | null;
  progress_steps_in_shard?: number | null;
  progress_global_step?: number | null;
  progress_global_steps?: number | null;
  progress_percent_complete?: number | null;
  progress_elapsed_ms?: number | null;
  progress_eta_ms?: number | null;
  progress_eta_iso?: string | null;
  loss_kd?: number | null;
  loss_triplet?: number | null;
  distill_temperature?: number | null;
  distill_alpha_kd?: number | null;
  distill_alpha_ce?: number | null;
  distill_loss_ce_aux?: number | null;
  distill_loss_total?: number | null;
  distill_triplet_margin?: number | null;
  distill_triplet_active_count?: number | null;
  distill_stage_a_step_count?: number | null;
  distill_stage_a_kd_mean?: number | null;
  objective?: string;
  loss_total?: number | null;
  coeff_ce?: number | null;
  coeff_prior?: number | null;
  coeff_decoder?: number | null;
  coeff_recon?: number | null;
  schedule_step_index?: number | null;
  latent_clean_mean?: number | null;
  latent_clean_std?: number | null;
  latent_noise_mean?: number | null;
  latent_noise_std?: number | null;
  latent_noisy_mean?: number | null;
  latent_noisy_std?: number | null;
  stage1_latent_count?: number | null;
  loss_prior?: number | null;
  loss_decoder?: number | null;
  loss_recon?: number | null;
  latent_bitrate_proxy?: number | null;
  [key: string]: unknown;
}

export interface TrainingRunnerCallbacks {
  onStep?: (entry: TrainingStepMetricsEntry) => Promise<void> | void;
  onEpoch?: (entry: { epoch: number; steps: number; loss: number }) => Promise<void> | void;
  onCheckpoint?: (entry: {
    key: string;
    defaultCheckpointKey: string | null;
    path: string | null;
    metadata: Record<string, unknown> | null;
    payload: unknown;
    step: number;
    epoch: number;
    batch: number;
  }) => Promise<void> | void;
}

export interface TrainingRunnerOptions extends TrainingRunnerCallbacks {
  optimizer?: TrainingOptimizer;
  crossEntropyLoss?: (
    logits: Tensor,
    targets: Tensor,
    config: TrainingConfigSchema,
    tape: unknown
  ) => Promise<Tensor>;
  clipGradients?: (
    grads: Map<Tensor, Tensor>,
    config: TrainingConfigSchema
  ) => Promise<ClipMetrics>;
  lossScaler?: DynamicLossScaler;
  trainingObjective?: TrainingObjective;
  resolveCheckpointKey?: (entry: {
    defaultCheckpointKey: string | null;
    step: number;
    epoch: number;
    batch: number;
  }) => Promise<string> | string;
}

export interface TrainingRunOptions {
  epochs?: number;
  batchSize?: number;
  shuffle?: boolean;
  maxSteps?: number | null;
  logEvery?: number;
  prepareBatch?: (batch: unknown) => Promise<TrainingBatch> | TrainingBatch;
  ulArtifactDir?: string | null;
  distillArtifactDir?: string | null;
  stageAArtifact?: string | null;
  stageAArtifactHash?: string | null;
  teacherModelId?: string | null;
  studentModelId?: string | null;
  distillDatasetId?: string | null;
  distillDatasetPath?: string | null;
  distillLanguagePair?: string | null;
  distillSourceLangs?: string[] | null;
  distillTargetLangs?: string[] | null;
  distillPairAllowlist?: string[] | null;
  strictPairContract?: boolean;
  distillShardIndex?: number | null;
  distillShardCount?: number | null;
  checkpointKey?: string | null;
  resumeFrom?: string | null;
  forceResume?: boolean;
  forceResumeReason?: string | null;
  forceResumeSource?: string | null;
  runtimeProfile?: string | null;
  seed?: number | null;
  kernelPathId?: string | null;
  tokenizerHash?: string | null;
  optimizerStepCount?: number | null;
  buildProvenance?: Record<string, unknown> | null;
  buildId?: string | null;
  buildCommitHash?: string | null;
  buildTimestamp?: string | number | Date | null;
  gpuAdapterInfo?: Record<string, unknown> | null;
  command?: string | null;
  surface?: string | null;
  checkpointOperator?: string | null;
  modelId?: string | null;
  modelUrl?: string | null;
  timestamp?: string | Date | null;
  persistStage1Latents?: boolean;
  checkpointEvery?: number | null;
}

export declare class TrainingRunner {
  constructor(config: TrainingConfigSchema, options?: TrainingRunnerOptions);
  lastArtifact: UlArtifactFinalizeResult | DistillArtifactFinalizeResult | null;
  lastCheckpoint: {
    key: string;
    defaultKey?: string | null;
    path?: string | null;
    metadata?: Record<string, unknown> | null;
    step: number;
    epoch: number;
    batch: number;
  } | null;
  resumeState: {
    step: number;
    epoch: number;
    batch: number;
    checkpointHash?: string | null;
    resumeAudits?: Array<Record<string, unknown>>;
    resumeAuditCount?: number;
    previousCheckpointHash?: string | null;
    checkpointKey?: string | null;
  } | null;
  run(
    model: {
      forward: (input: Tensor, tape: unknown) => Promise<Tensor>;
      loraParams?: () => Tensor[];
      paramGroups?: () => Record<string, Tensor[]>;
    },
    dataset: TrainingBatch[] | DataLoader<TrainingBatch> | unknown[],
    options?: TrainingRunOptions
  ): Promise<TrainingStepMetricsEntry[]>;
}

export declare function runTraining(
  model: {
    forward: (input: Tensor, tape: unknown) => Promise<Tensor>;
    loraParams?: () => Tensor[];
    paramGroups?: () => Record<string, Tensor[]>;
  },
  dataset: TrainingBatch[] | DataLoader<TrainingBatch> | unknown[],
  config: TrainingConfigSchema,
  options?: TrainingRunOptions & TrainingRunnerOptions
): Promise<TrainingStepMetricsEntry[]>;

export declare function createTrainingCheckpointPayload(
  model: {
    loraParams?: () => Tensor[];
    paramGroups?: () => Record<string, Tensor[]>;
  },
  optimizer: unknown,
  context: {
    step: number;
    epoch: number;
    batch: number;
    config: TrainingConfigSchema;
  }
): Promise<unknown>;

export declare function restoreTrainingCheckpointState(
  model: {
    loraParams?: () => Tensor[];
    paramGroups?: () => Record<string, Tensor[]>;
  },
  optimizer: unknown,
  checkpointRecord: unknown,
  config: TrainingConfigSchema
): Promise<{
  step: number;
  epoch: number;
  batch: number;
  checkpointHash: string | null;
  previousCheckpointHash: string | null;
  checkpointKey: string | null;
  resumeAudits: Array<Record<string, unknown>>;
  resumeAuditCount: number;
} | null>;