iris/host/devserver.go

package host

import (
	"bytes"
	"compress/gzip"
	"fmt"
	"io"
	"log"
	"net/http"
	"os"
	"path/filepath"
	"strings"
	"time"
)

// DevServer provides a development server with hot reload support.
// It watches source files, rebuilds WASM on changes, and notifies
// connected browsers to reload.
type DevServer struct {
	publicDir  string
	indexFile  string
	sourceDir  string
	outputPath string

	watcher    *Watcher
	builder    *Builder
	liveReload *LiveReload
	logger     *log.Logger
}

// DevServerOption configures a DevServer.
type DevServerOption func(*DevServer)

// WithLogger sets a custom logger for the dev server.
func WithLogger(logger *log.Logger) DevServerOption {
	return func(d *DevServer) {
		d.logger = logger
	}
}

// NewDevServer creates a new development server with hot reload.
//
// Parameters:
//   - publicDir: Directory containing static files (e.g., "public")
//   - indexFile: The fallback HTML file (e.g., "index.html")
//   - sourceDir: Directory containing Go source files to watch (e.g., ".")
//   - outputPath: Target path for compiled WASM file (e.g., "public/app.wasm")
func NewDevServer(publicDir, indexFile, sourceDir, outputPath string, opts ...DevServerOption) *DevServer {
	d := &DevServer{
		publicDir:  publicDir,
		indexFile:  indexFile,
		sourceDir:  sourceDir,
		outputPath: outputPath,
		logger:     log.New(os.Stdout, "[iris] ", log.LstdFlags),
	}

	for _, opt := range opts {
		opt(d)
	}

	d.liveReload = NewLiveReload()
	d.builder = NewBuilder(sourceDir, outputPath)

	d.watcher = NewWatcher(sourceDir, d.onFileChange,
		WithExtensions(".go"),
		WithInterval(500*time.Millisecond),
	)

	return d
}

// Start begins watching for file changes.
func (d *DevServer) Start() error {
	// Initial build
	d.logger.Println("Performing initial build...")
	result := d.builder.Build()
	if !result.Success {
		d.logger.Printf("Initial build failed: %v\n%s", result.Error, result.Output)
	} else {
		d.logger.Println("Initial build succeeded")
	}

	// Start watcher
	d.logger.Printf("Watching %s for changes...", d.sourceDir)
	return d.watcher.Start()
}

// Stop stops the file watcher.
func (d *DevServer) Stop() {
	d.watcher.Stop()
}

// ServeHTTP implements http.Handler interface.
func (d *DevServer) ServeHTTP(w http.ResponseWriter, r *http.Request) {
	// Handle live reload WebSocket endpoint
	if r.URL.Path == "/__livereload" {
		d.liveReload.ServeHTTP(w, r)
		return
	}

	p := filepath.Join(d.publicDir, filepath.Clean(r.URL.Path))

	if info, err := os.Stat(p); err != nil {
		d.serveHTMLWithInjection(w, r, filepath.Join(d.publicDir, d.indexFile))
		return
	} else if info.IsDir() {
		d.serveHTMLWithInjection(w, r, filepath.Join(d.publicDir, d.indexFile))
		return
	}

	// Check if this is an HTML file
	if strings.HasSuffix(p, ".html") {
		d.serveHTMLWithInjection(w, r, p)
		return
	}

	// Serve other files with compression
	d.serveFileWithCompression(w, r, p)
}

func (d *DevServer) onFileChange() {
	d.logger.Println("File change detected, rebuilding...")

	result := d.builder.Build()
	if !result.Success {
		d.logger.Printf("Build failed: %v\n%s", result.Error, result.Output)
		return
	}

	d.logger.Println("Build succeeded, reloading browsers...")
	d.liveReload.NotifyReload()
}

func (d *DevServer) serveHTMLWithInjection(w http.ResponseWriter, r *http.Request, filePath string) {
	// Read HTML file
	content, err := os.ReadFile(filePath)
	if err != nil {
		http.Error(w, "File not found", http.StatusNotFound)
		return
	}

	// Inject reload script
	content = InjectReloadScript(content)

	// Set headers
	w.Header().Set("Content-Type", "text/html; charset=utf-8")
	w.Header().Set("Cache-Control", "no-cache, no-store, must-revalidate")

	// Check if client accepts gzip
	if strings.Contains(r.Header.Get("Accept-Encoding"), "gzip") {
		w.Header().Set("Content-Encoding", "gzip")
		w.Header().Set("Vary", "Accept-Encoding")

		gz := gzip.NewWriter(w)
		defer gz.Close()

		gz.Write(content)
	} else {
		w.Write(content)
	}
}

func (d *DevServer) serveFileWithCompression(w http.ResponseWriter, r *http.Request, filePath string) {
	// Check if client accepts gzip
	if !strings.Contains(r.Header.Get("Accept-Encoding"), "gzip") {
		http.ServeFile(w, r, filePath)
		return
	}

	// Open and read file
	file, err := os.Open(filePath)
	if err != nil {
		http.Error(w, "File not found", http.StatusNotFound)
		return
	}
	defer file.Close()

	// Get file info for headers
	fileInfo, err := file.Stat()
	if err != nil {
		http.Error(w, "Error reading file", http.StatusInternalServerError)
		return
	}

	// Set content type based on file extension
	contentType := getContentType(filepath.Ext(filePath))
	w.Header().Set("Content-Type", contentType)

	// Disable caching in dev mode for WASM files
	if strings.HasSuffix(filePath, ".wasm") {
		w.Header().Set("Cache-Control", "no-cache, no-store, must-revalidate")
	}

	// Check if we should compress this file type
	if shouldCompress(contentType) {
		// Set gzip headers
		w.Header().Set("Content-Encoding", "gzip")
		w.Header().Set("Vary", "Accept-Encoding")

		// Create gzip writer
		gz := gzip.NewWriter(w)
		defer gz.Close()

		// Copy file content through gzip
		_, err = io.Copy(gz, file)
		if err != nil {
			return
		}
	} else {
		// Serve without compression
		http.ServeContent(w, r, fileInfo.Name(), fileInfo.ModTime(), file)
	}
}

// ListenAndServe starts the development server on the specified address.
func (d *DevServer) ListenAndServe(addr string) error {
	if err := d.Start(); err != nil {
		return fmt.Errorf("failed to start watcher: %w", err)
	}

	d.logger.Printf("Development server running on http://%s", addr)
	return http.ListenAndServe(addr, d)
}

// ResponseWriter wraps http.ResponseWriter to capture the response
type responseWriter struct {
	http.ResponseWriter
	buf        *bytes.Buffer
	statusCode int
}

func (rw *responseWriter) WriteHeader(code int) {
	rw.statusCode = code
}

func (rw *responseWriter) Write(b []byte) (int, error) {
	return rw.buf.Write(b)
}