import express from 'express';
import axios from 'axios';
import { v4 as uuidv4 } from 'uuid';

const app = express();

// 中间件
app.use(express.json());
app.use(express.text());

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",
};

// 认证中间件
const authMiddleware = (req, res, next) => {
  const authToken = process.env.AUTH_TOKEN;
  if (authToken) {
    const requestToken = req.headers.authorization || '';
    const token = requestToken.replace('Bearer ', '');
    if (token !== authToken) {
      return res.status(401).send('Access Denied');
    }
  }
  next();
};

// 获取临时用户ID
const getTempUserID = async () => {
  const response = await axios.get('https://playground.julius.ai/api/temp_user_id');
  return response.data.temp_user_id;
};

// 分割文本为块
const splitIntoChunks = (text, chunkSize) => {
  const chars = Array.from(text);
  const chunks = [];
  for (let i = 0; i < chars.length; i += chunkSize) {
    chunks.push(chars.slice(i, i + chunkSize).join(''));
  }
  return chunks;
};

// 发送消息到 Julius
const sendToJulius = async (tempUserID, message, model) => {
  const conversationID = uuidv4();
  const data = {
    message: {
      content: message,
      role: "user"
    },
    provider: "default",
    chat_mode: "auto",
    client_version: "20240130",
    theme: "dark",
    new_images: null,
    new_attachments: null,
    dataframe_format: "json",
    selectedModels: [model]
  };

  const response = await axios({
    method: 'post',
    url: 'https://playground.julius.ai/api/chat/message',
    data,
    headers: {
      'is-demo': tempUserID,
      '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',
      'interactive-charts': 'true',
      'use-dict': 'true',
      'Gcs': 'true',
      'Is-Native': 'false',
      'sec-ch-ua-platform': 'Windows'
    },
    responseType: 'text'
  });

  return response.data.split('\n')
    .filter(line => line.trim())
    .map(line => {
      try { return JSON.parse(line); } 
      catch { return null; }
    })
    .filter(Boolean)
    .reduce((acc, item) => acc + (item.content || ''), '');
};

// 处理聊天请求
app.post('/v1/chat/completions', authMiddleware, async (req, res) => {
  try {
    const { messages, model: originalModel, stream } = req.body;
    const model = modelMapping[originalModel] || originalModel;
    const tempUserID = await getTempUserID();
    const content = messages[messages.length - 1].content;

    // 获取 Julius 响应
    const juliusResp = await sendToJulius(tempUserID, content, model);

    // 流式响应
    if (stream) {
      res.setHeader('Content-Type', 'text/event-stream');
      res.setHeader('Cache-Control', 'no-cache');
      res.setHeader('Connection', 'keep-alive');

      const respId = `chatcmpl-${tempUserID}`;
      const created = Math.floor(Date.now() / 1000);
      const chunks = splitIntoChunks(juliusResp, 50);

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

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

      // 结束标志
      res.write('data: [DONE]\n\n');
      return res.end();
    }

    // 普通响应
    res.json({
      id: `chatcmpl-${tempUserID}`,
      object: "chat.completion",
      created: Math.floor(Date.now() / 1000),
      model,
      choices: [{
        message: {
          role: "assistant",
          content: juliusResp
        },
        finish_reason: "stop",
        index: 0
      }]
    });

  } catch (error) {
    console.error('Error:', error);
    res.status(500).json({ error: error.message });
  }
});

// 根路径响应
app.all('*', (req, res) => {
  res.json({
    status: "Julius2Api Service Running...",
    message: "MoLoveSze..."
  });
});

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