import { readFile, open as openFile } from "fs/promises";
import { readSync, fstatSync } from "fs";
import type {
  MainModule,
  DSPContext,
  FfmpegAudioBuffer,
  Timeline,
  Track,
  AudioClip,
} from "../../artifacts/dsp-engine.js";
import createDspEngine from "../../artifacts/dsp-engine.js";

import getRenderableState from "./getRenderableState.js";
import autoCrossfadeClips from "@suno/studiokit/autoCrossfadeClips.js";
import { getEffectiveMarkers } from "@suno/studiokit/warpUtils.js";
import { getWarpEnabledAndPopulated } from "@suno/studiokit/warpUtils.js";
import { StudioClip } from "@suno/studiokit/projectState/fixClip.js";
import { StudioProjectState } from "@suno/studiokit/projectState/fixStudioProjectState.js";
import { DEFAULT_TRACK_EQ, WarpMarker } from "./types.js";

export type PCMChunk = { interleaved: Float32Array; frameCount: number };

export async function initDspModule(
  sampleRate: number = 48_000,
  enableLimiter: boolean = true
): Promise<{
  dspModule: MainModule;
  dspContext: DSPContext;
}> {
  try {
    const dspModule = await (
      createDspEngine as unknown as () => Promise<MainModule>
    )();
    const dspContext = await dspModule.DSPContext.createOffline(
      "bounce",
      0.5,
      2,
      sampleRate,
      enableLimiter
    );

    return { dspModule, dspContext };
  } catch (error) {
    console.error("Failed to initialize DSP module:", error);
    throw error;
  }
}

export function createFfmpegDiskBufferFetcher(
  dspModule: MainModule,
  clipIdToPath: (clipId: string) => string
) {
  const cache = new Map<string, Promise<FfmpegAudioBuffer | null>>();

  const fetcher = async (clipId: string): Promise<FfmpegAudioBuffer | null> => {
    if (cache.has(clipId)) {
      const cachedBuffer = await cache.get(clipId);
      if (cachedBuffer) {
        return cachedBuffer;
      }
    }

    const promise = (async () => {
      const path = clipIdToPath(clipId);
      if (!path) {
        console.error(`Media not found for clip ID: ${clipId}`);
        return null;
      }
      const handle = await openFile(path, "r");
      try {
        const size = fstatSync(handle.fd).size;
        return dspModule.FfmpegAudioBuffer.createFromReadCallback(
          size,
          (buf: Uint8Array, offset: number) => {
            const bytesRed = readSync(handle.fd, buf, 0, buf.length, offset);
            if (bytesRed !== buf.length) {
              // this should never happen
              console.error(
                `Short read from ${path} at offset ${offset}: got ${bytesRed} bytes, expected ${buf.length}`
              );
            }
          },
          () => {
            (async () => {
              // be sure to retain handle until it's closed
              await handle.close();
            })();
          }
        );
      } catch (error) {
        await handle.close();
        throw error;
      }
    })();

    cache.set(clipId, promise);

    return await promise;
  };

  fetcher.clearCache = async () => {
    await Promise.all(
      cache.values().map(async (promise) => {
        const buffer = await promise;
        if (buffer) {
          buffer.delete();
        }
      })
    );
    cache.clear();
  };

  return fetcher;
}

export function getRenderLengthSeconds(
  state: StudioProjectState,
  dspModule: MainModule,
  startBeats: number,
  endBeats: number
): number {
  const renderableState = getRenderableState(state);
  const moduleTiming = dspModule.TimelineTempoMap.fromArray(
    renderableState.timing.bps,
    renderableState.timing.bpsAutomation
  );
  try {
    const startSeconds = moduleTiming.transformBeatsToSeconds(startBeats);
    const endSeconds = moduleTiming.transformBeatsToSeconds(endBeats);
    if (startSeconds === undefined || endSeconds === undefined) {
      throw new Error("Failed to transform beats to seconds");
    }
    return endSeconds - startSeconds;
  } finally {
    moduleTiming.delete();
  }
}

export async function getStateFileRenderLengthSeconds(
  stateFilePath: string,
  startBeats: number,
  endBeats: number
): Promise<number> {
  const state = JSON.parse(
    await readFile(stateFilePath, "utf-8")
  ) as StudioProjectState;
  const { dspModule, dspContext } = await initDspModule(48_000, false);
  try {
    return getRenderLengthSeconds(state, dspModule, startBeats, endBeats);
  } finally {
    dspContext.delete();
  }
}

export async function createTimelineFromState(
  state: StudioProjectState,
  dspModule: MainModule,
  dspContext: DSPContext,
  getFfmpegBuffer: (url: string) => Promise<FfmpegAudioBuffer | null>,
  onlyTrackId?: string
): Promise<Timeline> {
  const renderableState = getRenderableState(state);

  const moduleTiming = dspModule.TimelineTempoMap.fromArray(
    renderableState.timing.bps,
    renderableState.timing.bpsAutomation
  );

  const tracks: Track[] = [];

  let songStartBeats = Infinity;
  let songEndBeats = -Infinity;

  // Preload all clips
  const preloadedClips = new Map<string, FfmpegAudioBuffer | null>(
    await Promise.all(
      Object.entries(renderableState.tracksById).flatMap(([trackId, track]) => {
        if (!track) return [];

        if (onlyTrackId && trackId !== onlyTrackId) return [];

        return track.clips
          .filter((clip) => (clip as any).resolved !== false)
          .filter((clip) => clip.clipId !== null)
          .map(
            async (clip) =>
              [clip.clipId, await getFfmpegBuffer(clip.clipId)] satisfies [
                string,
                FfmpegAudioBuffer | null
              ]
          );
      })
    )
  );

  for (const [trackId, track] of Object.entries(renderableState.tracksById)) {
    if (!track) continue;

    const includeTrackInTimeline =
      onlyTrackId === undefined || trackId === onlyTrackId;

    const trackClips: AudioClip[] = [];

    const crossfadeClips = autoCrossfadeClips<StudioClip>(track.clips);

    for (const clip of crossfadeClips) {
      if ((clip as any).resolved === false) continue;
      if (!clip.clipId) continue;

      songStartBeats = Math.min(songStartBeats, clip.startBeats);
      songEndBeats = Math.max(songEndBeats, clip.endBeats);

      if (!includeTrackInTimeline) continue;

      const transformedMarkers: WarpMarker[] = Object.entries(
        getEffectiveMarkers(clip.warp)
      ).map(([seconds, beats]) => ({
        timeInUnderlyingBuffer: Number(seconds),
        timeInOutput: Number(beats),
      }));

      const ffmpegBuffer = preloadedClips.get(clip.clipId);
      if (!ffmpegBuffer) {
        console.error(`Failed to load audio for clip ID: ${clip.clipId}`);
        continue;
      }

      const warpMap = dspModule.WarpMap.fromArray(transformedMarkers);

      const audioClip = dspContext.createAudioClip(
        ffmpegBuffer,
        warpMap,
        clip.amplitude ?? 1, // gain
        clip.startBeats,
        clip.endBeats,
        clip.loop.startBeats,
        999999, // clip.loop.endBeats,
        clip.readStartBeats,
        clip.fadeInBeats ?? 0, // fadeInBeats
        1, // fadeInExponent
        clip.fadeOutBeats ?? 0, // fadeOutBeats
        1, // fadeOutExponent
        clip.transposition ?? 0, // transposition
        1, // clip BPS
        getWarpEnabledAndPopulated(clip.warp),
        ""
      );

      warpMap.delete();

      trackClips.push(audioClip);
    }

    if (!includeTrackInTimeline) continue;

    // Create and configure filter chain (6 filters total)
    const filterChain = dspContext.createFilterChain(6);

    // Set up EQ from track state using keys
    const trackEQ = track.eq || DEFAULT_TRACK_EQ;

    // Map EQ band types to DSP engine filter modes
    const FILTER_MODE_MAP = {
      highpass: dspModule.BiQuadFilterMode.HIGHPASS,
      lowshelf: dspModule.BiQuadFilterMode.LOW_SHELF,
      peaking: dspModule.BiQuadFilterMode.EQ_BAND,
      notch: dspModule.BiQuadFilterMode.NOTCH,
      highshelf: dspModule.BiQuadFilterMode.HIGH_SHELF,
      lowpass: dspModule.BiQuadFilterMode.LOWPASS,
    } as const;

    // Configure all 6 EQ bands
    const bandKeys: Array<
      "band1" | "band2" | "band3" | "band4" | "band5" | "band6"
    > = ["band1", "band2", "band3", "band4", "band5", "band6"];

    bandKeys.forEach((key, index) => {
      const band = trackEQ[key];

      // for backwards compatibility. should be safe to remove after rollout.
      if (!band || !band.type) return;

      const mode =
        FILTER_MODE_MAP[band.type as keyof typeof FILTER_MODE_MAP] ??
        dspModule.BiQuadFilterMode.EQ_BAND;

      // Determine if this band should be bypassed
      // Bypass if: overall EQ disabled OR band disabled OR (for gain-based bands, gain is zero)
      const isGainBased =
        band.type === "peaking" ||
        band.type === "lowshelf" ||
        band.type === "highshelf";
      const shouldBypass =
        !trackEQ.enabled || !band.enabled || (isGainBased && band.gain === 0);

      filterChain.setMode(
        index,
        shouldBypass ? dspModule.BiQuadFilterMode.BYPASS : mode
      );
      filterChain.setFrequency(index, band.frequency);
      filterChain.setQ(index, band.q);
      filterChain.setGain(index, band.gain);
    });

    const moduleTrack = dspContext.createTrack(
      track.amplitude,
      track.balance,
      !!renderableState.effectivelyMutedTracks[trackId],
      null,
      filterChain,
      trackClips
    );

    filterChain.delete();
    trackClips.forEach((clip) => clip.delete());

    tracks.push(moduleTrack);
  }

  const timeline = dspContext.createTimeline(moduleTiming, tracks);

  tracks.forEach((track) => track.delete());
  moduleTiming.delete();

  timeline.fadeInStartBeats = songStartBeats;
  timeline.fadeInLengthBeats = state.songFadeInBeats ?? 0;
  timeline.fadeInExponent = 2;
  timeline.fadeOutEndBeats = songEndBeats;
  timeline.fadeOutLengthBeats = state.songFadeOutBeats ?? 0;
  timeline.fadeOutExponent = 2;
  timeline.masterGain = state.amplitude ?? 1;

  return timeline;
}

export async function* streamPCMChunks(
  stateFilePath: string,
  startBeats: number = 0,
  endBeats: number = 100,
  sampleRate: number = 48_000,
  bufferSize: number = 4096,
  enableLimiter: boolean = true,
  clipIdToPath: (clipId: string) => string,
  onlyTrackId?: string
): AsyncGenerator<PCMChunk> {
  const { dspModule, dspContext } = await initDspModule(
    sampleRate,
    enableLimiter
  );
  const getFfmpegBuffer = createFfmpegDiskBufferFetcher(
    dspModule,
    clipIdToPath
  );
  const stateContent = await readFile(stateFilePath, "utf-8");
  const state = JSON.parse(stateContent) as StudioProjectState;
  const timeline = await createTimelineFromState(
    state,
    dspModule,
    dspContext,
    getFfmpegBuffer,
    onlyTrackId
  );
  dspContext.swapLiveTimeline(timeline);
  timeline.delete();

  const bouncer = dspContext.createBouncer(startBeats, endBeats, bufferSize);
  if (!bouncer) throw new Error("Failed to create bouncer");

  try {
    while (true) {
      const buf = bouncer.bounceNext();
      if (!buf) break;
      try {
        const outputChannelArray = Array.from(
          { length: 2 },
          () => new Float32Array(buf.frameCount)
        );
        buf.setInto(outputChannelArray);
        const interleavedArray = new Float32Array(buf.frameCount * 2);
        for (let i = 0; i < buf.frameCount; i++) {
          interleavedArray[i * 2] = outputChannelArray[0][i];
          interleavedArray[i * 2 + 1] = outputChannelArray[1][i];
        }
        yield { interleaved: interleavedArray, frameCount: buf.frameCount };
      } finally {
        buf.delete();
      }
    }
  } finally {
    bouncer.delete();
    await getFfmpegBuffer.clearCache();
    dspContext.delete();
  }
}