server.js

const express = require('express');
const { chromium } = require('playwright');

const app = express();
const port = Number(process.env.PORT) || 3000;
const defaultProvider = process.env.AI_PROVIDER || 'duck';
const providerUrls = { duck: process.env.DUCK_AI_URL || 'https://duck.ai/',  lumo: process.env.LUMO_AI_URL || 'https://lumo.proton.me/guest',
  mistral: process.env.MISTRAL_AI_URL || 'https://chat.mistral.ai/chat' };
const requestTimeoutMs = Number(process.env.PROMPT_TIMEOUT_MS) || 60_000;
const headed = process.env.HEADED_BROWSER !== 'false';
const keepBrowserOpen = process.env.KEEP_BROWSER_OPEN === 'true';
const keepBrowserOpenAfterResponse = process.env.KEEP_BROWSER_OPEN_AFTER_RESPONSE === 'true';
const allowedLanguages = new Set(['C++', 'Go', 'Java', 'Python', 'Rust', 'Swift', 'Typescript']);
const log = (message, details) => console.log(`[${new Date().toISOString()}] ${message}`, details || '');
let activeRequest = null;

app.use(express.json({ limit: '16kb' }));
app.use(express.static('public'));

async function findPromptInput(page) {
  const candidates = [page.locator('textarea[placeholder*="ask" i], textarea[placeholder*="message" i]'), page.locator('textarea'), page.locator('[contenteditable="true"]'), page.getByRole('textbox'), page.locator('input[type="text"]')];
  for (const candidate of candidates) { const input = candidate.first(); if (await input.count() && await input.isVisible().catch(() => false)) return input; }
  throw new Error('Could not find AI prompt input after the page loaded.');
}

async function waitForPromptInput(page) {
  await page.waitForLoadState('networkidle').catch(() => {});
  const deadline = Date.now() + requestTimeoutMs;
  while (Date.now() < deadline) { try { return await findPromptInput(page); } catch { await page.waitForTimeout(250); } }
  throw new Error('Could not find AI prompt input after the page loaded.');
}

async function acceptMistralTerms(page, provider) {
  if (provider !== 'mistral') return;
  const accept = page.getByRole('button', { name: /accept and continue/i });
  try {
    await accept.waitFor({ state: 'visible', timeout: 10_000 });
    log('Accepting Mistral Le Chat terms of service.');
    await accept.click();
    await page.waitForTimeout(500);
  } catch {
    log('Mistral terms dialog was not shown.');
  }
}

async function submitPrompt(page, input) {
  const buttons = [page.getByRole('button', { name: /send|submit|ask/i }), page.locator('button[aria-label*="send" i]'), page.locator('button[type="submit"]'), page.locator('form button')];
  for (const candidate of buttons) { const button = candidate.first(); if (await button.count() && await button.isVisible().catch(() => false)) { await button.click(); return; } }
  await input.press('Enter');
}

function extractCodeFromText(text, language) {
  const aliases = { 'C++': ['cpp', 'c\\+\\+'], Typescript: ['typescript', 'ts'] };
  const tags = aliases[language] || [language.toLowerCase()];
  const tagged = new RegExp('```(?:' + tags.join('|') + ')\\s*\\n([\\s\\S]*?)\\n?```', 'i').exec(text);
  if (tagged) return tagged[1].trim();
  const anyBlock = /```[^\n]*\n([\s\S]*?)\n?```/.exec(text);
  return anyBlock ? anyBlock[1].trim() : null;
}

function looksLikeCode(text, language) {
  const patterns = {
    'C++': /(?:#include|std::|\\b(?:int|void|string|class|namespace|template)\\b|[;{}])/i,
    Go: /(?:\\bpackage\\s+main\\b|\\bfunc\\s+\\w+\\s*\\(|\\b(?:fmt|log|http)\\.\\w+|:=|\\b(?:goroutine|goroutines|chan|defer|range|struct)\\b|[{}])/i,
    Java: /(?:^|[\\s(])(?:public|private|protected|class|static|void|new|package|import|extends|implements|System)(?:$|[\\s({.:])|[;{}]/i,
    Python: /(?:^|[\\s(])(?:from|import|def|class|if|elif|else|for|while|try|except|with|async|await|print|return)(?:$|[\\s(:])|^\\s*[@#]|:\\s*$/im,
    Rust: /(?:^|[\\s(])(?:fn|let|mut|struct|impl|use|pub|match|trait|enum|mod|crate|where|async|await|println)(?:$|[\\s({:])|[;{}]/i,
    Swift: /(?:^|[\\s(])(?:import|func|let|var|class|struct|enum|protocol|extension|guard|defer|print|if|else|for|while|switch|case|return|in)(?:$|[\\s({:])|[{}]/i,
    Typescript: /(?:^|[\\s(])(?:interface|type|const|let|var|function|export|import|class|enum|namespace|async|await|return)(?:$|[\\s({:])|[;{}]/i
  };
  return (patterns[language] || /[;{}]/).test(text);
}

function getCodeCandidate(text, language) {
  const fenced = extractCodeFromText(text, language);
  if (fenced) return fenced;
  if (language === 'Go' && /package\\s+main|func\\s+\\w+\\s*\\(/i.test(text)) {
    const start = text.search(/package\\s+main|import\\s+["']|func\\s+\\w+\\s*\\(/i);
    if (start >= 0) return text.slice(start).split(/\n\s*(?:This |These |The |It |Explanation:)/i)[0].trim();
  }
  const lines = text.split(/\r?\n/);
  const start = lines.findIndex((line) => looksLikeCode(line.trim(), language));
  if (start < 0) return null;
  let end = lines.length;
  for (let index = start + 1; index < lines.length; index += 1) {
    const line = lines[index].trim();
    if (/^(This |These |The |It |Each |First |Then |Finally |In this |Note:|Explanation:|Here’s|Here's)/i.test(line)) {
      end = index;
      break;
    }
  }
  const candidate = lines.slice(start, end).join('\n').trim();
  return candidate.split('\n').length >= 2 ? candidate : null;
}

async function readStableCodeResponse(page, prompt, language, provider) {
  log(`Waiting for ${provider} to finish streaming a ${language} response.`);
  const deadline = Date.now() + requestTimeoutMs;
  let latest = '';
  let stableSince = 0;

  while (Date.now() < deadline) {
    const candidates = [];
    const blocks = page.locator('pre code, pre');
    for (let index = (await blocks.count()) - 1; index >= 0; index -= 1) candidates.push(await blocks.nth(index).innerText().catch(() => ''));
    candidates.push(await page.locator('body').innerText().catch(() => ''));

    for (const text of candidates) {
      const code = getCodeCandidate(text.trim(), language);
      if (code && code !== prompt.trim() && code !== latest) { latest = code; stableSince = Date.now(); }
    }
    if (latest && Date.now() - stableSince >= 1_500) {
      log(`Playwright selected the completed ${language} code block from ${provider}.`, latest);
      return latest;
    }
    await page.waitForTimeout(300);
  }
  throw new Error(`Timed out waiting for ${provider} to finish the ${language} response.`);
}

async function askDuck(prompt, language, targetUrl, provider) {
  if (activeRequest) throw new Error('Another query is already running.');
  const browser = await chromium.launch({ headless: !headed });
  activeRequest = { browser, aborted: false };
  try {
    const page = await browser.newPage();
    page.setDefaultTimeout(requestTimeoutMs);
    await page.goto(targetUrl, { waitUntil: 'domcontentloaded', timeout: requestTimeoutMs });
    await acceptMistralTerms(page, provider);
    const input = await waitForPromptInput(page);
    const fullPrompt = `${prompt}\n\nPlease provide the answer as ${language} code in a code block.`;
    if (await input.getAttribute('contenteditable') === 'true') {
      await input.click();
      await page.keyboard.insertText(fullPrompt);
    } else {
      await input.fill(fullPrompt);
    }
    log(`Submitting one combined prompt to ${provider}.`);
    await submitPrompt(page, input);
    const output = await readStableCodeResponse(page, fullPrompt, language, provider);
    log('Server returning selected code to the web UI.', output);
    if (keepBrowserOpenAfterResponse || keepBrowserOpen) { log('Keeping browser open after returning the response.'); setTimeout(() => {}, 2 ** 31 - 1); }
    return output;
  } finally {
    if (!keepBrowserOpenAfterResponse && !keepBrowserOpen) await browser.close();
    if (activeRequest?.browser === browser) activeRequest = null;
  }
}

app.post('/api/abort', async (req, res) => {
  if (!activeRequest) return res.json({ aborted: false });
  activeRequest.aborted = true;
  await activeRequest.browser.close().catch(() => {});
  activeRequest = null;
  log('Aborted the active Playwright request.');
  res.json({ aborted: true });
});

app.post('/api/prompt', async (req, res) => {
  const prompt = typeof req.body?.prompt === 'string' ? req.body.prompt.trim() : '';
  const language = typeof req.body?.language === 'string' ? req.body.language : '';
  const provider = typeof req.body?.provider === 'string' ? req.body.provider : defaultProvider;
  if (!prompt) return res.status(400).json({ error: 'Enter a prompt first.' });
  if (!allowedLanguages.has(language)) return res.status(400).json({ error: 'Select a supported language.' });
  if (!providerUrls[provider]) return res.status(400).json({ error: 'Select a supported AI provider.' });
  try { res.json({ output: await askDuck(prompt, language, providerUrls[provider], provider) }); }
  catch (error) { console.error('AI request failed:', error); res.status(502).json({ error: error instanceof Error ? error.message : 'AI request failed.' }); }
});

app.listen(port, () => console.log(`AI code UI listening at http://localhost:${port}`));