const http = require('http');
const url = require('url');
const fetch = require('node-fetch');

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) {
  const url = `https://${location}-aiplatform.googleapis.com/v1/projects/${PROJECT_ID}/locations/${location}/publishers/anthropic/models/${MODEL}:streamRawPredict`;
  console.log('Constructed URL:', url); // 打印URL以验证其正确性
  return url;
}

async function handleRequest(request) {
  try {
    if (request.method === 'OPTIONS') {
      return handleOptions();
    }

    const apiKey = request.headers['x-api-key'];
    if (apiKey !== API_KEY) {
      return handleErrorResponse(403, {
        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 = JSON.parse(request.body);
    
    if (requestBody.anthropic_version) {
      delete requestBody.anthropic_version;
    }
    
    if (requestBody.model) {
      delete requestBody.model;
    }
    
    requestBody.anthropic_version = "vertex-2023-10-16";

    const modifiedHeaders = {
      ...request.headers,
      'Authorization': `Bearer ${accessToken}`,
      'Content-Type': 'application/json; charset=utf-8'
    };
    delete modifiedHeaders['anthropic-version'];

    const modifiedRequest = {
      headers: modifiedHeaders,
      method: request.method,
      body: JSON.stringify(requestBody),
      redirect: 'manual'
    };

    const response = await fetch(apiUrl, modifiedRequest);
    if (response.status === 301 || response.status === 302) {
      const location = response.headers.get('location');
      console.log('Redirected to:', location);
      const redirectedResponse = await fetch(location, modifiedRequest);
      return handleResponse(redirectedResponse);
    }
    return handleResponse(response);
  } catch (error) {
    console.error('Error handling request:', error);
    return handleErrorResponse(500, {
      type: "error",
      error: {
        type: "internal_error",
        message: "Internal server error."
      }
    });
  }
}

function handleResponse(response) {
  return {
    status: response.status,
    headers: {
      ...response.headers,
      'Access-Control-Allow-Origin': '*',
      'Access-Control-Allow-Methods': 'POST, GET, OPTIONS',
      'Access-Control-Allow-Headers': 'Content-Type, Authorization, x-api-key, anthropic-version, model'
    },
    body: response.body
  };
}

function handleErrorResponse(status, body) {
  return {
    status: status,
    headers: {
      'Content-Type': 'application/json',
      'Access-Control-Allow-Origin': '*',
      'Access-Control-Allow-Methods': 'POST, GET, OPTIONS, DELETE, HEAD',
      'Access-Control-Allow-Headers': 'Content-Type, Authorization, x-api-key, anthropic-version, model'
    },
    body: JSON.stringify(body)
  };
}

function handleOptions() {
  const headers = {
    'Access-Control-Allow-Origin': '*',
    'Access-Control-Allow-Methods': 'POST, GET, OPTIONS',
    'Access-Control-Allow-Headers': 'Content-Type, Authorization, x-api-key, anthropic-version, model'
  };

  return {
    status: 204,
    headers: headers
  };
}

const server = http.createServer(async (req, res) => {
  const parsedUrl = url.parse(req.url, true);
  if (parsedUrl.pathname === '/ai/v1/messages') {
    const request = {
      method: req.method,
      headers: req.headers,
      body: await getRequestBody(req)
    };

    const response = await handleRequest(request);
    res.writeHead(response.status, response.headers);
    res.end(response.body);
  } else {
    res.writeHead(404, { 'Content-Type': 'application/json' });
    res.end(JSON.stringify({ error: 'Not Found' }));
  }
});

function getRequestBody(req) {
  return new Promise((resolve, reject) => {
    let body = '';
    req.on('data', chunk => {
      body += chunk.toString();
    });
    req.on('end', () => {
      resolve(body);
    });
    req.on('error', err => {
      reject(err);
    });
  });
}

server.listen(8080, () => {
  console.log('Server is running on port 8080');
});