import { Hono } from 'hono';
import { serve } from '@hono/node-server';
import { cors } from 'hono/cors';
import * as playwright from 'playwright';
import { exec } from 'child_process';
import { promisify } from 'util';

const execAsync = promisify(exec);

// Config
const MAX_RETRY_COUNT = 3;
const RETRY_DELAY = 5000;

// Fake headers
const FAKE_HEADERS: Record<string, string> = {
  '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',
};

// Browser instance
let browser: playwright.Browser | null = null;

// Ensure browser is installed and launched
async function initBrowser(): Promise<playwright.Browser> {
  if (!browser) {
    try {
      // First try to launch
      browser = await playwright.chromium.launch({
        headless: true,
        args: ['--no-sandbox', '--disable-setuid-sandbox', '--disable-dev-shm-usage'],
      });
    } catch (error) {
      console.log('Installing browsers first...');
      // If that fails, install the browsers
      try {
        await execAsync('npx playwright install chromium', { stdio: 'inherit' });
        // Try launching again after installation
        browser = await playwright.chromium.launch({
          headless: true,
          args: ['--no-sandbox', '--disable-setuid-sandbox', '--disable-dev-shm-usage'],
        });
      } catch (installError) {
        console.error('Failed to install or launch browser:', installError);
        throw installError;
      }
    }
  }
  return browser;
}

// Clean headers by removing common headers that may be set by a proxy
function cleanHeaders(headers: Record<string, string>): Record<string, string> {
  const cleaned = { ...headers };
  delete cleaned['host'];
  delete cleaned['connection'];
  delete cleaned['content-length'];
  delete cleaned['accept-encoding'];
  delete cleaned['cdn-loop'];
  delete cleaned['cf-connecting-ip'];
  delete cleaned['cf-connecting-o2o'];
  delete cleaned['cf-ew-via'];
  delete cleaned['cf-ray'];
  delete cleaned['cf-visitor'];
  delete cleaned['cf-worker'];
  delete cleaned['x-direct-url'];
  delete cleaned['x-forwarded-for'];
  delete cleaned['x-forwarded-port'];
  delete cleaned['x-forwarded-proto'];
  return cleaned;
}

// Get token by sending a request with Playwright
async function playwrightRequestToken(clientIp?: string): Promise<string | undefined> {
  const browser = await initBrowser();
  const context = await browser.newContext();
  const page = await context.newPage();

  try {
    const url = 'https://duckduckgo.com/duckchat/v1/status';
    
    // Set up route handler to modify headers when intercepting the token request
    await page.route('**/*', async (route) => {
      if (route.request().url() === url) {
        const headers = cleanHeaders({
          ...FAKE_HEADERS,
          ...(clientIp ? { 'x-forwarded-for': clientIp } : {}),
        });
        await route.continue({
          headers,
        });
      } else {
        await route.continue();
      }
    });

    const response = await page.goto(url, {
      waitUntil: 'domcontentloaded',
      timeout: 30000,
    });

    if (!response || !response.ok()) {
      throw new Error('Request Token failed!');
    }

    const responseBody = await response.json();
    await context.close();
    return responseBody?.vqd;
  } catch (error) {
    await context.close();
    console.error('Request Token Error:', error);
    throw error;
  }
}

// Response interface for the completions endpoint
interface ApiResponse {
  status: number;
  headers: Record<string, string>;
  body: string | Buffer;
}

// Main Playwright function to create a chat completion
async function playwrightCreateCompletion(
  model: string,
  content: string,
  returnStream: boolean,
  clientIp?: string,
  retryCount = 0
): Promise<ApiResponse> {
  const browser = await initBrowser();
  const context = await browser.newContext();
  const page = await context.newPage();

  try {
    const token = await playwrightRequestToken(clientIp);
    const url = 'https://duckduckgo.com/duckchat/v1/chat';
    const body = JSON.stringify({
      model,
      messages: [{ role: 'user', content }],
    });

    let status: number;

    // Intercept the chat request to add token and proper headers.
    await page.route('**/*', async (route) => {
      if (route.request().url() === url) {
        const headers = cleanHeaders({
          ...FAKE_HEADERS,
          'x-vqd-4': token!,
          ...(clientIp ? { 'x-forwarded-for': clientIp } : {}),
        });
        await route.continue({
          method: 'POST',
          headers,
          postData: body,
        });
      } else {
        await route.continue();
      }
    });

    const response = await page.goto(url, {
      waitUntil: 'domcontentloaded',
      timeout: 30000,
    });

    if (!response) {
      throw new Error('No response received');
    }

    status = response.status();

    if (!response.ok()) {
      if (status === 418 && retryCount < MAX_RETRY_COUNT) {
        console.warn('Rate limit hit (418), retrying...');
        await context.close();
        await new Promise((resolve) => setTimeout(resolve, RETRY_DELAY));
        return playwrightCreateCompletion(model, content, returnStream, clientIp, retryCount + 1);
      }
      throw new Error(`Create Completion error! status: ${status}`);
    }

    const responseBody = await response.body();
    await context.close();

    if (!returnStream) {
      const text = responseBody.toString();
      return {
        status,
        headers: { 'content-type': 'application/json' },
        body: JSON.stringify({
          id: 'chatcmpl-' + Date.now(),
          object: 'chat.completion',
          created: Math.floor(Date.now() / 1000),
          model,
          choices: [
            {
              index: 0,
              message: { role: 'assistant', content: text },
              finish_reason: 'stop',
            },
          ],
        }),
      };
    }

    return {
      status,
      headers: {
        'content-type': 'text/event-stream',
        'cache-control': 'no-cache',
        'connection': 'keep-alive',
      },
      body: responseBody,
    };
  } catch (error: any) {
    await context.close();
    console.error('Create Completion Error:', error.message);
    throw error;
  }
}

// Message interface for incoming chat message objects
interface Message {
  role: 'user' | 'assistant';
  content: string;
}

// Hono Server Setup
const app = new Hono();

// Use CORS middleware
app.use('*', cors({ origin: '*', allowMethods: ['GET', 'POST', 'OPTIONS'] }));

// Root route to check if API is running
app.get('/', (c) => c.json({ message: 'API Service Running!' }));

// Ping route for health check
app.get('/ping', (c) => c.json({ message: 'pong' }));

// Models route to return the list of available models
app.get('/v1/models', (c) =>
  c.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' },
    ],
  })
);

// Chat completions route
app.post('/v1/chat/completions', async (c) => {
  try {
    const { model, messages, stream } = await c.req.json();
    const messageList = messages as Message[];
    const content = messageList
      .filter((msg: Message) => ['user', 'assistant'].includes(msg.role))
      .map((msg: Message) => msg.content)
      .join('\n');

    const clientIp = c.req.header('x-forwarded-for') || c.req.header('x-real-ip') || undefined;

    const result = await playwrightCreateCompletion(model, content, stream, clientIp);

    return new Response(result.body, {
      status: result.status,
      headers: result.headers,
    });
  } catch (error: any) {
    console.error('Error:', error);
    return c.json({ error: error.message }, 500);
  }
});

// Cleanup when the process ends
process.on('SIGINT', async () => {
  if (browser) await browser.close();
  process.exit(0);
});
process.on('SIGTERM', async () => {
  if (browser) await browser.close();
  process.exit(0);
});

// Start the server
const port = 7860;
serve({ fetch: app.fetch, port }, () => {
  console.log(`Server running at http://localhost:${port}`);
});