require("dotenv").config(); // Load environment variables from .env file
const express = require("express");
const axios = require("axios");
const bodyParser = require("body-parser");

const app = express();
const port = 3000; // Backend server will run on this port

// Middleware to parse JSON request bodies
app.use(bodyParser.json());

// API Key validation middleware
const validateApiKey = (req, res, next) => {
  const apiKey = req.headers["api-key"];
  const authHeader = req.headers["authorization"];

  // Try to extract token from Authorization: Bearer <token>
  let token = null;
  if (authHeader && authHeader.startsWith("Bearer ")) {
    token = authHeader.split(" ")[1];
  }

  const providedKey = apiKey || token;

  if (!providedKey) {
    return res.status(400).json({ error: "API key is missing" });
  }

  if (providedKey !== process.env.API_KEY) {
    return res.status(403).json({ error: "Invalid API key" });
  }

  next(); // Proceed if the API key is valid
};

// Forward request to localhost:11434 (ollama)

app.post("/api/generate/chat/completions", validateApiKey, async (req, res) => {
  try {
    const { model, messages } = req.body;

    const prompt = messages
      .map(
        (msg) =>
          `${msg.role === "system" ? "" : msg.role + ": "}${msg.content}`,
      )
      .join("\n");

    console.log("Generated prompt for Ollama:\n", prompt);

    const response = await axios.post(
      "http://localhost:11434/api/generate",
      {
        model: model || "deepseek-r1:latest",
        prompt,
        stream: true,
      },
      { responseType: "stream" },
    );

    let finalAnswer = "";
    let lastChunk = "";

    response.data.on("data", (chunk) => {
      const lines = chunk.toString("utf8").split("\n").filter(Boolean);

      for (const line of lines) {
        try {
          const json = JSON.parse(line);
          lastChunk = json.response || lastChunk;
          if (json.done) {
            finalAnswer = lastChunk;
          }
        } catch (err) {
          console.warn("Chunk parsing failed:", err);
        }
      }
    });

    response.data.on("end", () => {
      console.log("\n🎯 Final Answer Only:\n", finalAnswer);
      res.json({ response: finalAnswer.trim() });
    });

    response.data.on("error", (err) => {
      console.error("Ollama stream error:", err);
      res
        .status(500)
        .json({ error: "Ollama stream error", message: err.message });
    });
  } catch (error) {
    console.error(
      "Error communicating with Ollama:",
      error.response?.data || error.message,
    );
    res
      .status(500)
      .json({ error: "Internal Server Error", message: error.message });
  }
});

// Forward request to localhost:11434 (ollama)
app.post("/api/generate", validateApiKey, async (req, res) => {
  try {
    // Forwarding the request to localhost:11434 with the prompt
    const authHeader = req.headers["authorization"];
    console.log("Authorization header:", authHeader);
    console.log("checking api", apiKey !== process.env.API_KEY);
    console.log("Body: ", req.body);
    const response = await axios.post(
      "http://localhost:11434/api/generate",
      req.body,
    );

    // Send the response from localhost:11434 back to the client
    res.status(response.status).json(response.data);
  } catch (error) {
    // Enhanced error logging
    console.error(
      "Error forwarding request to localhost:11434:",
      error.response ? error.response.data : error.message,
    );
    res
      .status(500)
      .json({ error: "Internal Server Error", message: error.message });
  }
});

// Start the server
app.listen(port, () => {
  console.log(`Server running on http://localhost:${port}`);
});