import type { GenerationRequest } from "@/lib/types";

export type ComfyNode = {
  class_type: string;
  inputs: Record<string, unknown>;
  _meta?: { title: string };
};

export type ComfyPrompt = Record<string, ComfyNode>;

/**
 * The community Real-Audio decoder adapter. Its author reports v9 as the best
 * of the published ladder, so that is what this ships. It attaches to the
 * MODEL side only and was built against the bf16 layout, both of which the
 * graph below respects.
 *
 * The matching tokenizer head, which is what would actually read a recording
 * into real-audio tokens, has no ComfyUI node and runs only from the repo's
 * python scripts. So this adapter shifts how the renderer sounds; it does not
 * give the model a second channel into the reference.
 */
export const REAL_AUDIO_LORA = "nar_lora_joint_v9_comfyui.safetensors";
export const REAL_AUDIO_LORAS = [REAL_AUDIO_LORA];
export const REAL_AUDIO_BASE_CHECKPOINT = "yue2_3b_bf16.safetensors";

/**
 * The instrumental chain-of-thought adapter, trained on ~2.7k instrumental
 * tracks paired with chord-annotated ABC from SheetSage2. It steers the
 * autoregressive side, so it rides on CLIP, and its training assumed a full
 * score was attached: its author's note is `cot="full"` with a non-empty ABC
 * step, which is what INSTRUMENTAL_REQUIRES below encodes.
 *
 * Because every one of those training tracks was instrumental, it pulls away
 * from vocals. That is a feature when the goal is an arrangement and a
 * problem when there are lyrics to sing, so the UI says so rather than
 * letting it quietly mute a vocal take.
 */
export const INSTRUMENTAL_LORA = "ar_lora_inst_v3abc_comfyui.safetensors";
export const INSTRUMENTAL_LORAS = [INSTRUMENTAL_LORA];

/** Whether the graph can honour the instrumental adapter as trained. */
export function instrumentalRequirementsMet(req: {
  mode: GenerationRequest["mode"];
  scoreSource: GenerationRequest["scoreSource"];
}): boolean {
  return req.mode === "full" && req.scoreSource !== "none";
}

/** Node ids are stable so the websocket relay can map progress onto a stage. */
export const NODE = {
  checkpoint: "ckpt",
  audioEncoder: "audio_encoder",
  loadAudio: "load_audio",
  trim: "trim",
  transcribe: "transcribe",
  abcRelay: "abc_relay",
  loraNar: "lora_nar",
  loraInst: "lora_inst",
  planner: "planner",
  plannerRelay: "planner_relay",
  music: "music",
  latent: "latent",
  negative: "negative",
  sampler: "sampler",
  decode: "decode",
  save: "save",
} as const;

/**
 * Builds the API-format prompt for ComfyUI's /prompt endpoint.
 *
 * The wiring mirrors the official "YuE2: Text to Music" and "YuE2: Music Cover"
 * templates: CheckpointLoaderSimple feeds YuE2GenerateMusic, whose seconds
 * output sizes the empty latent, and KSampler runs the flow model at cfg 1
 * with dpm_2 / sgm_uniform. SheetSage2AudioToABC emits a list, so a PreviewAny
 * sits between it and the abc input, exactly as the template does. That relay
 * also surfaces the score in /history, which is how the app reads it back.
 */
export function buildPrompt(req: GenerationRequest, seed: number): ComfyPrompt {
  const a = req.advanced;
  const prompt: ComfyPrompt = {};

  prompt[NODE.checkpoint] = {
    class_type: "CheckpointLoaderSimple",
    inputs: { ckpt_name: req.checkpoint },
    _meta: { title: "YuE2 checkpoint" },
  };

  // The renderer adapter, when asked for. It changes the acoustic prior the
  // flow model samples from, which is why it sits on the MODEL branch alone
  // and leaves CLIP on the checkpoint. VAE is untouched either way.
  // The two adapters chain. Each loader has to be handed both branches, but
  // each advances only the one it actually modifies, so nothing downstream
  // reads a branch back out of a loader that was set to zero strength on it.
  // The repos are explicit that the other branch is untouched; this keeps the
  // graph saying the same thing. VAE never enters the chain.
  let modelRef: [string, number] = [NODE.checkpoint, 0];
  let clipRef: [string, number] = [NODE.checkpoint, 1];
  const narStrength = clampStrength(req.realAudioStrength);
  const instStrength = clampStrength(req.instrumentalStrength);

  if (req.realAudioMatch && narStrength > 0) {
    prompt[NODE.loraNar] = {
      class_type: "LoraLoader",
      inputs: {
        model: modelRef,
        clip: clipRef,
        lora_name: REAL_AUDIO_LORA,
        strength_model: narStrength,
        strength_clip: 0,
      },
      _meta: { title: "Real audio renderer" },
    };
    modelRef = [NODE.loraNar, 0];
  }

  if (req.instrumentalArrangement && instStrength > 0 && instrumentalRequirementsMet(req)) {
    prompt[NODE.loraInst] = {
      class_type: "LoraLoader",
      inputs: {
        model: modelRef,
        clip: clipRef,
        lora_name: INSTRUMENTAL_LORA,
        strength_model: 0,
        strength_clip: instStrength,
      },
      _meta: { title: "Instrumental arrangement" },
    };
    clipRef = [NODE.loraInst, 1];
  }

  let abcInput: unknown = req.abc;

  if (req.scoreSource === "reference" && req.referenceFilename) {
    prompt[NODE.audioEncoder] = {
      class_type: "AudioEncoderLoader",
      inputs: { audio_encoder_name: req.audioEncoder },
      _meta: { title: "SheetSage2" },
    };
    prompt[NODE.loadAudio] = {
      class_type: "LoadAudio",
      inputs: { audio: req.referenceFilename },
      _meta: { title: "Reference recording" },
    };
    // SheetSage2 quantizes what it hears onto a subbeat grid and gives up on
    // notes it cannot place, which a long or busy recording eventually
    // contains. Trimming first both avoids that and cuts transcription time.
    let audioSource: [string, number] = [NODE.loadAudio, 0];
    if (req.referenceStart > 0 || req.referenceDuration !== null) {
      prompt[NODE.trim] = {
        class_type: "TrimAudioDuration",
        inputs: {
          audio: [NODE.loadAudio, 0],
          start_index: req.referenceStart,
          // The node wants a number; a very large one means "to the end".
          duration: req.referenceDuration ?? 100000,
        },
        _meta: { title: "Trim reference" },
      };
      audioSource = [NODE.trim, 0];
    }

    prompt[NODE.transcribe] = {
      class_type: "SheetSage2AudioToABC",
      inputs: {
        audio_encoder: [NODE.audioEncoder, 0],
        audio: audioSource,
        mode: req.mode,
      },
      _meta: { title: "Transcribe reference" },
    };
    prompt[NODE.abcRelay] = {
      class_type: "PreviewAny",
      inputs: { source: [NODE.transcribe, 0] },
      _meta: { title: "Transcribed score" },
    };
    abcInput = [NODE.abcRelay, 0];
  } else if (req.scoreSource === "model") {
    prompt[NODE.planner] = {
      class_type: "YuE2GenerateABC",
      inputs: {
        clip: clipRef,
        style: req.style,
        lyrics: req.lyrics,
        seed,
        mode: req.mode,
        max_abc_tokens: a.maxAbcTokens,
        temperature: a.abcTemperature,
        top_p: a.abcTopP,
        top_k: a.abcTopK,
        repetition_penalty: a.abcRepetitionPenalty,
        penalty_window: 100,
      },
      _meta: { title: "Write the score" },
    };
    prompt[NODE.plannerRelay] = {
      class_type: "PreviewAny",
      inputs: { source: [NODE.planner, 0] },
      _meta: { title: "Generated score" },
    };
    abcInput = [NODE.plannerRelay, 0];
  }

  prompt[NODE.music] = {
    class_type: "YuE2GenerateMusic",
    inputs: {
      clip: clipRef,
      style: req.style,
      lyrics: req.lyrics,
      abc: abcInput,
      seed,
      mode: req.mode,
      max_duration: req.maxDuration,
      temperature: a.temperature,
      top_p: a.topP,
      top_k: a.topK,
      repetition_penalty: a.repetitionPenalty,
      cfg_scale: a.cfgScale,
    },
    _meta: { title: "Condition on style and lyrics" },
  };

  prompt[NODE.latent] = {
    class_type: "EmptyYuE2LatentAudio",
    inputs: { seconds: [NODE.music, 1], batch_size: 1 },
    _meta: { title: "Audio canvas" },
  };

  prompt[NODE.negative] = {
    class_type: "ConditioningZeroOut",
    inputs: { conditioning: [NODE.music, 0] },
    _meta: { title: "Null conditioning" },
  };

  prompt[NODE.sampler] = {
    class_type: "KSampler",
    inputs: {
      model: modelRef,
      positive: [NODE.music, 0],
      negative: [NODE.negative, 0],
      latent_image: [NODE.latent, 0],
      seed,
      steps: a.steps,
      cfg: 1,
      sampler_name: a.samplerName,
      scheduler: a.scheduler,
      denoise: 1,
    },
    _meta: { title: "Render" },
  };

  prompt[NODE.decode] = req.tiledDecode
    ? {
        class_type: "VAEDecodeAudioTiled",
        inputs: {
          samples: [NODE.sampler, 0],
          vae: [NODE.checkpoint, 2],
          tile_size: 1920,
          overlap: 128,
        },
        _meta: { title: "Decode (tiled)" },
      }
    : {
        class_type: "VAEDecodeAudio",
        inputs: { samples: [NODE.sampler, 0], vae: [NODE.checkpoint, 2] },
        _meta: { title: "Decode" },
      };

  // SaveAudioAdvanced takes quality only for the lossy formats.
  const format = req.outputFormat ?? "flac";
  const saveInputs: Record<string, unknown> = {
    audio: [NODE.decode, 0],
    filename_prefix: "audio/yue2-studio",
    format,
  };
  if (format === "mp3") saveInputs.quality = "320k";
  if (format === "opus") saveInputs.quality = "192k";

  prompt[NODE.save] = {
    class_type: "SaveAudioAdvanced",
    inputs: saveInputs,
    _meta: { title: "Save" },
  };

  return prompt;
}

function clampStrength(value: number): number {
  if (!Number.isFinite(value)) return 0;
  return Math.min(1.5, Math.max(0, value));
}

/** Every node class this app needs ComfyUI to have installed. */
export const REQUIRED_NODES = [
  "CheckpointLoaderSimple",
  "YuE2GenerateABC",
  "YuE2GenerateMusic",
  "EmptyYuE2LatentAudio",
  "SheetSage2AudioToABC",
  "AudioEncoderLoader",
  "LoadAudio",
  "TrimAudioDuration",
  "ConditioningZeroOut",
  "KSampler",
  "VAEDecodeAudio",
  "VAEDecodeAudioTiled",
  "SaveAudioAdvanced",
  "PreviewAny",
  "LoraLoader",
];