server.ts

import express from "express";
import path from "path";
import { createServer as createViteServer } from "vite";
import { GoogleGenAI, Type } from "@google/genai";
import dotenv from "dotenv";
import { createMidiFile, NoteEvent } from "./src/utils/midiWriter.js";

dotenv.config();

async function startServer() {
  const app = express();
  const PORT = process.env.PORT ? parseInt(process.env.PORT, 10) : 3000;

  // Set body boundaries to support larger base64 file payloads (e.g. up to 60MB)
  app.use(express.json({ limit: "60mb" }));
  app.use(express.urlencoded({ limit: "60mb", extended: true }));

  // API Route: Health check
  app.get("/api/health", (req, res) => {
    res.json({ status: "ok", message: "Stem to MIDI server is healthy." });
  });

  // API Route: Transcribe audio stem to MIDI parameters
  app.post("/api/transcribe", async (req, res) => {
    try {
      const {
        audioData, // base64 payload
        mimeType,
        stemType = "melodic / piano",
        blend = 0.6,
        margin = 2.0,
        onsetThreshold = 0.3,
        offsetThreshold = 0.3,
        frameThreshold = 0.1,
        bpm = 120,
        quantize = false,
        quantizeGrid = "1/16"
      } = req.body;

      if (!audioData) {
        return res.status(400).json({ error: "Audio data is required as a base64 string." });
      }

      const apiKey = process.env.GEMINI_API_KEY;
      if (!apiKey || apiKey === "MY_GEMINI_API_KEY" || apiKey.trim() === "") {
        return res.status(500).json({
          error: "GEMINI_API_KEY environment variable is not configured in this applet. Please provide your Gemini API Key in the Secrets / Settings panel of AI Studio."
        });
      }

      console.log(`[Transcriber] Starting transcription using Gemini for stem type: ${stemType}, BPM: ${bpm}`);

      // Initialize the GoogleGenAI client lazily to avoid crash on startup with correct User-Agent
      const ai = new GoogleGenAI({
        apiKey,
        httpOptions: {
          headers: {
            "User-Agent": "aistudio-build"
          }
        }
      });

      // Build a specialized, detailed prompt for Gemini
      const prompt = `
      You are an expert Audio-to-MIDI transcriber and Piano Roll extractor.
      Analyze the attached audio recording, which represents a "${stemType}" stem, separated with HPSS (Harmonic-Percussive separation blend: ${blend}, boundary margin: ${margin}).

      The audio is filtered to isolate specific frequencies:
      - Stem Type: ${stemType}
      - Target AMT (Automatic Music Transcription) onset threshold: ${onsetThreshold}
      - Target offset threshold: ${offsetThreshold}
      - Target frame salience threshold: ${frameThreshold}

      Your task is to transcribe all audible pitches/melody events from this recording.
      For stem type "bass", concentrate on sub-120Hz fundamental frequencies.
      For stem type "drums / percussive", focus on rhythmic impulses and map transients to standard general MIDI drum pitches (Kick drum = 36, Snare = 38, Closed Hihat = 42, Open Hihat = 46, Crash Cymbal = 49).
      For stem type "melodic / piano" or "vocal", transcribe the primary singing voice, leads, synth pad notes, harmonies, or keyboard chords.

      IMPORTANT LIMIT: Only transcribe clear, audible, prominent note events. Do not match background noise, echoes, or transient vibrations. Limit your transcription list to the most essential melodic, rhythmic, or chordal events up to a maximum of 200 notes. This keeps the performance fast, clean, and ensures highly usable piano rolls.

      Output your transcription as a JSON array of note events matching the schema precisely.
      `;

      // Call Gemini 3.5 flash with strict JSON schema to guarantee correct output structure
      const response = await ai.models.generateContent({
        model: "gemini-3.5-flash",
        contents: [
          {
            inlineData: {
              data: audioData.replace(/^data:audio\/[^;]+;base64,/, ""), // Strip the base64 prefix if present
              mimeType: mimeType || "audio/mp3"
            }
          },
          prompt
        ],
        config: {
          responseMimeType: "application/json",
          responseSchema: {
            type: Type.ARRAY,
            items: {
              type: Type.OBJECT,
              properties: {
                midi_note: {
                  type: Type.INTEGER,
                  description: "The MIDI note value (21 to 108 representing pitch, or 36/38/42/46/49 for drums)"
                },
                onset_time: {
                  type: Type.NUMBER,
                  description: "The timestamp in seconds when the note starts (starting from 0.0)"
                },
                offset_time: {
                  type: Type.NUMBER,
                  description: "The timestamp in seconds when the note stops. offset_time must be strictly greater than onset_time."
                },
                velocity: {
                  type: Type.INTEGER,
                  description: "The dynamic velocity velocity value (1 to 127)"
                }
              },
              required: ["midi_note", "onset_time", "offset_time", "velocity"]
            }
          }
        }
      });

      const responseText = response.text || "";
      console.log("[Transcriber] Raw response from Gemini received.");

      let notes: NoteEvent[] = [];
      try {
        notes = JSON.parse(responseText.trim());
      } catch (err: any) {
        console.error("[Transcriber] Failed to parse JSON from Gemini:", responseText);
        return res.status(502).json({
          error: "Failed to parse note data from Gemini model. The model did not output valid JSON. Please try again.",
          rawResponse: responseText
        });
      }

      if (!Array.isArray(notes)) {
        return res.status(502).json({
          error: "Gemini response parsed successfully but is not a JSON array. Please try again.",
          rawResponse: responseText
        });
      }

      // Filter and validate notes
      let validatedNotes: NoteEvent[] = notes
        .map(n => ({
          midi_note: Math.round(Number(n.midi_note)),
          onset_time: Math.max(0, Number(n.onset_time)),
          offset_time: Number(n.offset_time),
          velocity: Math.min(127, Math.max(1, Math.round(Number(n.velocity) || 80)))
        }))
        .filter(n => !isNaN(n.midi_note) && !isNaN(n.onset_time) && !isNaN(n.offset_time) && n.offset_time > n.onset_time);

      // Apply BPM Quantization if enabled
      if (quantize && validatedNotes.length > 0) {
        console.log(`[Transcriber] Applying BPM quantization grid: ${quantizeGrid} at ${bpm} BPM`);
        const beatDuration = 60.0 / bpm;
        let gridDivisor = 0.25; // default 1/16th note (4 subdivisions per beat)
        
        switch (quantizeGrid) {
          case "1/4": gridDivisor = 1.0; break;
          case "1/8": gridDivisor = 0.5; break;
          case "1/16": gridDivisor = 0.25; break;
          case "1/32": gridDivisor = 0.125; break;
        }

        const gridInterval = beatDuration * gridDivisor;

        validatedNotes = validatedNotes.map(n => {
          const quantizedOnset = Math.round(n.onset_time / gridInterval) * gridInterval;
          const quantizedOffset = Math.round(n.offset_time / gridInterval) * gridInterval;
          // Ensure min duration is one grid division
          const duration = Math.max(gridInterval, quantizedOffset - quantizedOnset);

          return {
            ...n,
            onset_time: quantizedOnset,
            offset_time: quantizedOnset + duration
          };
        });
      }

      // Sort notes chronologically by onset
      validatedNotes.sort((a, b) => a.onset_time - b.onset_time);

      // Generate the physical binary MIDI file
      const midiBytes = createMidiFile(validatedNotes, bpm);
      const midiBase64 = Buffer.from(midiBytes).toString("base64");

      console.log(`[Transcriber] Transcribed ${validatedNotes.length} notes successfully.`);

      return res.json({
        notes: validatedNotes,
        midiBase64: midiBase64,
        noteCount: validatedNotes.length,
        bpm: bpm
      });

    } catch (err: any) {
      console.error("[Transcriber Error]", err);
      return res.status(500).json({
        error: err.message || "An unexpected error occurred during audio transcription."
      });
    }
  });

  // API Route: Client sends note data, server serves as downloadable .mid file binary
  app.post("/api/download-midi", (req, res) => {
    try {
      const { notes, bpm = 120, fileName = "transcribed_stem.mid" } = req.body;
      if (!notes || !Array.isArray(notes)) {
        return res.status(400).json({ error: "A valid list of notes is required." });
      }

      const midiBytes = createMidiFile(notes, bpm);
      
      res.setHeader("Content-Type", "audio/midi");
      res.setHeader("Content-Disposition", `attachment; filename="${fileName.replace(/[^a-zA-Z0-9_\.-]/g, "")}"`);
      return res.send(Buffer.from(midiBytes));
    } catch (err: any) {
      console.error("[MIDI Download Error]", err);
      return res.status(500).json({ error: "Failed to generate MIDI file for download." });
    }
  });

  // Global Error Handler for Express and body parsing errors
  app.use((err: any, req: express.Request, res: express.Response, next: express.NextFunction) => {
    console.error("[Global Server Error Context]", err);
    res.status(err.status || 500).json({
      error: err.message || "An unexpected server-side error occurred."
    });
  });

  // Vite Integration for Spa or static assets
  if (process.env.NODE_ENV !== "production") {
    const vite = await createViteServer({
      server: { middlewareMode: true },
      appType: "spa"
    });
    app.use(vite.middlewares);
  } else {
    const distPath = path.join(process.cwd(), "dist");
    app.use(express.static(distPath));
    app.get("*", (req, res) => {
      res.sendFile(path.join(distPath, "index.html"));
    });
  }

  app.listen(PORT, "0.0.0.0", () => {
    console.log(`[Express Server] Server listening on http://0.0.0.0:${PORT} in ${process.env.NODE_ENV || "development"} mode.`);
  });
}

startServer();