const fs = require('fs')
const express = require('express')
const { promisify } = require('node:util')
const { exec } = require('node:child_process')
const puppeteer = require('puppeteer-core')
const freeport = require('freeport')
const ProxyChain = require('proxy-chain')

const app = express()

const freePortPromise = promisify(freeport)
const execPromise = promisify(exec)
const COOKIE_PATH = `${process.cwd()}/cookies.json`

async function delay(ms) {
    return new Promise(resolve => setTimeout(resolve, ms))
}

async function getChromiumPath() {
    if (process.env.PUPPETEER_EXECUTABLE_PATH) {
        return process.env.PUPPETEER_EXECUTABLE_PATH
    }
    const { stdout } = await execPromise('which chromium || which chromium-browser || which google-chrome')
    if (!stdout.trim()) throw new Error('Chromium not found')
    return stdout.trim()
}

async function waitForStableReply(page) {
    let previousText = ''
    let stableCount = 0

    while (stableCount < 3) {
        await delay(1000)

        const currentText = await page.evaluate(() => {
            const blocks = document.querySelectorAll('.break-anywhere')
            if (!blocks.length) return ''

            const last = blocks[blocks.length - 1]
            const parts = Array.from(
                last.querySelectorAll('.whitespace-pre-wrap, pre')
            )

            return parts.map(el => el.innerText.trim()).join('\n\n')
        })

        if (currentText && currentText === previousText) {
            stableCount++
        } else {
            stableCount = 0
        }

        previousText = currentText
    }

    return previousText
}

async function runPi(message) {
    let browser
    let proxyServer

    try {
        const port = await freePortPromise()
        proxyServer = new ProxyChain.Server({ port })
        await new Promise(resolve => proxyServer.listen(resolve))

        const chromiumPath = await getChromiumPath()

        browser = await puppeteer.launch({
            headless: true,
            executablePath: chromiumPath,
            args: [
                '--no-sandbox',
                '--disable-dev-shm-usage',
                `--proxy-server=127.0.0.1:${port}`
            ]
        })

        const page = await browser.newPage()

        if (fs.existsSync(COOKIE_PATH)) {
            const cookies = JSON.parse(fs.readFileSync(COOKIE_PATH, 'utf8'))
            await page.setCookie(...cookies)
        }

        await page.goto('https://pi.ai/talk', { waitUntil: 'networkidle2' })
        await page.waitForSelector('#composer-field textarea')

        await page.type('#composer-field textarea', message, { delay: 5 })
        await page.keyboard.press('Enter')

        await page.waitForSelector('.break-anywhere', { timeout: 30000 })

        const aiReply = await waitForStableReply(page)

        const updatedCookies = await page.cookies()
        fs.writeFileSync(COOKIE_PATH, JSON.stringify(updatedCookies, null, 2))

        return aiReply

    } finally {
        if (browser) await browser.close()
        if (proxyServer) await proxyServer.close(true)
    }
}

app.get('/chat', async (req, res) => {
    try {
        const message = req.query.message
        if (!message) {
            return res.status(400).json({ error: 'message query required' })
        }

        const reply = await runPi(message)

        res.json({
            success: true,
            reply
        })

    } catch (err) {
        console.error(err)
        res.status(500).json({
            success: false,
            error: err.message
        })
    }
})

const PORT = process.env.PORT || 7860

app.listen(PORT, '0.0.0.0', () => {
    console.log(`Running on ${PORT}`)
})