// ============================================================
// LivingLSAIs — Gemini Live Bridge Server  v1.0
// By ZEXIVA | Drop this in: Scripts/LLAIBridge/server.js
// Requires: node.exe in same folder
// Run:      node.exe server.js --parentPid <GTA_PID>
// ============================================================

const WebSocket = require("ws");
const https     = require("https");
const http      = require("http");

// ── Config ────────────────────────────────────────────────────
const BRIDGE_PORT   = 8765;          // C# connects here
const GEMINI_MODEL  = "gemini-3.1-flash-live-preview";
const GEMINI_WS_URL =
  `wss://generativelanguage.googleapis.com/ws/google.ai.generativelanguage.v1beta.GenerativeService.BidiGenerateContent`;

// ── Parent-PID watchdog (kills server when GTA exits) ─────────
const args      = process.argv.slice(2);
const pidIdx    = args.indexOf("--parentPid");
const parentPid = pidIdx >= 0 ? parseInt(args[pidIdx + 1]) : null;

if (parentPid) {
  setInterval(() => {
    try { process.kill(parentPid, 0); }
    catch (_) { console.log("[Bridge] GTA process gone — shutting down."); process.exit(0); }
  }, 3000);
}

// ── State per C# connection ───────────────────────────────────
// One NPC chat = one Gemini Live session.
// When C# sends { type:"startSession" } we open a new Gemini WS.
// When C# sends { type:"endSession"   } we close it.

let geminiWs       = null;   // active Gemini WebSocket
let csClient       = null;   // active C# WebSocket client
let apiKey         = null;   // set on first startSession
let sessionReady   = false;
let pendingMessage = null;   // 🔥 ZEXIVA FIX: queued message if C# was too fast
let turnTextAccum  = "";     // 🔥 TAG FIX: accumulates full turn text (with tags) before turnComplete fires
let pendingHistory = [];     // 🔥 MEMORY FIX: prior conversation turns to replay after setupComplete

// 🎬 SCENE SYSTEM — 2 separate Gemini Live sessions (one per NPC)
let sceneWsA       = null;   // NPC A Gemini WebSocket
let sceneWsB       = null;   // NPC B Gemini WebSocket
let sceneReadyA    = false;
let sceneReadyB    = false;
let sceneTextA     = "";     // Accumulated turn text for NPC A
let sceneTextB     = "";
let sceneAudioA    = [];     // PCM chunks for NPC A (buffered until turnComplete)
let sceneAudioB    = [];

// ── Local bridge WebSocket server ────────────────────────────
const wss = new WebSocket.Server({ port: BRIDGE_PORT });
console.log(`[Bridge] Listening on ws://127.0.0.1:${BRIDGE_PORT}`);

wss.on("connection", (ws) => {
  console.log("[Bridge] C# client connected");
  csClient = ws;

  ws.on("message", (data) => {
    try {
      const msg = JSON.parse(data.toString());
      handleCsMessage(msg, ws);
    } catch (e) {
      console.error("[Bridge] Bad JSON from C#:", e.message);
    }
  });

  ws.on("close", () => {
    console.log("[Bridge] C# client disconnected");
    csClient = null;
    closeGeminiSession();
  });

  ws.on("error", (e) => console.error("[Bridge] C# WS error:", e.message));
});

// ── Handle messages from C# ───────────────────────────────────
function handleCsMessage(msg, ws) {
  switch (msg.type) {

    // ── C# says: start a new NPC conversation ─────────────────
    case "startSession": {
      apiKey = msg.apiKey;
      if (!apiKey) {
        sendToCs({ type: "error", message: "No apiKey in startSession message" });
        return;
      }
      closeGeminiSession();            // close any previous session
      openGeminiSession(msg, ws);
      break;
    }

    // ── C# says: player typed / spoke ─────────────────────────
    case "playerMessage": {
      if (!geminiWs || geminiWs.readyState !== WebSocket.OPEN) {
        sendToCs({ type: "error", message: "No active Gemini session" });
        return;
      }
      if (!sessionReady) {
        // 🔥 ZEXIVA FIX: Queue karo, drop mat karo — silent voice bug fix
        console.warn("[Bridge] playerMessage queued — waiting for sessionReady");
        pendingMessage = { text: msg.text, contextUpdate: msg.contextUpdate || "" };
        return;
      }
      sendPlayerTurn(msg.text, msg.contextUpdate || "");
      break;
    }

    // ── C# says: end this NPC conversation ────────────────────
    case "endSession": {
      closeGeminiSession();
      break;
    }

    // 🎬 SCENE: Start 2 separate Gemini Live sessions (one per NPC)
    case "startSceneSession": {
      closeSceneSessions();
      apiKey = msg.apiKey;
      openSceneSession("A", msg.profileA, apiKey);
      openSceneSession("B", msg.profileB, apiKey);
      break;
    }

    // 🎬 SCENE: Send a turn for one NPC (A or B)
    case "sceneNpcTurn": {
      const ws = msg.npc === "A" ? sceneWsA : sceneWsB;
      const ready = msg.npc === "A" ? sceneReadyA : sceneReadyB;
      if (!ws || !ready) {
        sendToCs({ type: "sceneError", message: `Scene session NPC ${msg.npc} not ready` });
        return;
      }
      if (msg.npc === "A") { sceneTextA = ""; sceneAudioA = []; }
      else                  { sceneTextB = ""; sceneAudioB = []; }
      ws.send(JSON.stringify({
        clientContent: {
          turns: [{ role: "user", parts: [{ text: msg.text }] }],
          turn_complete: true
        }
      }));
      console.log(`[Scene] sceneNpcTurn NPC ${msg.npc}: "${msg.text.substring(0, 60)}"`);
      break;
    }

    // 🎬 SCENE: End both NPC sessions
    case "endSceneSession": {
      closeSceneSessions();
      break;
    }

    default:
      console.warn("[Bridge] Unknown message type:", msg.type);
  }
}

// ── Open a Gemini Live session ────────────────────────────────
function openGeminiSession(startMsg, ws) {
  sessionReady = false;
  const url    = `${GEMINI_WS_URL}?key=${apiKey}`;

  console.log(`[Bridge] Opening Gemini Live session for NPC: ${startMsg.npcName}`);

  geminiWs = new WebSocket(url);

  geminiWs.on("open", () => {
    console.log("[Bridge] Gemini WS open — sending setup");

    // ── BidiGenerateContent setup message ─────────────────────
    // This is the Gemini Live API "setup" frame.
    // systemInstruction  = teri full NpcProfile.ToSystemPrompt() string
    // voice              = Puck / Aoede / etc. (sent from C#)
    // output audioConfig = PCM 24kHz mono (raw bytes, no header)

    // History turns are handled dynamically in C# via system prompt injections to prevent Gemini schema crashes.
    console.log(`[Bridge] Initializing Gemini Live session for ${startMsg.npcName}`);

    const systemPromptWithAudioRule = `🚨 AUDIO MODALITY RULE (CRITICAL): You are absolutely BANNED from ever pronouncing, speaking, whispering, or saying the bracketed action tags (like [FIGHT], [FLEE], [DANCE], [HANDSUP], [KNEEL], [GIVEMONEY], [HANDSDOWN], [STOPSHOOTING], [STOPAIMING], [STANDUP], etc.) out loud in your audio output! The bracketed tags must ONLY exist in your text output. In your spoken voice, completely skip pronouncing the bracketed tags, as if they do not exist. Speak only the normal dialogue.

🚨 DESI LANGUAGE AUDIO RULE (CRITICAL): If speaking in Hinglish, Haryanvi, Bhojpuri, or Khari Boli, you MUST speak and pronounce the Romanized Hindi/dialect words naturally in your audio stream! Do NOT use Devanagari script. Do NOT use any '|' pipe separator. Just pronounce the Romanized script directly with natural spoken rhythm and cadence.

🚨 TAG POSITION LAW (ABSOLUTE — NO EXCEPTIONS): Action tags MUST ALWAYS appear at the VERY END of your text response — AFTER all spoken words and dialogue. NEVER place a tag at the START of a sentence (causes audio silence bug where voice is completely muted). NEVER place a tag in the middle of dialogue.
CORRECT: "You just made the worst mistake of your life, back off! [FIGHT]"
CORRECT: "Okay okay, please don't shoot! I'm complying! [HANDSUP]"
WRONG: "[FIGHT] You just made the worst mistake of your life!" — TAG AT START = AUDIO SILENCE GLITCH, voice muted.
WRONG: "Okay [HANDSUP] I'm complying!" — TAG IN MIDDLE = BROKEN.
This is ABSOLUTE. If your response has an action tag, it MUST be the very LAST thing after all dialogue. No exceptions. Ever.

🚨 STOP-ACTION PROTOCOL (MANDATORY): If you are currently in a hands-up state, you MUST append [HANDSDOWN] if told to put them down. If you are currently kneeling, you MUST append [STANDUP] to stand back up. If you are shooting, you MUST use [STOPSHOOTING] to stop. If you are aiming/warning, you MUST use [STOPAIMING] to holster. This ensures realistic transition animations and prevents buggy loops!

🚨 STRICT EXPLICIT PERMISSION LAW (STARTNPCTONPC & FOLLOW) — OBEY OR BE DELETED:
1. You are ABSOLUTELY BANNED from using [STARTNPCTONPC] or [FOLLOW] tags on your own! If you are proposing the action (e.g., 'Should I go talk to him?', 'Do you want me to follow you?'), you MUST output ONLY plain text. DO NOT use the tags yet.
2. You are ONLY allowed to append [STARTNPCTONPC] or [FOLLOW] if the player has EXPLICITLY commanded you, answered 'yes', 'do it', 'go', 'haan', 'chal', or 'baat kar' in their very last message. If the player did not explicitly give permission, these tags are strictly BANNED. Violation of this law triggers instant script termination.\n\n${startMsg.systemPrompt}`;

    const setup = {
      setup: {
        model: `models/${GEMINI_MODEL}`,
        generation_config: {
          response_modalities:  ["AUDIO"],
          speech_config: {
            voice_config: {
              prebuilt_voice_config: {
                voice_name: startMsg.voiceName || "Puck"
              }
            }
          }
        },
        input_audio_transcription: {},
        output_audio_transcription: {},
        system_instruction: {
          parts: [{ text: systemPromptWithAudioRule }]
        }
      }
    };

    geminiWs.send(JSON.stringify(setup));
  });

  geminiWs.on("message", (rawData) => {
    handleGeminiMessage(rawData);
  });

  geminiWs.on("error", (e) => {
    console.error("[Bridge] Gemini WS error:", e.message);
    sendToCs({ type: "error", message: "Gemini WS error: " + e.message });
  });

  const socket = geminiWs;
  socket.on("close", (code, reason) => {
    console.log(`[Bridge] Gemini WS closed. code=${code} reason=${reason}`);
    if (geminiWs === socket) {
      sessionReady = false;
      geminiWs     = null;
      sendToCs({ type: "sessionClosed" });
    }
  });
}

// ── Handle messages FROM Gemini ───────────────────────────────
function handleGeminiMessage(rawData) {
  let msg;
  try {
    msg = JSON.parse(rawData.toString());
  } catch (e) {
    console.error("[Bridge] Non-JSON from Gemini:", e.message);
    return;
  }

  // ── 1. Setup complete ── ─────────────────────────────────────
  if (msg.setupComplete !== undefined) {
    console.log("[Bridge] Gemini session ready!");
    sessionReady = true;
    sendToCs({ type: "sessionReady" });

    // 🔥 ZEXIVA FIX: Drain queued message if C# sent it before session was ready
    if (pendingMessage) {
      console.log("[Bridge] Draining queued message:", pendingMessage.text.substring(0, 50));
      sendPlayerTurn(pendingMessage.text, pendingMessage.contextUpdate);
      pendingMessage = null;
    }
    return;
  }

  // Extract output transcription (subtitles) - can be root or nested inside serverContent
  const transcriptionText = 
    (msg.outputTranscription && msg.outputTranscription.text) || 
    (msg.serverContent && msg.serverContent.outputTranscription && msg.serverContent.outputTranscription.text);

  // ── 2. Server content (audio + text + transcription) ─────────
  // 🔥 TAG FIX: Transcription text is the ONLY reliable source of tags.
  // We accumulate it into textAccum and send ONE "turnComplete" at the end.
  // We do NOT send separate "text" events from outputTranscription to avoid
  // duplicates — part.text (if present) covers subtitle streaming instead.
  if (msg.serverContent) {
    const sc = msg.serverContent;

    // 🔥 ACCUMULATE transcription text into our own buffer (tags included!)
    // outputTranscription is the authoritative text with tags — store it.
    if (transcriptionText) {
      const pipeIdx = transcriptionText.indexOf("|");
      const cleaned = pipeIdx !== -1
        ? transcriptionText.substring(0, pipeIdx).trim()
        : transcriptionText;
      // Only append if it's new content (not already in accum)
      if (!turnTextAccum.endsWith(cleaned)) {
        turnTextAccum += (turnTextAccum ? " " : "") + cleaned;
      }
      // Send real-time subtitle chunk (with tags — C# strips them for display)
      sendToCs({ type: "text", text: cleaned });
    }

    // Audio chunks — send raw PCM binary to C#
    if (sc.modelTurn && sc.modelTurn.parts) {
      for (const part of sc.modelTurn.parts) {

        // Audio PCM data (base64 encoded by Gemini)
        if (part.inlineData && part.inlineData.mimeType && part.inlineData.mimeType.startsWith("audio/pcm")) {
          const pcmBytes = Buffer.from(part.inlineData.data, "base64");
          if (csClient && csClient.readyState === WebSocket.OPEN) {
            csClient.send(pcmBytes, { binary: true });
          }
        }

        // 🔥 DUPLICATE FIX: part.text fires when outputTranscription is NOT present.
        // Only send it if we have no transcription text to avoid double-accumulation.
        if (part.text && !transcriptionText) {
          turnTextAccum += (turnTextAccum ? " " : "") + part.text;
          sendToCs({ type: "text", text: part.text });
        }
      }
    }

    // 🔥 TURN COMPLETE: Send accumulated full text (with tags!) THEN signal done.
    // C# CheckMailbox reads fullText from _bridgeTextAccum — tags must be present here!
    if (sc.turnComplete) {
      // Override the last "text" event with the full accumulated text so _bridgeTextAccum is complete
      if (turnTextAccum) {
        sendToCs({ type: "fullTurnText", text: turnTextAccum });
      }
      sendToCs({ type: "turnComplete" });
      console.log("[Bridge] NPC turn complete | text:", turnTextAccum.substring(0, 80));
      turnTextAccum = ""; // reset for next turn
    }

    // Interrupted (player sent new message mid-response)
    if (sc.interrupted) {
      turnTextAccum = "";
      sendToCs({ type: "interrupted" });
    }

    return;
  }

  // Transcription outside serverContent (rare) — just stream it
  if (transcriptionText && !msg.serverContent) {
    const pipeIdx = transcriptionText.indexOf("|");
    const cleaned = pipeIdx !== -1 ? transcriptionText.substring(0, pipeIdx).trim() : transcriptionText;
    sendToCs({ type: "text", text: cleaned });
  }

  // ── 3. Tool call (if any future tags use function calling) ───
  if (msg.toolCall) {
    // Not implemented yet — log and ignore
    console.log("[Bridge] Gemini toolCall received (not implemented):", JSON.stringify(msg.toolCall));
    return;
  }

  // ── 4. Rate limit / error ────────────────────────────────────
  if (msg.error) {
    console.error("[Bridge] Gemini error:", JSON.stringify(msg.error));
    sendToCs({ type: "error", message: msg.error.message || "Unknown Gemini error" });
    return;
  }
}

// ── Send a player turn to the live session ────────────────────
function sendPlayerTurn(playerText, contextUpdate) {
  if (!geminiWs || geminiWs.readyState !== WebSocket.OPEN) return;

  // 🔥 FIX: Split contextUpdate into two separate turns so Gemini treats it as
  // a real prior exchange in history — NOT a disposable wrapper it forgets.
  // Gemini Live stateful sessions treat multi-turn history more reliably than
  // a single user message with a [CONTEXT UPDATE] prefix block.
  const turns = [];

  if (contextUpdate && contextUpdate.trim()) {
    // Turn 1: inject context as a "user" system message (Gemini will anchor this in its state)
    turns.push({
      role:  "user",
      parts: [{ text: `[🧠 LIVE WORLD STATE — READ AND REMEMBER THIS PERMANENTLY. Do NOT forget this between messages.]
${contextUpdate.trim()}
[/LIVE WORLD STATE]` }]
    });
    // Turn 2: the actual player message
    turns.push({
      role:  "user",
      parts: [{ text: playerText }]
    });
  } else {
    turns.push({
      role:  "user",
      parts: [{ text: playerText }]
    });
  }

  const turn = {
    clientContent: {
      turns:         turns,
      turn_complete: true
    }
  };

  geminiWs.send(JSON.stringify(turn));
  console.log(`[Bridge] Sent player turn: "${playerText.substring(0, 60)}..."`);
}

// ── Close Gemini session cleanly ──────────────────────────────
function closeGeminiSession() {
  sessionReady   = false;
  pendingMessage = null;   // 🔥 ZEXIVA FIX: stale queue clear on session close
  pendingHistory = [];     // 🔥 MEMORY FIX: clear stale history queue on session close
  turnTextAccum  = "";     // 🔥 TAG FIX: clear accumulated text on session end
  if (geminiWs) {
    try { 
      // 🔥 ZEXIVA FIX: Remove all listeners so the closing handshake of the old socket doesn't nullify the NEW socket!
      geminiWs.removeAllListeners();
      if (geminiWs.readyState === WebSocket.OPEN || geminiWs.readyState === WebSocket.CONNECTING) {
        geminiWs.close(); 
      }
    } catch (_) {}
    geminiWs = null;
  }
}

// ── Send JSON to C# client ────────────────────────────────────
function sendToCs(obj) {
  if (csClient && csClient.readyState === WebSocket.OPEN) {
    csClient.send(JSON.stringify(obj));
  }
}

// ── Graceful shutdown ────────────────────────────────────────
process.on("SIGINT",  () => { closeGeminiSession(); closeSceneSessions(); process.exit(0); });
process.on("SIGTERM", () => { closeGeminiSession(); closeSceneSessions(); process.exit(0); });

// ══════════════════════════════════════════════════════════════
// 🎬 SCENE SYSTEM — 2 Gemini Live sessions (NPC A + NPC B)
// ══════════════════════════════════════════════════════════════

function openSceneSession(npc, profile, key) {
  const url = `${GEMINI_WS_URL}?key=${key}`;
  const ws  = new WebSocket(url);
  if (npc === "A") sceneWsA = ws; else sceneWsB = ws;

  ws.on("open", () => {
    const sysPrompt =
      `🎬 SCENE MODE RULES (ABSOLUTE):\n` +
      `You are ${profile.name} having a face-to-face conversation with ${profile.otherName}.\n` +
      `RULE 1: ALWAYS speak directly to ${profile.otherName} — never address the player/bystander as your main conversation partner.\n` +
      `RULE 2: If a bystander (the player) says something, you may glance and say ONE short line max, then immediately return attention to ${profile.otherName}.\n` +
      `RULE 3: [STOPARGUEMENT] IS MANDATORY to end the scene. When the conversation reaches a natural conclusion, you MUST output [STOPARGUEMENT] as the VERY LAST word of your final message. Example: "Aight, later. [STOPARGUEMENT]" — WITHOUT this the scene loops forever. This is not optional.\n` +
      `RULE 4: ALL other action tags ([FIGHT],[FLEE],[FOLLOW] etc.) are BANNED in this scene.\n` +
      `🚨 AUDIO MODALITY RULE: NEVER pronounce bracketed tags out loud. Tags exist in text only.\n` +
      `🚨 TAG POSITION LAW: Tags MUST appear at the VERY END — never at the start or middle of your response.\n\n` +
      profile.systemPrompt;

    ws.send(JSON.stringify({
      setup: {
        model: `models/${GEMINI_MODEL}`,
        generation_config: {
          response_modalities: ["AUDIO"],
          speech_config: { voice_config: { prebuilt_voice_config: { voice_name: profile.voiceName || "Puck" } } }
        },
        output_audio_transcription: {},
        system_instruction: { parts: [{ text: sysPrompt }] }
      }
    }));
    console.log(`[Scene] NPC ${npc} session opening — ${profile.name} (${profile.voiceName})`);
  });

  ws.on("message", (raw) => handleSceneGeminiMessage(npc, raw));
  ws.on("error",   (e)   => {
    console.error(`[Scene] NPC ${npc} WS error:`, e.message);
    sendToCs({ type: "sceneError", message: `NPC ${npc} error: ${e.message}` });
  });
  ws.on("close", () => {
    if (npc === "A") { sceneReadyA = false; sceneWsA = null; }
    else             { sceneReadyB = false; sceneWsB = null; }
    console.log(`[Scene] NPC ${npc} session closed`);
  });
}

function handleSceneGeminiMessage(npc, raw) {
  let msg;
  try { msg = JSON.parse(raw.toString()); }
  catch (_) { return; } // ignore non-JSON (shouldn't happen with this config)

  // Setup complete
  if (msg.setupComplete !== undefined) {
    if (npc === "A") sceneReadyA = true; else sceneReadyB = true;
    console.log(`[Scene] NPC ${npc} ready. A=${sceneReadyA} B=${sceneReadyB}`);
    if (sceneReadyA && sceneReadyB) {
      sendToCs({ type: "sceneSessionReady" });
      console.log("[Scene] Both NPC sessions ready — scene can begin!");
    }
    return;
  }

  if (msg.serverContent) {
    const sc = msg.serverContent;

    // Accumulate text (outputTranscription is authoritative — includes tags)
    const trans = (msg.outputTranscription && msg.outputTranscription.text) ||
                  (sc.outputTranscription && sc.outputTranscription.text);
    if (trans) {
      const cleaned = trans.indexOf("|") !== -1 ? trans.substring(0, trans.indexOf("|")).trim() : trans;
      if (npc === "A") sceneTextA += (sceneTextA ? " " : "") + cleaned;
      else             sceneTextB += (sceneTextB ? " " : "") + cleaned;
    }

    // Accumulate audio chunks
    if (sc.modelTurn && sc.modelTurn.parts) {
      for (const part of sc.modelTurn.parts) {
        if (part.inlineData && part.inlineData.mimeType && part.inlineData.mimeType.startsWith("audio/pcm")) {
          const pcm = Buffer.from(part.inlineData.data, "base64");
          if (npc === "A") sceneAudioA.push(pcm); else sceneAudioB.push(pcm);
        }
        if (part.text && !trans) {
          if (npc === "A") sceneTextA += (sceneTextA ? " " : "") + part.text;
          else             sceneTextB += (sceneTextB ? " " : "") + part.text;
        }
      }
    }

    // Turn complete — send buffered audio + full text to C# as one JSON message
    if (sc.turnComplete) {
      const fullText  = npc === "A" ? sceneTextA : sceneTextB;
      const audioBufs = npc === "A" ? sceneAudioA : sceneAudioB;
      const audioBuf  = audioBufs.length > 0 ? Buffer.concat(audioBufs) : Buffer.alloc(0);
      sendToCs({
        type:        "sceneTurnComplete",
        npc:         npc,
        fullText:    fullText,
        audioBase64: audioBuf.length > 0 ? audioBuf.toString("base64") : ""
      });
      console.log(`[Scene] NPC ${npc} turn done | ${fullText.substring(0, 80)} | audio ${audioBuf.length}B`);
      if (npc === "A") { sceneTextA = ""; sceneAudioA = []; }
      else             { sceneTextB = ""; sceneAudioB = []; }
    }
  }

  if (msg.error) {
    console.error(`[Scene] NPC ${npc} Gemini error:`, JSON.stringify(msg.error));
    sendToCs({ type: "sceneError", message: msg.error.message || "Unknown scene error" });
  }
}

function closeSceneSessions() {
  sceneReadyA = false; sceneReadyB = false;
  sceneTextA  = "";    sceneTextB  = "";
  sceneAudioA = [];    sceneAudioB = [];
  [sceneWsA, sceneWsB].forEach(ws => {
    if (ws) {
      try { ws.removeAllListeners(); if (ws.readyState === WebSocket.OPEN) ws.close(); } catch (_) {}
    }
  });
  sceneWsA = null; sceneWsB = null;
  console.log("[Scene] Both NPC sessions closed.");
}