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 getClientIP = (req) => {
  return req.headers['x-forwarded-for']?.split(',')[0]?.trim() || 
         req.headers['x-real-ip'] || 
         req.socket?.remoteAddress || 
         null;
};

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('/', (req) => json({ message: 'API Service Running!' }));
router.get('/ping', (req) => json({ message: 'pong' }));

router.get(config.API_PREFIX + 'v1/models', (req) =>
  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);
    const clientIp = getClientIP(request);

    return createCompletion(model, content, returnStream, clientIp);
  } catch (err) {
    console.error('Handle Completion Error:', err);
    return error(500, err.message);
  }
}

async function createCompletion(model, content, returnStream, clientIp, retryCount = 0) {
  try {
    const token = await requestToken(clientIp);
    const response = await fetch('https://duckduckgo.com/duckchat/v1/chat', {
      method: 'POST',
      headers: {
        ...config.FAKE_HEADERS,
        'x-vqd-4': token,
        ...(clientIp ? { 'x-forwarded-for': clientIp } : {}),
      },
      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, clientIp, 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(clientIp) {
  try {
    const response = await fetch('https://duckduckgo.com/duckchat/v1/status', {
      method: 'GET',
      headers: {
        ...config.FAKE_HEADERS,
        ...(clientIp ? { 'x-forwarded-for': clientIp } : {}),
      },
    });
    if (!response.ok) {
      throw new Error('Request Token failed!');
    }
    const data = await response.json();
    return data?.vqd;
  } catch (err) {
    console.error('Request Token Error:', err);
    throw err;
  }
}

function convertModel(model) {
  return model; // Adjust if needed
}

function newChatCompletionWithModel(content, model) {
  return {
    id: 'chatcmpl-' + Date.now(),
    object: 'chat.completion',
    created: Math.floor(Date.now() / 1000),
    model,
    choices: [
      {
        index: 0,
        message: { role: 'assistant', content },
        finish_reason: 'stop',
      },
    ],
  };
}

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

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

// Create Server
createServer(createServerAdapter(router.fetch)).listen(config.PORT, () => {
  console.log('Server running at http://localhost:' + config.PORT);
});