import { createServerAdapter } from '@whatwg-node/server';
import { AutoRouter, json, error, cors } from 'itty-router';
import { createServer } from 'http';
import dotenv from 'dotenv';

dotenv.config();

class Config {
  constructor() {
    this.PORT = process.env.PORT || 7860;
    this.API_PREFIX = process.env.API_PREFIX || '/';
    this.MAX_RETRY_COUNT = parseInt(process.env.MAX_RETRY_COUNT) || 3;
    this.RETRY_DELAY = parseInt(process.env.RETRY_DELAY) || 5000;
    this.FAKE_HEADERS = {
      'Accept': 'application/json',
      'Content-Type': 'application/json',
      'User-Agent': 'Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/130.0.0.0 Safari/537.36',
      'Referer': 'https://duckduckgo.com/',
      'Origin': 'https://duckduckgo.com',
      'x-vqd-accept': '1',
      ...JSON.parse(process.env.FAKE_HEADERS || '{}'),
    };
  }
}

const config = new Config();

const { preflight, corsify } = cors({
  origin: '*',
  allowMethods: '*',
  exposeHeaders: '*',
});

const withBenchmarking = (request) => {
  request.start = Date.now();
};

const logger = (res, req) => {
  console.log(req.method, res.status, req.url, Date.now() - req.start, 'ms');
};

const router = AutoRouter({
  before: [withBenchmarking, preflight],
  missing: () => error(404, '404 Not Found. Please check whether the calling URL is correct.'),
  finally: [corsify, logger],
});

router.get('/', () => json({ message: 'API Service Running!' }));
router.get('/ping', () => json({ message: 'pong' }));
router.get(config.API_PREFIX + 'v1/models', () =>
  json({
    object: 'list',
    data: [
      { id: 'gpt-4o-mini', object: 'model', owned_by: 'ddg' },
      { id: 'claude-3-haiku', object: 'model', owned_by: 'ddg' },
      { id: 'llama-3.1-70b', object: 'model', owned_by: 'ddg' },
      { id: 'mixtral-8x7b', object: 'model', owned_by: 'ddg' },
      { id: 'o3-mini', object: 'model', owned_by: 'ddg' },
    ],
  })
);

router.post(config.API_PREFIX + 'v1/chat/completions', (req) => handleCompletion(req));

async function handleCompletion(request) {
  try {
    const { model: inputModel, messages, stream: returnStream } = await request.json();
    const model = convertModel(inputModel);
    const content = messagesPrepare(messages);
    return createCompletion(model, content, returnStream);
  } catch (err) {
    console.error('Handle Completion Error:', err);
    return error(500, err.message);
  }
}

async function createCompletion(model, content, returnStream, retryCount = 0) {
  try {
    const token = await requestToken();
    const response = await fetch('https://duckduckgo.com/duckchat/v1/chat', {
      method: 'POST',
      headers: {
        ...config.FAKE_HEADERS,
        'x-vqd-4': token,
      },
      body: JSON.stringify({
        model,
        messages: [{ role: 'user', content }],
      }),
    });

    if (!response.ok) {
      if (response.status === 418) {
        console.warn('Rate limit hit (418), retrying...');
        throw new Error('Rate limit exceeded');
      }
      throw new Error(`Create Completion error! status: ${response.status}, message: ${await response.text()}`);
    }

    return handlerStream(model, response.body, returnStream);
  } catch (err) {
    console.error('Create Completion Error:', err.message);
    if (retryCount < config.MAX_RETRY_COUNT && (err.message.includes('Rate limit') || err.message.includes('418'))) {
      console.log('Retrying... count', ++retryCount);
      await new Promise((resolve) => setTimeout(resolve, config.RETRY_DELAY));
      return createCompletion(model, content, returnStream, retryCount);
    }
    throw err;
  }
}

async function handlerStream(model, body, returnStream) {
  const reader = body.getReader();
  const decoder = new TextDecoder();
  const encoder = new TextEncoder();
  let previousText = '';

  const stream = new ReadableStream({
    async start(controller) {
      while (true) {
        const { done, value } = await reader.read();
        if (done) {
          if (!returnStream) {
            controller.enqueue(encoder.encode(JSON.stringify(newChatCompletionWithModel(previousText, model))));
          }
          return controller.close();
        }

        const chunk = decoder.decode(value).trim();
        const lines = chunk.split('\n').filter((line) => line.startsWith('data: '));

        for (const line of lines) {
          const data = line.slice(6);
          if (data === '[DONE]') {
            if (returnStream) {
              controller.enqueue(encoder.encode(`data: ${JSON.stringify(newStopChunkWithModel('stop', model))}\n\n`));
            }
            return controller.close();
          }

          try {
            const parsed = JSON.parse(data);
            if (parsed.message) {
              previousText += parsed.message;
              if (returnStream) {
                controller.enqueue(
                  encoder.encode(`data: ${JSON.stringify(newChatCompletionChunkWithModel(parsed.message, model))}\n\n`)
                );
              }
            }
          } catch (err) {
            console.error('Stream parse error:', err);
          }
        }
      }
    },
  });

  return new Response(stream, {
    headers: {
      'Content-Type': returnStream ? 'text/event-stream' : 'application/json',
      'Cache-Control': 'no-cache',
      'Connection': 'keep-alive',
    },
  });
}

function messagesPrepare(messages) {
  return messages
    .filter((msg) => ['user', 'assistant'].includes(msg.role))
    .map((msg) => msg.content)
    .join('\n');
}

async function requestToken() {
  try {
    const response = await fetch('https://duckduckgo.com/duckchat/v1/status', {
      method: 'GET',
      headers: config.FAKE_HEADERS,
    });
    const token = response.headers.get('x-vqd-4');
    if (!token) {
      console.error('No x-vqd-4 token found in response headers');
      throw new Error('Failed to retrieve x-vqd-4 token');
    }
    console.log('Token retrieved:', token);
    return token;
  } catch (err) {
    console.error('Request token error:', err);
    throw err;
  }
}

function convertModel(inputModel) {
  const modelMap = {
    'claude-3-haiku': 'claude-3-haiku-20240307',
    'llama-3.1-70b': 'meta-llama/Meta-Llama-3.1-70B-Instruct-Turbo',
    'mixtral-8x7b': 'mistralai/Mixtral-8x7B-Instruct-v0.1',
    'o3-mini': 'o3-mini', // Fallback to default if unsupported
  };
  const selectedModel = modelMap[inputModel.toLowerCase()] || 'gpt-4o-mini';
  console.log(`Converted model: ${inputModel} -> ${selectedModel}`);
  return selectedModel;
}

function newChatCompletionChunkWithModel(text, model) {
  return {
    id: 'chatcmpl-' + Math.random().toString(36).slice(2),
    object: 'chat.completion.chunk',
    created: Math.floor(Date.now() / 1000),
    model,
    choices: [{ index: 0, delta: { content: text }, finish_reason: null }],
  };
}

function newStopChunkWithModel(reason, model) {
  return {
    id: 'chatcmpl-' + Math.random().toString(36).slice(2),
    object: 'chat.completion.chunk',
    created: Math.floor(Date.now() / 1000),
    model,
    choices: [{ index: 0, finish_reason: reason }],
  };
}

function newChatCompletionWithModel(text, model) {
  return {
    id: 'chatcmpl-' + Math.random().toString(36).slice(2),
    object: 'chat.completion',
    created: Math.floor(Date.now() / 1000),
    model,
    usage: { prompt_tokens: 0, completion_tokens: 0, total_tokens: 0 },
    choices: [{ message: { content: text, role: 'assistant' }, index: 0 }],
  };
}

(async () => {
  if (typeof addEventListener === 'function') return;
  const ittyServer = createServerAdapter(router.fetch);
  console.log(`Listening on http://0.0.0.0:${config.PORT}`);
  const httpServer = createServer(ittyServer);
  httpServer.listen(config.PORT, '0.0.0.0');
})();