const express = require('express');
const fetch = require('node-fetch');
const { v4: uuidv4 } = require('uuid');
const app = express();
require('dotenv').config();

const modelMapping = {
  "gpt-4o-mini": "GPT-4o mini",
  "claude-haiku": "Claude Haiku",
  "llama-3": "Llama 3",
  "gemini-1.5": "Gemini 1.5",
  "gemini-flash": "Gemini Flash",
  "command-r": "Command R"
};

app.use(express.json());

const getCurrentTimestamp = () => Math.floor(Date.now() / 1000);

const splitIntoChunks = (text, chunkSize) => {
  const chunks = [];
  for (let i = 0; i < text.length; i += chunkSize) {
    chunks.push(text.slice(i, i + chunkSize));
  }
  return chunks;
};

const getTempUserID = async () => {
  try {
    const response = await fetch('https://playground.julius.ai/api/temp_user_id');
    const data = await response.json();
    return data.temp_user_id;
  } catch (error) {
    throw new Error('Failed to get temp user ID');
  }
};

const sendToJulius = async (tempUserID, message, model) => {
  const conversationID = uuidv4();
  const juliusReq = {
    message: {
      content: message,
      role: "user"
    },
    provider: "default",
    chat_mode: "auto",
    client_version: "20240130",
    theme: "dark",
    selectedModels: [model]
  };

  try {
    const response = await fetch('https://playground.julius.ai/api/chat/message', {
      method: 'POST',
      headers: {
        'is-demo': tempUserID,
        'Content-Type': 'application/json',
        'Platform': 'web',
        'conversation-id': conversationID,
        'User-Agent': 'Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/133.0.0.0 Safari/537.36'
      },
      body: JSON.stringify(juliusReq)
    });

    const reader = response.body.pipeThrough(new TextDecoderStream()).getReader();
    let fullResponse = '';

    while (true) {
      const { done, value } = await reader.read();
      if (done) break;
      
      const lines = value.split('\n');
      for (const line of lines) {
        try {
          const jsonResp = JSON.parse(line);
          if (jsonResp.content) {
            fullResponse += jsonResp.content;
          }
        } catch (e) {
          // 忽略解析错误
        }
      }
    }

    return fullResponse;
  } catch (error) {
    throw new Error('Failed to send to Julius');
  }
};

app.all('*', async (req, res) => {
  // 身份验证
  const authToken = process.env.AUTH_TOKEN;
  if (authToken) {
    const requestToken = req.headers.authorization?.replace('Bearer ', '');
    if (!requestToken || requestToken !== authToken) {
      return res.status(401).send('Access Denied');
    }
  }

  // 健康检查
  if (req.path !== '/v1/chat/completions') {
    return res.json({
      status: "Julius2Api Service Running...",
      message: "MoLoveSze..."
    });
  }

  if (req.method !== 'POST') {
    return res.status(405).send('Method not allowed');
  }

  try {
    const openAIReq = req.body;
    const isStream = openAIReq.stream || false;

    // 模型映射
    const mappedModel = modelMapping[openAIReq.model] || 'GPT-4o mini';
    openAIReq.model = mappedModel;

    const tempUserID = await getTempUserID();
    const juliusResp = await sendToJulius(tempUserID, 
      openAIReq.messages[openAIReq.messages.length - 1].content, 
      mappedModel
    );

    const respId = `chatcmpl-${tempUserID}`;
    const created = getCurrentTimestamp();

    if (isStream) {
      res.writeHead(200, {
        'Content-Type': 'text/event-stream',
        'Cache-Control': 'no-cache',
        'Connection': 'keep-alive'
      });

      // 发送初始响应
      const firstResponse = {
        id: respId,
        object: "chat.completion.chunk",
        created: created,
        model: mappedModel,
        choices: [{
          delta: { role: "assistant" },
          index: 0
        }]
      };
      res.write(`data: ${JSON.stringify(firstResponse)}\n\n`);

      // 分块发送内容
      const chunks = splitIntoChunks(juliusResp, 50);
      chunks.forEach((chunk, index) => {
        const response = {
          id: respId,
          object: "chat.completion.chunk",
          created: created,
          model: mappedModel,
          choices: [{
            delta: { content: chunk },
            index: 0,
            finish_reason: index === chunks.length - 1 ? 'stop' : null
          }]
        };
        res.write(`data: ${JSON.stringify(response)}\n\n`);
      });

      res.write('data: [DONE]\n\n');
      res.end();
    } else {
      const response = {
        id: respId,
        object: "chat.completion",
        created: created,
        model: mappedModel,
        choices: [{
          message: {
            role: "assistant",
            content: juliusResp
          },
          finish_reason: "stop"
        }]
      };
      res.json(response);
    }
  } catch (error) {
    console.error(error);
    res.status(500).send(error.message);
  }
});

const PORT = process.env.PORT || 3000;
app.listen(PORT, () => {
  console.log(`Server running on port ${PORT}`);
});