#!/usr/bin/env bun

import { mkdir } from 'node:fs/promises'
import { join } from 'node:path'
import { type BrowserContext, chromium, type Page } from 'playwright'
import { db } from '../src/lib/db'

const SITE_IMAGES_DIR = join(process.cwd(), 'site-images')
const VIEWPORT_WIDTH = 1920
const VIEWPORT_HEIGHT = 1080
const TIMEOUT = 10_000
const MAX_RETRIES = 1
const CONCURRENCY = 10

const COLOR_SCHEMES = ['light', 'dark'] as const
type ColorScheme = (typeof COLOR_SCHEMES)[number]

interface DomainBackedSite {
	did: string
	rkey: string
	url: string
	domain: string
	domainType: 'custom' | 'wisp'
}

interface ScreenshotResult {
	success: boolean
	error?: string
}

interface DomainBackedSiteRow {
	did: string
	rkey: string
	domain: string
	domain_type: DomainBackedSite['domainType']
}

async function getDomainBackedSites(): Promise<DomainBackedSite[]> {
	const rows = await db`
		SELECT
			s.did,
			s.rkey,
			COALESCE(cd.domain, d.domain) AS domain,
			CASE
				WHEN cd.domain IS NOT NULL THEN 'custom'
				ELSE 'wisp'
			END AS domain_type
		FROM sites s
		LEFT JOIN LATERAL (
			SELECT domain
			FROM custom_domains
			WHERE did = s.did
				AND rkey = s.rkey
				AND verified = true
			ORDER BY created_at ASC
			LIMIT 1
		) cd ON true
		LEFT JOIN LATERAL (
			SELECT domain
			FROM domains
			WHERE did = s.did
				AND rkey = s.rkey
			ORDER BY created_at ASC
			LIMIT 1
		) d ON cd.domain IS NULL
		WHERE cd.domain IS NOT NULL OR d.domain IS NOT NULL
		ORDER BY s.created_at DESC
	`

	return (rows as DomainBackedSiteRow[]).map((row) => ({
		did: row.did as string,
		rkey: row.rkey as string,
		domain: row.domain as string,
		domainType: row.domain_type as DomainBackedSite['domainType'],
		url: `https://${row.domain}`,
	}))
}

function sanitizeFilename(value: string): string {
	return value.replace(/[^a-z0-9_-]/gi, '_').toLowerCase()
}

async function screenshotSite(
	page: Page,
	site: DomainBackedSite,
	scheme: ColorScheme,
	retries: number = MAX_RETRIES,
): Promise<ScreenshotResult> {
	const filename = `${sanitizeFilename(site.domain)}.png`
	const filepath = join(SITE_IMAGES_DIR, scheme, filename)

	for (let attempt = 0; attempt <= retries; attempt++) {
		try {
			await page.goto(site.url, {
				waitUntil: 'networkidle',
				timeout: TIMEOUT,
			})

			await page.waitForTimeout(1000)

			await page.screenshot({
				path: filepath,
				fullPage: false,
				type: 'png',
			})

			return { success: true }
		} catch (error) {
			if (attempt < retries) {
				continue
			}

			return {
				success: false,
				error: error instanceof Error ? error.message : String(error),
			}
		}
	}

	return { success: false, error: 'Unknown error' }
}

async function main() {
	console.log('Starting site image poller')
	for (const scheme of COLOR_SCHEMES) {
		await mkdir(join(SITE_IMAGES_DIR, scheme), { recursive: true })
	}
	console.log(`Saving screenshots to ${SITE_IMAGES_DIR}/{light,dark}`)

	const sites = await getDomainBackedSites()
	console.log(`Found ${sites.length} sites with custom domains or wisp subdomains`)

	if (sites.length === 0) {
		return
	}

	const browser = await chromium.launch({
		headless: true,
	})

	const contexts: Record<ColorScheme, BrowserContext> = {
		light: await browser.newContext({
			viewport: { width: VIEWPORT_WIDTH, height: VIEWPORT_HEIGHT },
			colorScheme: 'light',
			userAgent:
				'Mozilla/5.0 (Macintosh; Intel Mac OS X 10_15_7) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/120.0.0.0 Safari/537.36 WispSiteImageBot/1.0',
		}),
		dark: await browser.newContext({
			viewport: { width: VIEWPORT_WIDTH, height: VIEWPORT_HEIGHT },
			colorScheme: 'dark',
			userAgent:
				'Mozilla/5.0 (Macintosh; Intel Mac OS X 10_15_7) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/120.0.0.0 Safari/537.36 WispSiteImageBot/1.0',
		}),
	}

	const results = {
		success: 0,
		failed: 0,
		errors: [] as Array<{ site: string; error: string }>,
	}

	for (let i = 0; i < sites.length; i += CONCURRENCY) {
		const batch = sites.slice(i, i + CONCURRENCY)
		const batchNum = Math.floor(i / CONCURRENCY) + 1
		const totalBatches = Math.ceil(sites.length / CONCURRENCY)

		console.log(`Batch ${batchNum}/${totalBatches}: ${batch.length} sites × 2 schemes`)

		const batchResults = await Promise.all(
			batch.flatMap((site, index) =>
				COLOR_SCHEMES.map(async (scheme) => {
					const page = await contexts[scheme].newPage()
					const globalIndex = i + index + 1
					console.log(`  [${globalIndex}/${sites.length}] ${site.url} (${site.domainType}, ${scheme})`)

					const result = await screenshotSite(page, site, scheme)
					await page.close()

					return { site, scheme, result }
				}),
			),
		)

		for (const { site, scheme, result } of batchResults) {
			if (result.success) {
				results.success++
				continue
			}

			results.failed++
			results.errors.push({
				site: `${site.did}/${site.rkey} (${site.domain}, ${scheme})`,
				error: result.error || 'Unknown error',
			})
		}
	}

	await browser.close()

	console.log(`Successful: ${results.success}`)
	console.log(`Failed: ${results.failed}`)

	if (results.errors.length > 0) {
		console.log('Failed sites:')
		for (const { site, error } of results.errors) {
			console.log(`  - ${site}: ${error}`)
		}
	}
}

main().catch((error) => {
	console.error('Fatal error:', error)
	process.exit(1)
})