// Command wasmgen preprocesses component directives into glue code so that
// server-component calls look identical to client-component calls, and routes
// can opt into SSR with a tag.
//
// Directives (as doc comments on functions in the target package):
//
//	//gowasm:page <path> [static] [layout=<name>]
//	                                a page factory `func(Deps) func() *vdom.VNode`.
//	                                `static` => the server SSRs that route;
//	                                `layout=` => wrap it in a named //gowasm:layout.
//	//gowasm:layout <name>          a layout `func(Deps, *vdom.VNode) *vdom.VNode`.
//	//gowasm:server                 a server component `func() func() *vdom.VNode`.
//
// It emits three files in the package:
//   - routes.gen.go  (neutral)  Routes(Deps) + StaticPaths + RouteLayout/LayoutFor
//   - server.gen.go  (native)   init() registering each server component
//   - client.gen.go  (wasm)     a client stub per server component that
//     mounts it over /rsc
package main

import (
	"fmt"
	"go/ast"
	"go/format"
	"go/parser"
	"go/token"
	"os"
	"path/filepath"
	"sort"
	"strings"
	"text/template"
)

type page struct {
	Func   string
	Path   string
	Static bool
	Layout string // name of the //gowasm:layout wrapper (empty => default)
}

type layoutDecl struct {
	Name string
	Func string
}

type genData struct {
	Pages         []page
	Servers       []string
	Layouts       []layoutDecl
	DefaultLayout string // func name used when a page declares no layout
}

func main() {
	dir := "app"
	if len(os.Args) > 1 {
		dir = os.Args[1]
	}

	fset := token.NewFileSet()
	pkgs, err := parser.ParseDir(fset, dir, func(fi os.FileInfo) bool {
		return !strings.HasSuffix(fi.Name(), ".gen.go")
	}, parser.ParseComments)
	if err != nil {
		fmt.Fprintln(os.Stderr, "wasmgen: parse:", err)
		os.Exit(1)
	}

	var data genData
	for _, pkg := range pkgs {
		for _, file := range pkg.Files {
			for _, decl := range file.Decls {
				fn, ok := decl.(*ast.FuncDecl)
				if !ok || fn.Doc == nil {
					continue
				}
				for _, c := range fn.Doc.List {
					line := strings.TrimSpace(strings.TrimPrefix(c.Text, "//"))
					switch {
					case strings.HasPrefix(line, "gowasm:page "):
						fields := strings.Fields(line[len("gowasm:page "):])
						if len(fields) == 0 {
							continue
						}
						p := page{Func: fn.Name.Name, Path: fields[0]}
						for _, f := range fields[1:] {
							switch {
							case f == "static":
								p.Static = true
							case strings.HasPrefix(f, "layout="):
								p.Layout = strings.TrimPrefix(f, "layout=")
							}
						}
						data.Pages = append(data.Pages, p)
					case strings.HasPrefix(line, "gowasm:layout "):
						fields := strings.Fields(line[len("gowasm:layout "):])
						if len(fields) == 0 {
							continue
						}
						data.Layouts = append(data.Layouts, layoutDecl{Name: fields[0], Func: fn.Name.Name})
					case line == "gowasm:server":
						data.Servers = append(data.Servers, fn.Name.Name)
					}
				}
			}
		}
	}
	sort.Slice(data.Pages, func(i, j int) bool { return data.Pages[i].Path < data.Pages[j].Path })
	sort.Strings(data.Servers)
	sort.Slice(data.Layouts, func(i, j int) bool { return data.Layouts[i].Name < data.Layouts[j].Name })

	// Default layout: prefer one literally named "app", else the first declared.
	var defaultName string
	if len(data.Layouts) > 0 {
		data.DefaultLayout, defaultName = data.Layouts[0].Func, data.Layouts[0].Name
		for _, l := range data.Layouts {
			if l.Name == "app" {
				data.DefaultLayout, defaultName = l.Func, l.Name
				break
			}
		}
	}
	for i := range data.Pages {
		if data.Pages[i].Layout == "" {
			data.Pages[i].Layout = defaultName
		}
	}

	write(filepath.Join(dir, "routes.gen.go"), routesTmpl, data)
	write(filepath.Join(dir, "server.gen.go"), serverTmpl, data)
	write(filepath.Join(dir, "client.gen.go"), clientTmpl, data)

	fmt.Printf("wasmgen: %d page(s), %d server component(s)\n", len(data.Pages), len(data.Servers))
}

func write(path, tmpl string, data genData) {
	var b strings.Builder
	if err := template.Must(template.New("g").Parse(tmpl)).Execute(&b, data); err != nil {
		panic(err)
	}
	src, err := format.Source([]byte(b.String()))
	if err != nil {
		fmt.Fprintf(os.Stderr, "wasmgen: format %s: %v\n---\n%s\n", path, err, b.String())
		os.Exit(1)
	}
	// Only write when changed, so re-running codegen doesn't churn mtimes (which
	// would make the dev server's file watcher loop).
	if existing, err := os.ReadFile(path); err == nil && string(existing) == string(src) {
		return
	}
	if err := os.WriteFile(path, src, 0o644); err != nil {
		panic(err)
	}
}

const routesTmpl = `// Code generated by wasmgen. DO NOT EDIT.
package app

import "kjol/wasmruntime/vdom"

// Routes maps each //gowasm:page path to its instantiated render function.
func Routes(d Deps) map[string]func() *vdom.VNode {
	return map[string]func() *vdom.VNode{
{{range .Pages}}		"{{.Path}}": {{.Func}}(d),
{{end}}	}
}

// StaticPaths are the routes the server pre-renders (SSR); others render client-side.
var StaticPaths = map[string]bool{
{{range .Pages}}{{if .Static}}	"{{.Path}}": true,
{{end}}{{end}}}

// RouteLayout maps each route to the name of the layout that wraps it.
var RouteLayout = map[string]string{
{{range .Pages}}	"{{.Path}}": "{{.Layout}}",
{{end}}}

// LayoutFor wraps a page's content in the layout declared for its route.
func LayoutFor(d Deps, path string, content *vdom.VNode) *vdom.VNode {
{{if .Layouts}}	switch RouteLayout[path] {
{{range .Layouts}}	case "{{.Name}}":
		return {{.Func}}(d, content)
{{end}}	}
	return {{.DefaultLayout}}(d, content)
{{else}}	return content
{{end}}}
`

const serverTmpl = `// Code generated by wasmgen. DO NOT EDIT.

//go:build !(js && wasm)

package app

import "kjol/wasmruntime/rsc"

func init() {
{{range .Servers}}	rsc.Register("{{.}}", {{.}})
{{end}}}
`

const clientTmpl = `// Code generated by wasmgen. DO NOT EDIT.

//go:build js && wasm

package app

import (
	"kjol/wasmruntime/rsc"
	"kjol/wasmruntime/vdom"
)

{{range .Servers}}// {{.}} is a generated client stub for the server component of the same name.
func {{.}}() func() *vdom.VNode { return rsc.Mount("{{.}}") }
{{end}}`