import fs from 'node:fs/promises'
import { TemplateHtml, getPathByRoot, clientRoot, ssrManifest, entryServer } from "helper/path"
import { parseCookieHeader } from "helper/cookie"
import { Env } from "helper/env"
import type { ViteDevServer } from 'vite'
import Send from 'koa-send'
import type Koa from 'koa'
import c2k from 'koa-connect'
import { transformHtmlTemplate } from "unhead/server";

const base = Env.base

export async function SsrMiddleWare(app: Koa, options?: { onDevViteClose?: Function }) {
    let vite: ViteDevServer
    if (process.env.NODE_ENV !== 'production') {
        // Dev mode: create Vite server in middleware mode.
        const { createServer } = await import('vite')
        vite = await createServer({
            server: { middlewareMode: true },
            configFile: getPathByRoot('packages', 'client/vite.config.ts'),
            root: getPathByRoot('packages', 'client'),
            appType: 'custom',
            base,
        })
        app.use(c2k(vite.middlewares))
        vite.httpServer?.on("close", () => {
            vite.close()
            options?.onDevViteClose?.()
        })
    } else {
        // Production mode: serve pre-built static assets.
        app.use(async (ctx, next) => {
            try {
                await Send(ctx, ctx.path, { root: clientRoot, index: false });
                if (ctx.status === 404) {
                    await next()
                }
            } catch (error) {
                if (ctx.status === 404) {
                    await next()
                } else {
                    throw error
                }
            }
        })
    }

    // Handle every other route with SSR.
    app.use(async (ctx, next) => {
        if (!ctx.originalUrl.startsWith(base)) return await next()

        try {
            const url = ctx.originalUrl.replace(base, '')
            let template
            let render
            let manifest
            if (process.env.NODE_ENV !== 'production') {
                // Always read fresh template in development
                template = await fs.readFile(getPathByRoot('packages', 'client/index.html'), 'utf-8')
                template = await vite.transformIndexHtml(url, template)
                manifest = {}
                render = (await vite.ssrLoadModule(getPathByRoot('packages', 'client/src/entry-server.ts'))).render
            } else {
                manifest = await fs.readFile(ssrManifest, 'utf-8')
                template = TemplateHtml
                // @ts-ignore
                render = (await import(entryServer)).render
            }

            const cookies = parseCookieHeader(ctx.request.headers['cookie'] as string)

            const rendered = await render(url, manifest, { cookies })

            const html = await transformHtmlTemplate(
                rendered.unHead,
                template
                    .replace(`<!--app-head-->`, rendered.head ?? '')
                    .replace(`<!--app-html-->`, rendered.html ?? '')
            )
            ctx.status = 200
            ctx.set({ 'Content-Type': 'text/html' })
            ctx.body = html

            // 设置服务端渲染期间收集到的 Set-Cookie
            const setCookies: string[] = (rendered as any).setCookies || []
            if (setCookies.length > 0) {
                ctx.set('Set-Cookie', setCookies)
            }
        } catch (e: Error | any) {
            vite?.ssrFixStacktrace(e)
            ctx.status = 500
            console.error(e)
            ctx.body = e.stack
        }
        await next()
    })
}