const express = require('express');
const axios = require('axios');
const fs = require('fs');
const path = require('path');
const os = require('os');
const bytes = require('bytes');
const cp = require('child_process');
//const swaggerUi = require('swagger-ui-express');

const app = express();
app.set('json spaces', 2);

const geminiApiKey = 'AIzaSyCRpX4SnNQl__BNLcVWqgzKdbIlt06vod0';
const geminiApiUrl = 'https://generativelanguage.googleapis.com/v1/models/gemini-pro:generateContent';

const geminiTmpDir = path.join(require('os').tmpdir(), 'gemini_sessions');
const deepseekTmpDir = path.join(require('os').tmpdir(), 'deepseek_sessions');
const qwenTmpDir = path.join(require('os').tmpdir(), 'qwen_sessions');
const mistralTmpDir = path.join(require('os').tmpdir(), 'mistral_sessions');

if (!fs.existsSync(geminiTmpDir)) fs.mkdirSync(geminiTmpDir);
if (!fs.existsSync(deepseekTmpDir)) fs.mkdirSync(deepseekTmpDir);
if (!fs.existsSync(qwenTmpDir)) fs.mkdirSync(qwenTmpDir);
if (!fs.existsSync(mistralTmpDir)) fs.mkdirSync(mistralTmpDir);

function getSessionFile(sessionId, model) {
  const dir = model === 'gemini' ? geminiTmpDir : deepseekTmpDir || model === 'qwen' ? qwenTmpDir : mistralTmpDir;
  return path.join(dir, `${sessionId}.json`);
}

async function gemini(sessionId, prompt) {
  const sessionFile = getSessionFile(sessionId, 'gemini');
  let history = [];
  if (fs.existsSync(sessionFile)) {
    history = JSON.parse(fs.readFileSync(sessionFile, 'utf8'));
  }
  history.push({ role: 'user', content: { parts: [{ text: prompt }] } });

  const response = await axios.post(`${geminiApiUrl}?key=${geminiApiKey}`, {  
    contents: history.map(entry => ({  
        role: entry.role === 'assistant' ? 'model' : 'user',  
        parts: entry.content.parts  
    }))  
  }, {  
    headers: { 'Content-Type': 'application/json' }  
  });  
  
  const reply = response.data.candidates?.[0]?.content?.parts?.[0]?.text || 'No response';  
  history.push({ role: 'assistant', content: { parts: [{ text: reply }] } });  
  fs.writeFileSync(sessionFile, json.stringify(history));  
  return reply;
}

async function deepseek(sessionId, prompt, r1 = true) {
  const sessionFile = getSessionFile(sessionId, 'deepseek');
  let history = [];
  if (fs.existsSync(sessionFile)) {
    history = JSON.parse(fs.readFileSync(sessionId, 'utf8'));
  }
  history.push({ role: 'user', content: prompt });

  let { data } = await axios.post("https://ai.clauodflare.workers.dev/chat", {  
    "model": r1 ? "@cf/deepseek-ai/deepseek-r1-distill-qwen-32b" : "@cf/deepseek-ai/deepseek-math-7b-instruct",  
    "messages": history  
  }).catch(e => e.response);  
  
  if (!data.success) throw new Error(JSON.stringify(data, null, 2));  
  let response = data.data.response  
  history.push({ role: 'assistant', content: r1 ? response.split("</think>")[1]?.trim() : response });  
  fs.writeFileSync(sessionFile, JSON.stringify(history));  
  return response;
}

async function qwen(sessionId, prompt, selModel = 0) {
  const model = ["@cf/qwen/qwen1.5-0.5b-chat", "@cf/qwen/qwen1.5-1.8b-chat", "@cf/qwen/qwen1.5-14b-chat-awq", "@cf/qwen/qwen1.5-7b-chat-awq"]
  const sessionFile = getSessionFile(sessionId, 'qwen');
  let history = [];
  if (fs.existsSync(sessionFile)) {
    history = JSON.parse(fs.readFileSync(sessionFile, 'utf8'));
  }
  history.push({ role: 'user', content: prompt });

  let { data } = await axios.post("https://ai.clauodflare.workers.dev/chat", {  
    "model": model[selModel],  
    "messages": history  
  }).catch(e => e.response);  
  
  if (!data.success) throw new Error(JSON.stringify(data, null, 2));  
  let response = data.data.response  
  history.push({ role: 'assistant', content: response });  
  fs.writeFileSync(sessionFile, JSON.stringify(history));  
  return response;
}

async function google(sessionId, prompt, selModel = 0) {
  const model = ["@hf/google/gemma-7b-it", "@cf/google/gemma-7b-it-lora", "@cf/google/gemma-2b-it-lora"]
  const sessionFile = getSessionFile(sessionId, 'qwen');
  let history = [];
  if (fs.existsSync(sessionFile)) {
    history = JSON.parse(fs.readFileSync(sessionFile, 'utf8'));
  }
  history.push({ role: 'user', content: prompt });

  let { data } = await axios.post("https://ai.clauodflare.workers.dev/chat", {  
    "model": model[selModel],  
    "messages": history  
  }).catch(e => e.response);  
  
  if (!data.success) throw new Error(JSON.stringify(data, null, 2));  
  let response = data.data.response  
  history.push({ role: 'assistant', content: response });  
  fs.writeFileSync(sessionFile, JSON.stringify(history));  
  return response;
}

async function mistral(sessionId, prompt, selModel = 0) {
  const model = ["@hf/mistral/mistral-7b-instruct-v0.2", "@cf/mistral/mistral-7b-instruct-v0.2-lora", "@cf/mistral/mistral-7b-instruct-v0.1"]
  const sessionFile = getSessionFile(sessionId, 'mistral');
  let history = [];
  if (fs.existsSync(sessionFile)) {
    history = JSON.parse(fs.readFileSync(sessionFile, 'utf8'));
  }
  history.push({ role: 'user', content: prompt });

  let { data } = await axios.post("https://ai.clauodflare.workers.dev/chat", {  
    "model": model[selModel],  
    "messages": history  
  }).catch(e => e.response);  
  
  if (!data.success) throw new Error(JSON.stringify(data, null, 2));  
  let response = data.data.response  
  history.push({ role: 'assistant', content: response });  
  fs.writeFileSync(sessionFile, JSON.stringify(history));  
  return response;
}

const utils = {
  formatSize: (n) => bytes(+n, { unitSeparator: ' ' })
}

app.all('/', (_, res) => {
	const status = {}
	status['diskUsage'] = cp.execSync('du -sh').toString().split('M')[0] + ' MB'

	const used = process.memoryUsage()
	for (let x in used) status[x] = utils.formatSize(used[x])

	const totalmem = os.totalmem()
	const freemem = os.freemem()
	status['memoryUsage'] =
		`${utils.formatSize(totalmem - freemem)} / ${utils.formatSize(totalmem)}`
	res.json({
		ai: [
          {
            'google': [
              'gemini-1.5-pro',
              'gemma-7b-it',
              'gemma-7b-it-lora',
              'gemma-2b-it-lora'
            ]
          },
          {
            'deepseek-ai': [
              'deepseek-r1-distill-qwen-32b',
              'deepseek-math-7b-instruct'
            ]
          },
          {
            'qwen': [
              'qwen1.5-0.5b-chat',
              'qwen1.5-1.8b-chat', 
              'qwen1.5-14b-chat-awq', 
              'qwen1.5-7b-chat-awq'
            ]
          }, 
          {
            'mistral': [
              'mistral-7b-instruct-v0.2', 
              'mistral-7b-instruct-v0.2-lora', 
              'mistral-7b-instruct-v0.1'
            ]
          }
        ],
		uptime: new Date(process.uptime() * 1000).toUTCString().split(' ')[4],
		status
	})
})

/* Google */
app.get('/google/gemini-1.5-pro', async (req, res) => {
  try {
    const { question, sessionId } = req.query;
    if (!question || !sessionId) return res.status(400).json({ error: 'Missing parameters' });
    const reply = await gemini(sessionId, question);
    res.json({ reply });
  } catch (error) {
    res.status(500).json({ error: error.response?.data || error.message });
  }
});

app.get('/google/gemma-7b-it', async (req, res) => {
  try {
    const { question, sessionId } = req.query;
    if (!question || !sessionId) return res.status(400).json({ error: 'Missing parameters' });
    const reply = await google(sessionId, question);
    res.json({ reply });
  } catch (error) {
    res.status(500).json({ error: error.response?.data || error.message });
  }
});

app.get('/google/gemma-7b-it-lora', async (req, res) => {
  try {
    const { question, sessionId } = req.query;
    if (!question || !sessionId) return res.status(400).json({ error: 'Missing parameters' });
    const reply = await gemini(sessionId, question, 1);
    res.json({ reply });
  } catch (error) {
    res.status(500).json({ error: error.response?.data || error.message });
  }
});

app.get('/google/gemma-2b-it-lora', async (req, res) => {
  try {
    const { question, sessionId } = req.query;
    if (!question || !sessionId) return res.status(400).json({ error: 'Missing parameters' });
    const reply = await gemini(sessionId, question, 2);
    res.json({ reply });
  } catch (error) {
    res.status(500).json({ error: error.response?.data || error.message });
  }
});

/* DeepSeek-AI */
app.get('/deepseek-ai/deepseek-r1-distill-qwen-32b', async (req, res) => {
  try {
    const { question, sessionId } = req.query;
    if (!question || !sessionId) return res.status(400).json({ error: 'Missing parameters' });
    const reply = await deepseek(sessionId, question);
    res.json({ reply: reply.split("</think>")?.pop()?.trim() || reply.split("</think>")[1]?.trim(), think: reply.split('</think>')[0]?.trim() });
  } catch (error) {
    res.status(500).json({ error: error.response?.data || error.message });
  }
});

app.get('/depseek-ai/deepseek-math-7b-instruct', async (req, res) => {
  try {
    const { question, sessionId } = req.query;
    if (!question || !sessionId) return res.status(400).json({ error: 'Missing parameters' });
    const reply = await deepseek(sessionId, question, false);
    res.json({ reply });
  } catch (error) {
    res.status(500).json({ error: error.response?.data || error.message });
  }
});

/* QWEN */
app.get('/qwen/qwen1.5-0.5b-chat', async (req, res) => {
  try {
    const { question, sessionId } = req.query;
    if (!question || !sessionId) return res.status(400).json({ error: 'Missing parameters' });
    const reply = await qwen(sessionId, question);
    res.json({ reply });
  } catch (error) {
    res.status(500).json({ error: error.response?.data || error.message });
  }
});

app.get('/qwen/qwen1.5-1.8b-chat', async (req, res) => {
  try {
    const { question, sessionId } = req.query;
    if (!question || !sessionId) return res.status(400).json({ error: 'Missing parameters' });
    const reply = await qwen(sessionId, question, 1);
    res.json({ reply });
  } catch (error) {
    res.status(500).json({ error: error.response?.data || error.message });
  }
});

app.get('/qwen/qwen1.5-14b-chat-awq', async (req, res) => {
  try {
    const { question, sessionId } = req.query;
    if (!question || !sessionId) return res.status(400).json({ error: 'Missing parameters' });
    const reply = await qwen(sessionId, question, 2);
    res.json({ reply });
  } catch (error) {
    res.status(500).json({ error: error.response?.data || error.message });
  }
});

app.get('/qwen/qwen1.5-7b-chat-awq', async (req, res) => {
  try {
    const { question, sessionId } = req.query;
    if (!question || !sessionId) return res.status(400).json({ error: 'Missing parameters' });
    const reply = await qwen(sessionId, question, 3);
    res.json({ reply });
  } catch (error) {
    res.status(500).json({ error: error.response?.data || error.message });
  }
});

/* Mistral */
app.get('/mistral/mistral-7b-instruct-v0.2', async (req, res) => {
  try {
    const { question, sessionId } = req.query;
    if (!question || !sessionId) return res.status(400).json({ error: 'Missing parameters' });
    const reply = await mistral(sessionId, question);
    res.json({ reply });
  } catch (error) {
    res.status(500).json({ error: error.response?.data || error.message });
  }
});

app.get('/mistral/mistral-7b-instruct-v0.2-lora', async (req, res) => {
  try {
    const { question, sessionId } = req.query;
    if (!question || !sessionId) return res.status(400).json({ error: 'Missing parameters' });
    const reply = await mistral(sessionId, question, 1);
    res.json({ reply });
  } catch (error) {
    res.status(500).json({ error: error.response?.data || error.message });
  }
});

app.get('/mistral/mistral-7b-instruct-v0.1', async (req, res) => {
  try {
    const { question, sessionId } = req.query;
    if (!question || !sessionId) return res.status(400).json({ error: 'Missing parameters' });
    const reply = await mistral(sessionId, question, 2);
    res.json({ reply });
  } catch (error) {
    res.status(500).json({ error: error.response?.data || error.message });
  }
});

app.listen(7860);