index.js

const express = require("express");
const { orchestrate, prepareOrchestration } = require("./lib/orchestrator");
const { updateActivity, startRequest, endRequest } = require("./lib/idle");
const { getProvider } = require("./lib/provider-selector");
const { addMessage } = require("./lib/db");
require("dotenv").config();

const provider = getProvider();
require("dotenv").config();

const app = express();
app.use(express.json());

const PORT = process.env.PORT || 3000;
const MODEL_NAME = process.env.MODEL_NAME || "stable-unified";
const DECISION_MODEL = process.env.DECISION_MODEL || "llama3.1:70b";
const CHOICE_MODEL = process.env.CHOICE_MODEL || "llama3.1:70b";

// OpenAI Compatible Models Endpoint
app.get("/v1/models", (req, res) => {
  res.json({
    object: "list",
    data: [
      {
        id: MODEL_NAME,
        object: "model",
        created: 1677610602,
        owned_by: "stable"
      }
    ]
  });
});

// Health check endpoint
app.get("/health", (req, res) => {
  res.json({ status: "ok" });
});

// OpenAI Compatible Endpoint
app.post("/v1/chat/completions", async (req, res) => {
  const chatId = `chatcmpl-${Date.now()}`;
  try {
    const { messages, stream } = req.body;

    if (!messages || !Array.isArray(messages)) {
      return res.status(400).json({ error: "Invalid request: messages must be an array" });
    }

    startRequest();

    if (stream) {
      // 1. Send headers immediately to prevent timeouts
      res.setHeader('Content-Type', 'text/event-stream');
      res.setHeader('Cache-Control', 'no-cache');
      res.setHeader('Connection', 'keep-alive');

      // 2. Send initial role+empty content chunk to fully satisfy OpenAI pattern match
      const initialData = {
        id: chatId,
        object: "chat.completion.chunk",
        created: Math.floor(Date.now() / 1000),
        model: MODEL_NAME,
        choices: [{ index: 0, delta: { role: "assistant", content: "" }, finish_reason: null }]
      };
      res.write(`data: ${JSON.stringify(initialData)}\n\n`);

      // 3. Keep-alive heartbeat loop while orchestration is thinking
      const heartbeat = setInterval(() => {
        if (!res.writableEnded) {
          res.write(': keep-alive\n\n');
        }
      }, 5000);

      try {
        // 4. Perform heavy orchestration (Best of N)
        const response = await orchestrate(messages);
        clearInterval(heartbeat);

        const fullContent = response.message.content;
        // OpenAI-like chunking: send in smaller pieces to simulate streaming
        const chunkSize = 20;
        for (let i = 0; i < fullContent.length; i += chunkSize) {
          const chunk = fullContent.slice(i, i + chunkSize);
          const data = {
            id: chatId,
            object: "chat.completion.chunk",
            created: Math.floor(Date.now() / 1000),
            model: MODEL_NAME,
            choices: [
              {
                index: 0,
                delta: { content: chunk },
                finish_reason: (i + chunkSize >= fullContent.length) ? "stop" : null
              }
            ]
          };
          res.write(`data: ${JSON.stringify(data)}\n\n`);
          // Small delay to make it look like streaming
          await new Promise(resolve => setTimeout(resolve, 5));
        }

        res.write("data: [DONE]\n\n");
        res.end();
      } catch (innerErr) {
        clearInterval(heartbeat);
        throw innerErr;
      }
      return;
    }

    const response = await orchestrate(messages);

    // Format to OpenAI standard
    res.json({
      id: chatId,
      object: "chat.completion",
      created: Math.floor(Date.now() / 1000),
      model: MODEL_NAME,
      choices: [
        {
          index: 0,
          message: {
            role: "assistant",
            content: response.message.content
          },
          finish_reason: "stop"
        }
      ],
      usage: {
        prompt_tokens: -1,
        completion_tokens: -1,
        total_tokens: -1
      }
    });

  } catch (err) {
    console.error("API Error:", err);
    if (!res.headersSent) {
      res.status(500).json({ error: "Internal Server Error", message: err.message });
    } else {
      const errorData = {
        id: chatId,
        object: "chat.completion.chunk",
        created: Math.floor(Date.now() / 1000),
        model: MODEL_NAME,
        choices: [{
          index: 0,
          delta: { content: `\n\n[API Error]: ${err.message}` },
          finish_reason: "stop"
        }]
      };
      res.write(`data: ${JSON.stringify(errorData)}\n\n`);
      res.write("data: [DONE]\n\n");
      res.end();
    }
  } finally {
    endRequest();
  }
});

app.listen(PORT, async () => {
  console.log(`Stable API running on port ${PORT}`);
  console.log(`Decision Model: ${DECISION_MODEL}`);
  console.log(`Choice Model: ${CHOICE_MODEL}`);

  // Ensure models are available on startup
  await provider.isAvailable();
});