export type PlanMode = "full" | "melody";

export type OutputFormat = "flac" | "mp3" | "opus";

/**
 * How the symbolic plan (ABC score) that YuE2 conditions on is produced.
 * - "model"      YuE2GenerateABC writes the score from style + lyrics
 * - "reference"  SheetSage2AudioToABC transcribes an uploaded recording
 * - "manual"     the operator supplies or edits the score by hand
 * - "none"       no score at all, YuE2 falls back to off mode
 */
export type ScoreSource = "model" | "reference" | "manual" | "none";

export interface GenerationRequest {
  style: string;
  lyrics: string;
  mode: PlanMode;
  scoreSource: ScoreSource;
  abc: string;
  referenceFilename: string | null;
  /** Seconds into the reference to start transcribing. */
  referenceStart: number;
  /** Seconds of reference to transcribe, or null for all of it. */
  referenceDuration: number | null;
  maxDuration: number;
  seed: number;
  randomizeSeed: boolean;
  checkpoint: string;
  audioEncoder: string;
  tiledDecode: boolean;
  /** Load the community Real-Audio decoder adapter (MODEL branch). */
  realAudioMatch: boolean;
  /** How hard that adapter pulls, 0 to 1.5. */
  realAudioStrength: number;
  /** Load the instrumental chain-of-thought adapter (CLIP branch). */
  instrumentalArrangement: boolean;
  /** How hard that adapter pulls, 0 to 1.5. */
  instrumentalStrength: number;
  outputFormat: OutputFormat;
  advanced: AdvancedParams;
}

export interface AdvancedParams {
  temperature: number;
  topP: number;
  topK: number;
  repetitionPenalty: number;
  cfgScale: number;
  steps: number;
  samplerName: string;
  scheduler: string;
  abcTemperature: number;
  abcTopP: number;
  abcTopK: number;
  abcRepetitionPenalty: number;
  maxAbcTokens: number;
}

export type JobStage =
  | "queued"
  | "transcribing"
  | "planning"
  | "conditioning"
  | "rendering"
  | "decoding"
  | "done"
  | "failed"
  | "cancelled";

export interface JobAudio {
  filename: string;
  subfolder: string;
  type: string;
  url: string;
}

export interface Job {
  id: string;
  promptId: string | null;
  createdAt: number;
  finishedAt: number | null;
  stage: JobStage;
  /** 0 to 1 within the current stage. */
  stageProgress: number;
  /** False when the websocket is unreachable and stages come from polling. */
  liveProgress: boolean;
  queuePosition: number | null;
  request: GenerationRequest;
  seedUsed: number;
  abc: string | null;
  audio: JobAudio | null;
  durationSeconds: number | null;
  error: string | null;
  title: string;
}

export interface ComfyCapabilities {
  connected: boolean;
  url: string;
  mock: boolean;
  checkpoints: string[];
  audioEncoders: string[];
  loras: string[];
  missingNodes: string[];
  error: string | null;
  /** Tail of the engine log, so a remote failure explains itself. */
  engineLog: string | null;
  /** Identifies the running image, so a stale container is obvious. */
  buildStamp: string | null;
}

export const STAGE_LABEL: Record<JobStage, string> = {
  queued: "Waiting for the GPU",
  transcribing: "Reading the reference",
  planning: "Writing the score",
  conditioning: "Encoding style and lyrics",
  rendering: "Rendering audio",
  decoding: "Decoding to stereo",
  done: "Ready",
  failed: "Failed",
  cancelled: "Cancelled",
};