// app.js
const express = require('express');
const fetch = require('node-fetch');
const app = express();
const port = 8080;

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`;
}

app.use(express.json());

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

app.post('/ai/v1/messages', async (req, 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 === "claude-3-5-sonnet-20240620") {
    requestBody.model = "claude-3-5-sonnet@20240620";
  }
  
  requestBody.anthropic_version = "vertex-2023-10-16";

  try {
    const response = await fetch(apiUrl, {
      method: 'POST',
      headers: {
        'Authorization': `Bearer ${accessToken}`,
        'Content-Type': 'application/json; charset=utf-8'
      },
      body: JSON.stringify(requestBody)
    });

    const data = await response.text();
    res.status(response.status).send(data);
  } catch (error) {
    console.error('Error:', error);
    res.status(500).json({ error: 'Internal Server Error' });
  }
});

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