const express = require('express');
const fetch = require('node-fetch');
const bodyParser = require('body-parser');

const app = express();
app.use(bodyParser.json());

const MODEL = 'claude-3-5-sonnet@20240620';

const PROJECT_ID = process.env.PROJECT_ID;
const CLIENT_ID = process.env.CLIENT_ID;
const CLIENT_SECRET = process.env.CLIENT_SECRET;
const REFRESH_TOKEN = process.env.REFRESH_TOKEN;
const API_KEY = process.env.API_KEY;

const TOKEN_URL = 'https://www.googleapis.com/oauth2/v4/token';

let tokenCache = {
  accessToken: '',
  expiry: 0,
  refreshPromise: null
};

async function getAccessToken() {
  const now = Date.now() / 1000;

  if (tokenCache.accessToken && now < tokenCache.expiry - 120) {
    return tokenCache.accessToken;
  }

  if (tokenCache.refreshPromise) {
    await tokenCache.refreshPromise;
    return tokenCache.accessToken;
  }

  tokenCache.refreshPromise = (async () => {
    try {
      const response = await fetch(TOKEN_URL, {
        method: 'POST',
        headers: {
          'Content-Type': 'application/json'
        },
        body: JSON.stringify({
          client_id: CLIENT_ID,
          client_secret: CLIENT_SECRET,
          refresh_token: REFRESH_TOKEN,
          grant_type: 'refresh_token'
        })
      });

      const data = await response.json();
      tokenCache.accessToken = data.access_token;
      tokenCache.expiry = now + data.expires_in;
    } finally {
      tokenCache.refreshPromise = null;
    }
  })();

  await tokenCache.refreshPromise;
  return tokenCache.accessToken;
}

function getLocation() {
  const currentSeconds = new Date().getSeconds();
  return currentSeconds < 30 ? 'europe-west1' : 'us-east5';
}

function constructApiUrl(location) {
  return `https://${location}-aiplatform.googleapis.com/v1/projects/${PROJECT_ID}/locations/${location}/publishers/anthropic/models/${MODEL}:streamRawPredict`;
}

async function handleRequest(req, res) {
  if (req.method === 'OPTIONS') {
    return handleOptions(res);
  }

  const apiKey = req.headers['x-api-key'];
  if (apiKey !== API_KEY) {
    return res.status(403).json({
      type: "error",
      error: {
        type: "permission_error",
        message: "Your API key does not have permission to use the specified resource."
      }
    });
  }

  const accessToken = await getAccessToken();
  const location = getLocation();
  const apiUrl = constructApiUrl(location);

  let requestBody = req.body;

  if (requestBody.anthropic_version) {
    delete requestBody.anthropic_version;
  }

  if (requestBody.model) {
    delete requestBody.model;
  }

  requestBody.anthropic_version = "vertex-2023-10-16";

  const modifiedHeaders = {
    'Authorization': `Bearer ${accessToken}`,
    'Content-Type': 'application/json; charset=utf-8'
  };

  const response = await fetch(apiUrl, {
    headers: modifiedHeaders,
    method: 'POST',
    body: JSON.stringify(requestBody)
  });

  const responseBody = await response.text();
  res.status(response.status).set(response.headers).send(responseBody);
}

function handleOptions(res) {
  res.setHeader('Access-Control-Allow-Origin', '*');
  res.setHeader('Access-Control-Allow-Methods', 'POST, GET, OPTIONS');
  res.setHeader('Access-Control-Allow-Headers', 'Content-Type, Authorization, x-api-key, anthropic-version, model');
  res.sendStatus(204);
}

app.post('/ai/v1/messages', handleRequest);
app.options('/ai/v1/messages', handleOptions);

app.get('/', (req, res) => {
  res.status(200).send('Vertex Claude API Proxy');
});

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