const express = require('express');
const fetch = require('node-fetch');
const crypto = require('crypto');

const app = express();
const port = 8000;

const NOTDIAMOND_URL = 'https://chat.notdiamond.ai/mini-chat';
const NOTDIAMOND_HEADERS = {
  'Content-Type': 'application/json',
  'next-action': '8189eb37107121e024940f588629a394a594e6a4'
};
const SUPABASE_URL = 'https://spuckhogycrxcbomznwo.supabase.co';
const SUPABASE_KEY = 'eyJhbGciOiJIUzI1NiIsInR5cCI6IkpXVCJ9.eyJpc3MiOiJzdXBhYmFzZSIsInJlZiI6InNwdWNraG9neWNyeGNib216bndvIiwicm9sZSI6ImFub24iLCJpYXQiOjE3MDcyNDYwMzksImV4cCI6MjAyMjgyMjAzOX0.tvlGT7NZY8bijMjNIu1WhAtPnSKuDeYhtveo4DRt6xg';

const DEFAULT_MODEL = 'gpt-4o';

const MODEL_MAPPING = {
    "gpt-4o-mini": {
        "provider": "openai",
        "mapping": "gpt-4o-mini"
    },
    "gpt-4o": {
        "provider": "openai",
        "mapping": "gpt-4o"
    },
    "gpt-4-turbo": {
        "provider": "openai",
        "mapping": "gpt-4-turbo-2024-04-09"
    },
    "gemini-1.5-pro-latest": {
        "provider": "google",
        "mapping": "gemini-1.5-pro-latest"
    },
    "gemini-1.5-flash-latest": {
        "provider": "google",
        "mapping": "gemini-1.5-flash-latest"
    },
    "llama-3.1-70b-instruct": {
        "provider": "meta",
        "mapping": "Meta-Llama-3.1-70B-Instruct-Turbo"
    },
    "llama-3.1-405b-instruct": {
        "provider": "meta",
        "mapping": "Meta-Llama-3.1-405B-Instruct-Turbo"
    },
    "claude-3-5-sonnet-20240620": {
        "provider": "anthropic",
        "mapping": "claude-3-5-sonnet-20240620"
    },
    "claude-3-haiku-20240307": {
        "provider": "anthropic",
        "mapping": "claude-3-haiku-20240307"
    },
    "perplexity": {
        "provider": "perplexity",
        "mapping": "llama-3.1-sonar-large-128k-online"
    },
    "mistral-large-2407": {
        "provider": "mistral",
        "mapping": "mistral.mistral-large-2407-v1:0"
    }
};

const authCache = new Map();

async function getAuthCookie(req) {
  const authHeader = req.headers['authorization'];
  if (authHeader && authHeader.startsWith('Bearer ')) {
    const [email, password] = authHeader.slice(7).split('|');
    if (!email || !password) {
      throw new Error('Invalid authentication format');
    }

    const cacheKey = `${email}:${password}`;
    const cachedAuth = authCache.get(cacheKey);
    
    if (cachedAuth && cachedAuth.expiresAt > Date.now()) {
      return {
        cookie: cachedAuth.cookie,
        cacheHit: true,
        expiresAt: cachedAuth.expiresAt
      };
    }

    const response = await fetch(`${SUPABASE_URL}/auth/v1/token?grant_type=password`, {
      method: 'POST',
      headers: {
        'accept': '*/*',
        'accept-language': 'zh-CN,zh;q=0.9',
        'apikey': SUPABASE_KEY,
        'authorization': `Bearer ${SUPABASE_KEY}`,
        'content-type': 'application/json;charset=UTF-8',
        'origin': 'https://chat.notdiamond.ai',
        'referer': 'https://chat.notdiamond.ai/',
        'user-agent': 'Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/128.0.0.0 Safari/537.36',
        'x-client-info': 'supabase-ssr/0.4.0',
        'x-supabase-api-version': '2024-01-01'
      },
      body: JSON.stringify({
        email: email,
        password: password,
        gotrue_meta_security: {}
      })
    });

    if (!response.ok) {
      throw new Error('Authentication failed');
    }

    const responseData = await response.json();
    const baseCookie = Buffer.from(JSON.stringify(responseData)).toString('base64');
    const cookie = `sb-spuckhogycrxcbomznwo-auth-token=base64-${baseCookie}`;
    const expiresAt = Date.now() + (responseData.expires_in * 1000);

    authCache.set(cacheKey, {
      cookie: cookie,
      expiresAt: expiresAt
    });

    return {
      cookie: cookie,
      cacheHit: false,
      expiresAt: expiresAt
    };
  }
  return null;
}

function createOpenAIChunk(content, model, finishReason = null) {
  return {
    id: `chatcmpl-${crypto.randomUUID()}`,
    object: "chat.completion.chunk",
    created: Math.floor(Date.now() / 1000),
    model: model,
    system_fingerprint: "fp_4e2b2da518",
    choices: [
      {
        index: 0,
        delta: content ? { content: content } : {},
        logprobs: null,
        finish_reason: finishReason
      }
    ],
    usage: null
  };
}

async function* streamNotdiamondResponse(response, model) {
  let buffer = "";
  let lastContent = "";

  function processDollars(str) {
    return str.replace(/\${2,}/, match => match.slice(1));
  }

  for await (const chunk of response.body) {
    buffer += chunk.toString();
    let lines = buffer.split('\n');
    buffer = lines.pop() || "";

    for (const line of lines) {
      if (line.trim() === '') continue;
      
      try {
        const jsonMatch = line.match(/\{.*\}/);
        if (jsonMatch) {
          const data = JSON.parse(jsonMatch[0]);
          let content = '';
          
          if (data.output?.curr !== undefined) {
            content = processDollars(data.output.curr);
            lastContent = content;
          } else if (data.curr !== undefined) {
            content = processDollars(data.curr);
            lastContent = content;
          } else if (data.diff !== undefined && Array.isArray(data.diff) && data.diff.length > 1) {
            const newContent = processDollars(data.diff[1]);
            content = lastContent + newContent;
            lastContent = content;
          } else if (data.diff !== undefined && Array.isArray(data.diff) && data.diff.length === 1) {
            content = lastContent;  // If diff is empty, use the last content
          }
          
          if (content) {
            yield createOpenAIChunk(content, model);
          }
        }
      } catch (error) {
        console.error('Error processing line:', error, 'Raw line:', line);
      }
    }
  }

  if (buffer.trim()) {
    try {
      const jsonMatch = buffer.match(/\{.*\}/);
      if (jsonMatch) {
        const data = JSON.parse(jsonMatch[0]);
        let content = '';
        if (data.output?.curr !== undefined) {
          content = processDollars(data.output.curr);
        } else if (data.curr !== undefined) {
          content = processDollars(data.curr);
        } else if (data.diff !== undefined && Array.isArray(data.diff) && data.diff.length > 1) {
          const newContent = processDollars(data.diff[1]);
          content = lastContent + newContent;
        } else if (data.diff !== undefined && Array.isArray(data.diff) && data.diff.length === 1) {
          content = lastContent;
        }
        if (content) {
          yield createOpenAIChunk(content, model);
        }
      }
    } catch (error) {
      console.error('Error processing final buffer:', error);
    }
  }

  yield createOpenAIChunk('', model, 'stop');
}

app.use(express.json());

app.get('/', (req, res) => {
  res.json({
    service: "AI Chat Completion Proxy",
    usage: {
      endpoint: "/ai/v1/chat/completions",
      method: "POST",
      headers: {
        "Content-Type": "application/json",
        "Authorization": "Bearer YOUR_AUTH_KEY"
      },
      body: {
        model: "One of: " + Object.keys(MODEL_MAPPING).join(", "),
        messages: [
          { role: "system", content: "You are a helpful assistant." },
          { role: "user", content: "Hello, who are you?" }
        ],
        stream: false,
        temperature: 0.7
      }
    },
    availableModels: Object.keys(MODEL_MAPPING),
    note: "Replace YOUR_AUTH_KEY with your actual authentication key."
  });
});

app.get('/ai/v1/models', async (req, res) => {
  try {
    const authCookie = await getAuthCookie(req);
    if (!authCookie) {
      return res.status(401).json({ error: 'Unauthorized' });
    }

    res.json({
      object: "list",
      data: Object.entries(MODEL_MAPPING).map(([id, info]) => ({
        id: id,
        object: "model",
        created: Math.floor(Date.now() / 1000),
        owned_by: info.provider,
        permission: [],
        root: id,
        parent: null,
      }))
    });

    console.log(`${new Date().toISOString()} - ${req.ip || req.connection.remoteAddress} - GET /ai/v1/models - 200`);
  } catch (error) {
    res.status(401).json({ error: 'Authentication failed', details: error.message });
  }
});

app.post('/ai/v1/chat/completions', async (req, res) => {
  const startTime = new Date();
  const clientIP = req.ip || req.connection.remoteAddress;
  let userModel = req.body.model || DEFAULT_MODEL;
  let status = 200;
  let cookieCacheHit = false;
  let cookieExpiresAt = null;

  try {
    const authResult = await getAuthCookie(req);
    if (!authResult) {
      status = 401;
      throw new Error('Unauthorized');
    }

    cookieCacheHit = authResult.cacheHit;
    cookieExpiresAt = new Date(authResult.expiresAt).toISOString();

    let requestData = req.body;
    let messages = requestData.messages;
    let modelInfo = MODEL_MAPPING[userModel];
    if (!modelInfo) {
      status = 400;
      throw new Error('Invalid model specified');
    }
    let model = modelInfo.mapping;
    const stream = requestData.stream || false;

    if (!messages || !Array.isArray(messages)) {
      status = 400;
      throw new Error('Invalid request body: messages should be an array of message objects');
    }

    const payload = {
      messages: messages,
      model: model,
      stream: stream,
      frequency_penalty: requestData.frequency_penalty || 0,
      presence_penalty: requestData.presence_penalty || 0,
      temperature: requestData.temperature || 0.6,
      top_p: requestData.top_p || 1
    };

    const headers = {
      'Content-Type': 'application/json',
      'next-action': '4e63dabc37fef18cae74cbfd41d1bace49acf47e',
      'Cookie': authResult.cookie
    };

    const response = await fetch(NOTDIAMOND_URL, {
      method: 'POST',
      headers: headers,
      body: JSON.stringify([payload])
    });

    const generator = streamNotdiamondResponse(response, userModel);
    
    if (stream) {
      res.writeHead(200, { 
        'Content-Type': 'text/event-stream',
        'Cache-Control': 'no-cache',
        'Connection': 'keep-alive'
      });

      for await (const chunk of generator) {
        res.write(`data: ${JSON.stringify(chunk)}\n\n`);
      }
      res.write("data: [DONE]\n\n");
      res.end();
    } else {
      let fullContent = "";
      for await (const chunk of generator) {
        if (chunk.choices[0].delta.content) {
          fullContent += chunk.choices[0].delta.content;
        }
      }

      res.json({
        id: `chatcmpl-${crypto.randomUUID()}`,
        object: "chat.completion",
        created: Math.floor(Date.now() / 1000),
        model: userModel,
        system_fingerprint: "fp_4e2b2da518",
        choices: [
          {
            index: 0,
            message: {
              role: "assistant",
              content: fullContent
            },
            finish_reason: "stop"
          }
        ],
        usage: {
          prompt_tokens: Math.floor(fullContent.length / 4),
          completion_tokens: Math.floor(fullContent.length / 4),
          total_tokens: Math.floor(fullContent.length / 2)
        }
      });
    }
  } catch (error) {
    status = status === 200 ? 500 : status;
    res.status(status).json({ 
      error: 'Error',
      details: error.message
    });
  } finally {
    const endTime = new Date();
    const duration = endTime - startTime;
    console.log(`${endTime.toISOString()} - ${clientIP} - ${userModel} - ${status} - ${duration}ms - CacheHit:${cookieCacheHit} - ExpiresAt:${cookieExpiresAt}`);
  }
});

app.listen(port, () => {
  console.log(`Server running on port ${port}`);
});