package caddysnake

import (
        "net/http"
        "os"
        "path/filepath"
        "sync"
        "time"

        "github.com/fsnotify/fsnotify"
        "go.uber.org/zap"
)

// watchDirRecursive adds all directories under root to the fsnotify watcher.
// It is used by both AutoreloadableApp and DynamicApp.
func watchDirRecursive(watcher *fsnotify.Watcher, root string, logger *zap.Logger) {
        filepath.Walk(root, func(path string, info os.FileInfo, err error) error {
                if err != nil {
                        return err
                }
                if !info.IsDir() {
                        return nil
                }
                if addErr := watcher.Add(path); addErr != nil {
                        logger.Warn("autoreload: failed to watch directory",
                                zap.String("path", path),
                                zap.Error(addErr),
                        )
                }
                return nil
        })
}

// isPythonFileEvent returns true if the event is a write/create/remove/rename
// of a .py file.
func isPythonFileEvent(event fsnotify.Event) bool {
        if filepath.Ext(event.Name) != ".py" {
                return false
        }
        return event.Has(fsnotify.Write) || event.Has(fsnotify.Create) ||
                event.Has(fsnotify.Remove) || event.Has(fsnotify.Rename)
}

// handleNewDirEvent checks if the event is a newly created directory and adds
// it to the watcher if appropriate.
func handleNewDirEvent(event fsnotify.Event, watcher *fsnotify.Watcher) {
        if !event.Has(fsnotify.Create) {
                return
        }
        info, err := os.Stat(event.Name)
        if err != nil || !info.IsDir() {
                return
        }
        watcher.Add(event.Name)
}

// AutoreloadableApp wraps an AppServer to support hot-reloading when Python
// files in the working directory change. It watches for .py file modifications
// and reloads the app after a debounce period to group rapid changes.
type AutoreloadableApp struct {
        mu                   sync.RWMutex
        app                  AppServer
        factory              func() (AppServer, error)
        watcher              *fsnotify.Watcher
        stopCh               chan struct{}
        logger               *zap.Logger
        workingDir           string
        exitOnReloadFailure   func(code int) // if set, process exits on reload failure instead of serving 500
}

// NewAutoreloadableApp creates an AutoreloadableApp that wraps the given app and
// starts a filesystem watcher on the working directory. When any .py file changes,
// the app is reloaded after a 500ms debounce window.
// If exitOnReloadFailure is non-nil, it is called with code 1 when a reload fails
// (e.g. app deleted), so the process can terminate and stop serving requests.
func NewAutoreloadableApp(
        app AppServer,
        workingDir string,
        factory func() (AppServer, error),
        logger *zap.Logger,
        exitOnReloadFailure func(code int),
) (*AutoreloadableApp, error) {
        watcher, err := fsnotify.NewWatcher()
        if err != nil {
                return nil, err
        }

        a := &AutoreloadableApp{
                app:                 app,
                factory:             factory,
                watcher:             watcher,
                stopCh:              make(chan struct{}),
                logger:              logger,
                workingDir:          workingDir,
                exitOnReloadFailure: exitOnReloadFailure,
        }

        watchDirRecursive(watcher, workingDir, logger)

        go a.watch()

        logger.Info("autoreload enabled", zap.String("working_dir", workingDir))

        return a, nil
}

// watch runs in a goroutine and listens for filesystem events.
// It debounces rapid changes (e.g. editor save + format) into a single reload.
func (a *AutoreloadableApp) watch() {
        var debounceTimer *time.Timer
        const debounceDuration = 500 * time.Millisecond

        for {
                select {
                case event, ok := <-a.watcher.Events:
                        if !ok {
                                return
                        }
                        if !isPythonFileEvent(event) {
                                handleNewDirEvent(event, a.watcher)
                                continue
                        }
                        a.logger.Debug("python file changed",
                                zap.String("file", event.Name),
                                zap.String("op", event.Op.String()),
                        )
                        if debounceTimer != nil {
                                debounceTimer.Stop()
                        }
                        debounceTimer = time.AfterFunc(debounceDuration, func() {
                                a.reload()
                        })
                case err, ok := <-a.watcher.Errors:
                        if !ok {
                                return
                        }
                        a.logger.Error("autoreload watcher error", zap.Error(err))
                case <-a.stopCh:
                        if debounceTimer != nil {
                                debounceTimer.Stop()
                        }
                        return
                }
        }
}

// reload performs the actual app reload by stopping the old worker processes
// and starting new ones via the factory function.
func (a *AutoreloadableApp) reload() {
        a.logger.Info("reloading python app due to file changes")

        // Create new app OUTSIDE lock to avoid blocking requests
        newApp, err := a.factory()
        if err != nil {
                a.logger.Error("failed to reload python app", zap.Error(err))
                if a.exitOnReloadFailure != nil {
                        a.exitOnReloadFailure(1)
                }
                a.mu.Lock()
                a.app = &errorApp{err: err}
                a.mu.Unlock()
                return
        }

        // Swap under lock (fast operation)
        a.mu.Lock()
        oldApp := a.app
        a.app = newApp
        a.mu.Unlock()

        a.logger.Info("python app reloaded successfully")

        // Cleanup old app OUTSIDE lock. The write lock above guarantees all
        // in-flight requests using oldApp have completed before the swap.
        if err := oldApp.Cleanup(); err != nil {
                a.logger.Error("failed to cleanup old python app during reload", zap.Error(err))
        }
}

// HandleRequest forwards the request to the underlying app while holding a read
// lock. This ensures the app isn't swapped mid-request.
func (a *AutoreloadableApp) HandleRequest(w http.ResponseWriter, r *http.Request) error {
        a.mu.RLock()
        defer a.mu.RUnlock()
        return a.app.HandleRequest(w, r)
}

// Cleanup stops the filesystem watcher and cleans up the underlying app.
func (a *AutoreloadableApp) Cleanup() error {
        close(a.stopCh)
        a.watcher.Close()
        a.mu.RLock()
        app := a.app
        a.mu.RUnlock()
        return app.Cleanup()
}

// errorApp is a stub AppServer returned when a reload fails.
// It returns HTTP 503 for all requests until the next successful reload.
type errorApp struct {
        err error
}

func (e *errorApp) HandleRequest(w http.ResponseWriter, r *http.Request) error {
        w.WriteHeader(http.StatusServiceUnavailable)
        w.Write([]byte("Service temporarily unavailable"))
        return nil
}

func (e *errorApp) Cleanup() error {
        return nil
}

// Caddy plugin to serve Python apps.
package caddysnake

import (
        "context"
        _ "embed"
        "encoding/json"
        "errors"
        "fmt"
        "log"
        "net"
        "net/http"
        "net/http/httputil"
        "os"
        "os/exec"
        "path/filepath"
        "runtime"
        "strconv"
        "strings"
        "sync"
        "sync/atomic"
        "syscall"
        "time"

        "github.com/caddyserver/caddy/v2"
        "github.com/caddyserver/caddy/v2/caddyconfig"
        "github.com/caddyserver/caddy/v2/caddyconfig/caddyfile"
        "github.com/caddyserver/caddy/v2/caddyconfig/httpcaddyfile"
        caddycmd "github.com/caddyserver/caddy/v2/cmd"
        "github.com/caddyserver/caddy/v2/modules/caddyhttp"
        "github.com/caddyserver/certmagic"
        "github.com/spf13/cobra"
        "go.uber.org/zap"
        "go.uber.org/zap/zapcore"

        "github.com/caddyserver/caddy/v2/modules/caddyhttp/encode"
        _ "github.com/caddyserver/caddy/v2/modules/caddyhttp/encode/gzip"
        _ "github.com/caddyserver/caddy/v2/modules/caddyhttp/encode/zstd"
        _ "github.com/caddyserver/caddy/v2/modules/caddyhttp/fileserver"
)

//go:embed caddysnake.py
var caddysnake_py string

// AppServer defines the interface to interacting with a WSGI or ASGI server
type AppServer interface {
        Cleanup() error
        HandleRequest(w http.ResponseWriter, r *http.Request) error
}

// CaddySnake module that communicates with a Python app
type CaddySnake struct {
        ModuleWsgi string `json:"module_wsgi,omitempty"`
        ModuleAsgi string `json:"module_asgi,omitempty"`
        Lifespan   string `json:"lifespan,omitempty"`
        WorkingDir string `json:"working_dir,omitempty"`
        VenvPath   string `json:"venv_path,omitempty"`
        Workers    string `json:"workers,omitempty"`
        Autoreload string `json:"autoreload,omitempty"`
        PythonPath string `json:"python_path,omitempty"`
        logger     *zap.Logger
        app        AppServer
}

// UnmarshalCaddyfile implements caddyfile.Unmarshaler.
func (f *CaddySnake) UnmarshalCaddyfile(d *caddyfile.Dispenser) error {
        for d.Next() {
                args := d.RemainingArgs()
                if len(args) == 1 {
                        f.ModuleWsgi = args[0]
                } else if len(args) == 0 {
                        for nesting := d.Nesting(); d.NextBlock(nesting); {
                                switch d.Val() {
                                case "module_asgi":
                                        if !d.Args(&f.ModuleAsgi) {
                                                return d.Errf("expected exactly one argument for module_asgi")
                                        }
                                case "module_wsgi":
                                        if !d.Args(&f.ModuleWsgi) {
                                                return d.Errf("expected exactly one argument for module_wsgi")
                                        }
                                case "lifespan":
                                        if !d.Args(&f.Lifespan) || (f.Lifespan != "on" && f.Lifespan != "off") {
                                                return d.Errf("expected exactly one argument for lifespan: on|off")
                                        }
                                case "working_dir":
                                        if !d.Args(&f.WorkingDir) {
                                                return d.Errf("expected exactly one argument for working_dir")
                                        }
                                case "venv":
                                        if !d.Args(&f.VenvPath) {
                                                return d.Errf("expected exactly one argument for venv")
                                        }
                                case "workers":
                                        if !d.Args(&f.Workers) {
                                                return d.Errf("expected exactly one argument for workers")
                                        }
                                case "autoreload":
                                        f.Autoreload = "on"
                                case "python_path":
                                        if !d.Args(&f.PythonPath) {
                                                return d.Errf("expected exactly one argument for python_path")
                                        }
                                default:
                                        return d.Errf("unknown subdirective: %s", d.Val())
                                }
                        }
                } else {
                        return d.ArgErr()
                }
        }
        return nil
}

// CaddyModule returns the Caddy module information.
func (CaddySnake) CaddyModule() caddy.ModuleInfo {
        return caddy.ModuleInfo{
                ID:  "http.handlers.python",
                New: func() caddy.Module { return new(CaddySnake) },
        }
}

// Provision sets up the module.
func (f *CaddySnake) Provision(ctx caddy.Context) error {
        var err error
        f.logger = ctx.Logger(f)
        workers, _ := strconv.Atoi(f.Workers)
        if workers <= 0 {
                workers = runtime.GOMAXPROCS(0)
        }

        isDynamic := containsPlaceholder(f.ModuleWsgi) || containsPlaceholder(f.ModuleAsgi) ||
                containsPlaceholder(f.WorkingDir) || containsPlaceholder(f.VenvPath)

        if isDynamic {
                return f.provisionDynamic(workers)
        }

        pythonBin := resolvePythonInterpreter(f.PythonPath, f.VenvPath)

        if f.ModuleWsgi != "" {
                f.app, err = NewPythonWorkerGroup("wsgi", f.ModuleWsgi, f.WorkingDir, f.VenvPath, f.Lifespan, workers, pythonBin)
                if err != nil {
                        return err
                }
                if f.Lifespan != "" {
                        f.logger.Warn("lifespan attribute is ignored in WSGI mode", zap.String("lifespan", f.Lifespan))
                }
                f.logger.Info("serving wsgi app", zap.String("module_wsgi", f.ModuleWsgi), zap.String("working_dir", f.WorkingDir), zap.String("venv_path", f.VenvPath), zap.String("python", pythonBin))
        } else if f.ModuleAsgi != "" {
                f.app, err = NewPythonWorkerGroup("asgi", f.ModuleAsgi, f.WorkingDir, f.VenvPath, f.Lifespan, workers, pythonBin)
                if err != nil {
                        return err
                }
                f.logger.Info("serving asgi app", zap.String("module_asgi", f.ModuleAsgi), zap.String("working_dir", f.WorkingDir), zap.String("venv_path", f.VenvPath), zap.String("python", pythonBin))
        } else {
                return errors.New("asgi or wsgi app needs to be specified")
        }

        if f.Autoreload == "on" {
                watchDir := f.WorkingDir
                if watchDir == "" {
                        watchDir = "."
                }
                absDir, absErr := filepath.Abs(watchDir)
                if absErr != nil {
                        return fmt.Errorf("autoreload: %w", absErr)
                }

                var factory func() (AppServer, error)
                if f.ModuleWsgi != "" {
                        factory = func() (AppServer, error) {
                                return NewPythonWorkerGroup("wsgi", f.ModuleWsgi, f.WorkingDir, f.VenvPath, f.Lifespan, workers, pythonBin)
                        }
                } else {
                        factory = func() (AppServer, error) {
                                return NewPythonWorkerGroup("asgi", f.ModuleAsgi, f.WorkingDir, f.VenvPath, f.Lifespan, workers, pythonBin)
                        }
                }

                // Keep Caddy running on reload errors; failed app serves 503 until recovery.
                f.app, err = NewAutoreloadableApp(f.app, absDir, factory, f.logger, nil)
                if err != nil {
                        return fmt.Errorf("autoreload: %w", err)
                }
        }

        return nil
}

// provisionDynamic sets up the module in dynamic mode where Caddy placeholders
// in module_wsgi/module_asgi, working_dir, or venv are resolved per-request.
func (f *CaddySnake) provisionDynamic(workers int) error {
        autoreload := f.Autoreload == "on"
        pythonPath := f.PythonPath

        if f.ModuleWsgi != "" {
                lifespan := f.Lifespan
                factory := func(module, dir, venv string) (AppServer, error) {
                        pythonBin := resolvePythonInterpreter(pythonPath, venv)
                        return NewPythonWorkerGroup("wsgi", module, dir, venv, lifespan, workers, pythonBin)
                }
                var err error
                f.app, err = NewDynamicApp(f.ModuleWsgi, f.WorkingDir, f.VenvPath, factory, f.logger, autoreload, nil)
                if err != nil {
                        return err
                }
                if f.Lifespan != "" {
                        f.logger.Warn("lifespan attribute is ignored in WSGI mode", zap.String("lifespan", f.Lifespan))
                }
                f.logger.Info("serving dynamic wsgi app",
                        zap.String("module_wsgi", f.ModuleWsgi),
                        zap.String("working_dir", f.WorkingDir),
                        zap.String("venv_path", f.VenvPath),
                )
        } else if f.ModuleAsgi != "" {
                lifespan := f.Lifespan
                factory := func(module, dir, venv string) (AppServer, error) {
                        pythonBin := resolvePythonInterpreter(pythonPath, venv)
                        return NewPythonWorkerGroup("asgi", module, dir, venv, lifespan, workers, pythonBin)
                }
                var err error
                f.app, err = NewDynamicApp(f.ModuleAsgi, f.WorkingDir, f.VenvPath, factory, f.logger, autoreload, nil)
                if err != nil {
                        return err
                }
                f.logger.Info("serving dynamic asgi app",
                        zap.String("module_asgi", f.ModuleAsgi),
                        zap.String("working_dir", f.WorkingDir),
                        zap.String("venv_path", f.VenvPath),
                )
        } else {
                return errors.New("asgi or wsgi app needs to be specified")
        }
        return nil
}

// Validate implements caddy.Validator.
func (m *CaddySnake) Validate() error {
        if m.ModuleWsgi != "" && m.ModuleAsgi != "" {
                return errors.New("cannot specify both module_wsgi and module_asgi")
        }
        if m.ModuleWsgi == "" && m.ModuleAsgi == "" {
                return errors.New("one of module_wsgi or module_asgi is required")
        }
        if m.Workers != "" {
                w, err := strconv.Atoi(m.Workers)
                if err != nil || w < 0 {
                        return fmt.Errorf("invalid workers value: %s", m.Workers)
                }
        }
        if m.Lifespan != "" && m.Lifespan != "on" && m.Lifespan != "off" {
                return fmt.Errorf("lifespan must be 'on' or 'off', got: %s", m.Lifespan)
        }
        return nil
}

// Cleanup frees resources uses by module
func (m *CaddySnake) Cleanup() error {
        if m != nil && m.app != nil {
                m.logger.Info("cleaning up module")
                return m.app.Cleanup()
        }
        return nil
}

// ServeHTTP implements caddyhttp.MiddlewareHandler.
func (f CaddySnake) ServeHTTP(w http.ResponseWriter, r *http.Request, next caddyhttp.Handler) error {
        if err := f.app.HandleRequest(w, r); err != nil {
                return err
        }
        return nil
}

// Interface guards
var (
        _ caddy.Provisioner           = (*CaddySnake)(nil)
        _ caddy.Validator             = (*CaddySnake)(nil)
        _ caddy.CleanerUpper          = (*CaddySnake)(nil)
        _ caddyhttp.MiddlewareHandler = (*CaddySnake)(nil)
        _ caddyfile.Unmarshaler       = (*CaddySnake)(nil)
)

func parsePythonDirective(h httpcaddyfile.Helper) (caddyhttp.MiddlewareHandler, error) {
        var app CaddySnake
        if err := app.UnmarshalCaddyfile(h.Dispenser); err != nil {
                return nil, err
        }
        return app, nil
}

// resolvePythonInterpreter determines the Python interpreter to use.
// Priority: explicit python_path > venv/bin/python > system python3.
func resolvePythonInterpreter(pythonPath, venvPath string) string {
        if pythonPath != "" {
                return pythonPath
        }

        if venvPath != "" {
                var binDir string
                if runtime.GOOS == "windows" {
                        binDir = "Scripts"
                } else {
                        binDir = "bin"
                }
                venvPython := filepath.Join(venvPath, binDir, "python3")
                if runtime.GOOS == "windows" {
                        venvPython = filepath.Join(venvPath, binDir, "python.exe")
                }
                if _, err := os.Stat(venvPython); err == nil {
                        return venvPython
                }
                venvPython2 := filepath.Join(venvPath, binDir, "python")
                if _, err := os.Stat(venvPython2); err == nil {
                        return venvPython2
                }
        }

        return "python3"
}

// writeCaddysnakePy writes the embedded caddysnake.py to a temp file and returns the path.
func writeCaddysnakePy() (string, error) {
        f, err := os.CreateTemp("", "caddysnake-*.py")
        if err != nil {
                return "", err
        }
        if _, err := f.WriteString(caddysnake_py); err != nil {
                f.Close()
                os.Remove(f.Name())
                return "", err
        }
        f.Close()
        return f.Name(), nil
}

// proxyBufferPool implements httputil.BufferPool using sync.Pool to reduce GC pressure.
type proxyBufferPool struct {
        pool sync.Pool
}

func (p *proxyBufferPool) Get() []byte {
        b := p.pool.Get()
        if b == nil {
                return make([]byte, 32*1024)
        }
        return b.([]byte)
}

func (p *proxyBufferPool) Put(b []byte) {
        p.pool.Put(b)
}

var sharedProxyBufferPool = &proxyBufferPool{}

type PythonWorker struct {
        Interface  string
        App        string
        WorkingDir string
        Venv       string
        Lifespan   string
        PythonBin  string
        Socket     *os.File
        SockDir    string // private directory containing the socket (Unix only)
        ScriptPath string
        DialNet    string // "unix" or "tcp"
        DialAddr   string // socket path or host:port

        Cmd       *exec.Cmd
        Transport *http.Transport
        Proxy     *httputil.ReverseProxy
}

func NewPythonWorker(iface, app, workingDir, venv, lifespan, pythonBin, scriptPath string) (*PythonWorker, error) {
        var socket *os.File
        var sockDir string
        var err error
        if runtime.GOOS == "windows" {
                socket, err = os.CreateTemp("", "caddysnake-worker.port*")
        } else {
                sockDir, err = os.MkdirTemp("", "caddysnake-*")
                if err != nil {
                        return nil, err
                }
                if chErr := os.Chmod(sockDir, 0700); chErr != nil {
                        os.RemoveAll(sockDir)
                        return nil, chErr
                }
                socket, err = os.Create(filepath.Join(sockDir, "worker.sock"))
        }
        if err != nil {
                if sockDir != "" {
                        os.RemoveAll(sockDir)
                }
                return nil, err
        }
        path := socket.Name()
        socket.Close()

        dialNet := "unix"
        dialAddr := path
        if runtime.GOOS == "windows" {
                dialNet = "tcp"
                dialAddr = "" // set after Python writes port to path
        }

        w := &PythonWorker{
                Interface:  iface,
                App:        app,
                WorkingDir: workingDir,
                Venv:       venv,
                Lifespan:   lifespan,
                PythonBin:  pythonBin,
                Socket:     socket,
                SockDir:    sockDir,
                ScriptPath: scriptPath,
                DialNet:    dialNet,
                DialAddr:   dialAddr,
        }
        err = w.Start()
        return w, err
}

func (w *PythonWorker) Start() error {
        w.Transport = &http.Transport{
                DialContext: func(ctx context.Context, network, addr string) (net.Conn, error) {
                        return w.dialWithRetry(ctx)
                },
                MaxIdleConns:        1024,
                MaxIdleConnsPerHost: 256,
                IdleConnTimeout:     90 * time.Second,
                DisableCompression:  true,
        }
        w.Proxy = &httputil.ReverseProxy{
                Rewrite: func(req *httputil.ProxyRequest) {
                        req.Out.URL.Scheme = "http"
                        req.Out.URL.Host = w.DialAddr
                },
                Transport:  w.Transport,
                BufferPool: sharedProxyBufferPool,
        }

        workingDir := w.WorkingDir
        if workingDir == "" {
                workingDir, _ = os.Getwd()
        }

        args := []string{
                w.ScriptPath,
                "--interface", w.Interface,
                "--app", w.App,
                "--socket", w.Socket.Name(),
        }
        if workingDir != "" {
                args = append(args, "--working-dir", workingDir)
        }
        if w.Venv != "" {
                args = append(args, "--venv", w.Venv)
        }
        if w.Lifespan != "" {
                args = append(args, "--lifespan", w.Lifespan)
        }

        w.Cmd = exec.Command(w.PythonBin, args...)
        w.Cmd.Stdout = os.Stdout
        w.Cmd.Stderr = os.Stderr
        w.Cmd.Env = append(os.Environ(), "PYTHONUNBUFFERED=1")
        setSysProcAttr(w.Cmd)

        if err := w.Cmd.Start(); err != nil {
                return err
        }
        if runtime.GOOS == "windows" {
                port, err := waitForPortFile(w.Socket.Name(), 10*time.Second)
                if err != nil {
                        w.Cmd.Process.Kill()
                        _ = w.Cmd.Wait()
                        return fmt.Errorf("waiting for Python worker port file: %w", err)
                }
                w.DialAddr = "127.0.0.1:" + strconv.Itoa(port)
        }
        return nil
}

// waitForPortFile polls the given file path until it contains a valid port number.
func waitForPortFile(path string, timeout time.Duration) (int, error) {
        deadline := time.Now().Add(timeout)
        for time.Now().Before(deadline) {
                data, err := os.ReadFile(path)
                if err == nil && len(data) > 0 {
                        port, err := strconv.Atoi(strings.TrimSpace(string(data)))
                        if err == nil && port > 0 && port < 65536 {
                                return port, nil
                        }
                }
                time.Sleep(50 * time.Millisecond)
        }
        return 0, fmt.Errorf("port file %s not ready within %v", path, timeout)
}

// dialWithRetry attempts to establish a connection with retry logic
func (w *PythonWorker) dialWithRetry(ctx context.Context) (net.Conn, error) {
        const maxRetries = 5
        const baseDelay = 100 * time.Millisecond

        for attempt := 0; attempt < maxRetries; attempt++ {
                conn, err := net.Dial(w.DialNet, w.DialAddr)
                if err == nil {
                        return conn, nil
                }

                if attempt == maxRetries-1 {
                        return nil, fmt.Errorf("failed to connect after %d attempts: %w", maxRetries, err)
                }

                delay := baseDelay * time.Duration(1<<attempt)

                select {
                case <-ctx.Done():
                        return nil, ctx.Err()
                case <-time.After(delay):
                }
        }

        return nil, fmt.Errorf("unexpected error in dialWithRetry")
}

func (w *PythonWorker) Cleanup() error {
        if w.Transport != nil {
                w.Transport.CloseIdleConnections()
        }
        if w.Cmd != nil && w.Cmd.Process != nil {
                // On Windows, Signal(SIGTERM) is not supported; only Kill works.
                // Send SIGTERM on Unix for graceful shutdown (ASGI lifespan), Kill on Windows.
                if runtime.GOOS == "windows" {
                        w.Cmd.Process.Kill()
                } else {
                        _ = w.Cmd.Process.Signal(syscall.SIGTERM)
                }
                done := make(chan error, 1)
                go func() {
                        _, err := w.Cmd.Process.Wait()
                        done <- err
                }()
                select {
                case err := <-done:
                        if err != nil {
                                return err
                        }
                case <-time.After(5 * time.Second):
                        w.Cmd.Process.Kill()
                        <-done
                }
        }
        if w.Socket != nil {
                w.Socket.Close()
                os.Remove(w.Socket.Name())
                if w.SockDir != "" {
                        os.RemoveAll(w.SockDir)
                }
        }
        return nil
}

func (w *PythonWorker) HandleRequest(rw http.ResponseWriter, req *http.Request) error {
        w.Proxy.ServeHTTP(rw, req)
        return nil
}

type PythonWorkerGroup struct {
        Workers    []*PythonWorker
        roundRobin atomic.Uint64
        ScriptPath string
}

func NewPythonWorkerGroup(iface, app, workingDir, venv, lifespan string, count int, pythonBin string) (*PythonWorkerGroup, error) {
        scriptPath, err := writeCaddysnakePy()
        if err != nil {
                return nil, fmt.Errorf("failed to write caddysnake.py: %w", err)
        }

        errs := make([]error, count)
        workers := make([]*PythonWorker, count)
        for i := 0; i < count; i++ {
                workers[i], errs[i] = NewPythonWorker(iface, app, workingDir, venv, lifespan, pythonBin, scriptPath)
        }
        wg := &PythonWorkerGroup{
                Workers:    workers,
                ScriptPath: scriptPath,
        }
        if err := errors.Join(errs...); err != nil {
                return nil, errors.Join(wg.Cleanup(), err)
        }
        return wg, nil
}

func (wg *PythonWorkerGroup) Cleanup() error {
        if wg == nil {
                return nil
        }
        errs := make([]error, len(wg.Workers))
        for i, worker := range wg.Workers {
                if worker != nil {
                        errs[i] = worker.Cleanup()
                }
        }
        if wg.ScriptPath != "" {
                os.Remove(wg.ScriptPath)
        }
        return errors.Join(errs...)
}

func (wg *PythonWorkerGroup) HandleRequest(rw http.ResponseWriter, req *http.Request) error {
        n := wg.roundRobin.Add(1)
        idx := int(n % uint64(len(wg.Workers)))
        wg.Workers[idx].HandleRequest(rw, req)
        return nil
}

func init() {
        caddy.RegisterModule(CaddySnake{})
        httpcaddyfile.RegisterHandlerDirective("python", parsePythonDirective)
        caddycmd.RegisterCommand(caddycmd.Command{
                Name: "python-server",
                Usage: "--server-type wsgi|asgi --app <module> " +
                        "[--domain <example.com>] [--listen <addr>] [--workers <count>] " +
                        "[--python-path <path>] [--working-dir <path>] [--venv <path>] " +
                        "[--static-path <path>] [--static-route <route>] " +
                        "[--debug] [--access-logs] [--autoreload]",
                Short: "Spins up a Python server",
                Long: `
A Python WSGI or ASGI server designed for apps and frameworks.

You can specify a custom socket address using the '--listen' option. You can also specify the number of workers to spawn.

Providing a domain name with the '--domain' flag enables HTTPS and sets the listener to the appropriate secure port.
Ensure DNS A/AAAA records are correctly set up if using a public domain for secure connections.
`,
                CobraFunc: func(cmd *cobra.Command) {
                        cmd.Flags().StringP("server-type", "t", "", "Required. The type of server to use: wsgi|asgi")
                        cmd.Flags().StringP("app", "a", "", "Required. App module to be imported")
                        cmd.Flags().StringP("domain", "d", "", "Domain name at which to serve the files")
                        cmd.Flags().StringP("listen", "l", "", "The address to which to bind the listener")
                        cmd.Flags().StringP("workers", "w", "0", "The number of workers to spawn")
                        cmd.Flags().String("python-path", "", "Path to the Python interpreter")
                        cmd.Flags().String("working-dir", "", "Working directory for the Python app")
                        cmd.Flags().String("venv", "", "Path to a Python virtual environment to use")
                        cmd.Flags().String("static-path", "", "Path to a static directory to serve: path/to/static")
                        cmd.Flags().String("static-route", "/static", "Route to serve the static directory: /static")
                        cmd.Flags().Bool("debug", false, "Enable debug logs")
                        cmd.Flags().Bool("access-logs", false, "Enable access logs")
                        cmd.Flags().String("lifespan", "off", "Enable ASGI lifespan support (ignored in WSGI mode)")
                        cmd.Flags().Bool("autoreload", false, "Watch .py files and reload on changes")
                        cmd.RunE = caddycmd.WrapCommandFuncForCobra(pythonServer)
                },
        })
}

// pythonServer is inspired on the php-server command of the Frankenphp project (MIT License)
func pythonServer(fs caddycmd.Flags) (int, error) {
        caddy.TrapSignals()

        domain := fs.String("domain")
        app := fs.String("app")
        listen := fs.String("listen")
        workers := fs.String("workers")
        debug := fs.Bool("debug")
        accessLogs := fs.Bool("access-logs")
        autoreload := fs.Bool("autoreload")
        staticPath := fs.String("static-path")
        staticRoute := fs.String("static-route")
        serverType := fs.String("server-type")
        pythonPath := fs.String("python-path")
        workingDir := fs.String("working-dir")
        venv := fs.String("venv")
        lifespan := fs.String("lifespan")

        if serverType == "" {
                return caddy.ExitCodeFailedStartup, errors.New("--server-type is required")
        }
        if app == "" {
                return caddy.ExitCodeFailedStartup, errors.New("--app is required")
        }

        gzip, err := caddy.GetModule("http.encoders.gzip")
        if err != nil {
                return caddy.ExitCodeFailedStartup, err
        }

        zstd, err := caddy.GetModule("http.encoders.zstd")
        if err != nil {
                return caddy.ExitCodeFailedStartup, err
        }

        encodings := caddy.ModuleMap{
                "zstd": caddyconfig.JSON(zstd.New(), nil),
                "gzip": caddyconfig.JSON(gzip.New(), nil),
        }
        prefer := []string{"zstd", "gzip"}

        pythonHandler := CaddySnake{}
        if serverType == "wsgi" {
                pythonHandler.ModuleWsgi = app
        } else {
                pythonHandler.ModuleAsgi = app
        }
        if venv != "" {
                pythonHandler.VenvPath = venv
        } else if venv := os.Getenv("VIRTUAL_ENV"); venv != "" {
                pythonHandler.VenvPath = venv
        }

        pythonHandler.Workers = workers
        pythonHandler.PythonPath = pythonPath
        if autoreload {
                pythonHandler.Autoreload = "on"
        }
        pythonHandler.WorkingDir = workingDir
        pythonHandler.Lifespan = lifespan

        routes := caddyhttp.RouteList{}

        if staticPath != "" {
                if strings.HasSuffix(staticRoute, "/") {
                        staticRoute = staticRoute + "*"
                } else if !strings.HasSuffix(staticRoute, "/*") {
                        staticRoute = staticRoute + "/*"
                }
                staticRoute := caddyhttp.Route{
                        MatcherSetsRaw: []caddy.ModuleMap{
                                {
                                        "path": caddyconfig.JSON(caddyhttp.MatchPath{staticRoute}, nil),
                                },
                        },
                        HandlersRaw: []json.RawMessage{
                                caddyconfig.JSONModuleObject(encode.Encode{
                                        EncodingsRaw: encodings,
                                        Prefer:       prefer,
                                }, "handler", "encode", nil),
                                caddyconfig.JSON(map[string]interface{}{
                                        "handler": "file_server",
                                        "root":    staticPath,
                                }, nil),
                        },
                }
                routes = append(routes, staticRoute)
        }

        mainRoute := caddyhttp.Route{
                MatcherSetsRaw: []caddy.ModuleMap{
                        {
                                "path": caddyconfig.JSON(caddyhttp.MatchPath{"/*"}, nil),
                        },
                },
                HandlersRaw: []json.RawMessage{
                        caddyconfig.JSONModuleObject(encode.Encode{
                                EncodingsRaw: encodings,
                                Prefer:       prefer,
                        }, "handler", "encode", nil),
                        caddyconfig.JSONModuleObject(pythonHandler, "handler", "python", nil),
                },
        }
        routes = append(routes, mainRoute)

        subroute := caddyhttp.Subroute{
                Routes: routes,
        }

        route := caddyhttp.Route{
                HandlersRaw: []json.RawMessage{caddyconfig.JSONModuleObject(subroute, "handler", "subroute", nil)},
        }

        if domain != "" {
                route.MatcherSetsRaw = []caddy.ModuleMap{
                        {
                                "host": caddyconfig.JSON(caddyhttp.MatchHost{domain}, nil),
                        },
                }
        }

        server := &caddyhttp.Server{
                ReadHeaderTimeout: caddy.Duration(10 * time.Second),
                IdleTimeout:       caddy.Duration(30 * time.Second),
                MaxHeaderBytes:    1024 * 10,
                Routes:            caddyhttp.RouteList{route},
        }
        if listen == "" {
                if domain == "" {
                        listen = ":9080"
                } else {
                        listen = ":" + strconv.Itoa(certmagic.HTTPSPort)
                }
        }
        server.Listen = []string{listen}

        if accessLogs {
                server.Logs = &caddyhttp.ServerLogConfig{}
        }

        httpApp := caddyhttp.App{
                Servers: map[string]*caddyhttp.Server{"srv0": server},
        }

        var f bool
        cfg := &caddy.Config{
                Admin: &caddy.AdminConfig{
                        Disabled: false,
                        Config: &caddy.ConfigSettings{
                                Persist: &f,
                        },
                },
                AppsRaw: caddy.ModuleMap{
                        "http": caddyconfig.JSON(httpApp, nil),
                },
        }

        if debug {
                cfg.Logging = &caddy.Logging{
                        Logs: map[string]*caddy.CustomLog{
                                "default": {
                                        BaseLog: caddy.BaseLog{Level: zapcore.DebugLevel.CapitalString()},
                                },
                        },
                }
        }

        if err := caddy.Run(cfg); err != nil {
                return caddy.ExitCodeFailedStartup, err
        }

        log.Printf("Serving Python app on %s", listen)

        select {}
}

// findSitePackagesInVenv searches for the site-packages directory in a given venv.
func findSitePackagesInVenv(venvPath string) (string, error) {
        var sitePackagesPath string
        if runtime.GOOS == "windows" {
                sitePackagesPath = filepath.Join(venvPath, "Lib\\site-packages")
        } else {
                libPath := filepath.Join(venvPath, "lib")
                pythonDir, err := findPythonDirectory(libPath)
                if err != nil {
                        return "", err
                }
                sitePackagesPath = filepath.Join(libPath, pythonDir, "site-packages")
        }
        fileInfo, err := os.Stat(sitePackagesPath)
        if err != nil {
                if os.IsNotExist(err) {
                        return "", fmt.Errorf("site-packages directory does not exist in: %s", sitePackagesPath)
                }
                return "", err
        }
        if !fileInfo.IsDir() {
                return "", fmt.Errorf("found site-packages is not a directory: %s", sitePackagesPath)
        }
        return sitePackagesPath, nil
}

// findWorkingDirectory checks if the directory exists and returns the absolute path
func findWorkingDirectory(workingDir string) (string, error) {
        workingDirAbs, err := filepath.Abs(workingDir)
        if err != nil {
                return "", err
        }
        fileInfo, err := os.Stat(workingDirAbs)
        if err != nil {
                if os.IsNotExist(err) {
                        return "", fmt.Errorf("working_dir directory does not exist in: %s", workingDirAbs)
                }
                return "", err
        }
        if !fileInfo.IsDir() {
                return "", fmt.Errorf("working_dir is not a directory: %s", workingDirAbs)
        }
        return workingDirAbs, nil
}

// findPythonDirectory searches for a directory that matches "python3.*" inside the given libPath.
func findPythonDirectory(libPath string) (string, error) {
        entries, err := os.ReadDir(libPath)
        if err != nil {
                return "", fmt.Errorf("unable to read venv lib directory: %w", err)
        }
        for _, e := range entries {
                if e.IsDir() && strings.HasPrefix(e.Name(), "python3") {
                        return e.Name(), nil
                }
        }
        return "", errors.New("unable to find a python3.* directory in the venv")
}

package caddysnake

import (
        "errors"
        "fmt"
        "net/http"
        "os"
        "path/filepath"
        "regexp"
        "strings"
        "sync"
        "time"

        "github.com/caddyserver/caddy/v2"
        "github.com/fsnotify/fsnotify"
        "go.uber.org/zap"
)

var validModulePattern = regexp.MustCompile(`^[a-zA-Z_][\w.]*:[a-zA-Z_]\w*$`)

func validateResolvedValues(module, dir, venv string) error {
        if !validModulePattern.MatchString(module) {
                return fmt.Errorf("invalid module name: %q", module)
        }
        if dir != "" {
                absDir, err := filepath.Abs(dir)
                if err != nil {
                        return fmt.Errorf("invalid working directory: %w", err)
                }
                if strings.Contains(absDir, "..") {
                        return fmt.Errorf("working directory contains path traversal: %q", dir)
                }
        }
        if venv != "" {
                absVenv, err := filepath.Abs(venv)
                if err != nil {
                        return fmt.Errorf("invalid venv path: %w", err)
                }
                if strings.Contains(absVenv, "..") {
                        return fmt.Errorf("venv path contains path traversal: %q", venv)
                }
        }
        return nil
}

// containsPlaceholder checks if a string contains Caddy placeholders (e.g. {host.labels.0}).
func containsPlaceholder(s string) bool {
        return strings.Contains(s, "{") && strings.Contains(s, "}")
}

// appFactory is a function that creates a new AppServer for a resolved
// module, working directory and venv path combination.
type appFactory func(resolvedModule, resolvedDir, resolvedVenv string) (AppServer, error)

// DynamicApp implements AppServer by lazily importing Python apps based on
// Caddy placeholders resolved at request time. For example, when working_dir
// contains {host.labels.2}, each subdomain gets its own Python app instance
// imported from the corresponding directory.
type DynamicApp struct {
        mu            sync.RWMutex
        apps          map[string]AppServer
        modulePattern string
        workingDir    string
        venvPath      string
        factory       appFactory
        logger        *zap.Logger

        // Autoreload fields
        autoreload          bool
        watcher             *fsnotify.Watcher
        dirToKeys           map[string][]string // abs working dir -> cache keys that use it
        stopCh              chan struct{}
        exitOnReloadFailure func(code int) // if set and autoreload, process exits when app creation fails
}

// NewDynamicApp creates a DynamicApp that resolves placeholders from
// modulePattern, workingDir, and venvPath at request time and lazily creates
// Python app instances via the supplied factory function.
// When autoreload is true, if exitOnReloadFailure is non-nil it is called with
// code 1 when app creation fails (e.g. app deleted), so the process can terminate.
func NewDynamicApp(modulePattern, workingDir, venvPath string, factory appFactory, logger *zap.Logger, autoreload bool, exitOnReloadFailure func(code int)) (*DynamicApp, error) {
        d := &DynamicApp{
                apps:                 make(map[string]AppServer),
                modulePattern:        modulePattern,
                workingDir:           workingDir,
                venvPath:             venvPath,
                factory:              factory,
                logger:               logger,
                autoreload:           autoreload,
                exitOnReloadFailure:  exitOnReloadFailure,
        }

        if autoreload {
                watcher, err := fsnotify.NewWatcher()
                if err != nil {
                        return nil, err
                }
                d.watcher = watcher
                d.dirToKeys = make(map[string][]string)
                d.stopCh = make(chan struct{})
                go d.watchForChanges()
                logger.Info("autoreload enabled for dynamic app")
        }

        return d, nil
}

// resolve uses the Caddy replacer from the request context to substitute
// placeholders in the module pattern, working directory, and venv path.
func (d *DynamicApp) resolve(r *http.Request) (key, module, dir, venv string) {
        module = d.modulePattern
        dir = d.workingDir
        venv = d.venvPath

        if repl, ok := r.Context().Value(caddy.ReplacerCtxKey).(*caddy.Replacer); ok && repl != nil {
                module = repl.ReplaceAll(module, "")
                dir = repl.ReplaceAll(dir, "")
                venv = repl.ReplaceAll(venv, "")
        }

        key = module + "|" + dir + "|" + venv
        return
}

// getOrCreateApp returns an existing app for the given key, or creates one
// using the factory if it doesn't exist yet.
func (d *DynamicApp) getOrCreateApp(key, module, dir, venv string) (AppServer, error) {
        if err := validateResolvedValues(module, dir, venv); err != nil {
                return nil, err
        }

        d.mu.RLock()
        app, ok := d.apps[key]
        d.mu.RUnlock()
        if ok {
                return app, nil
        }

        d.mu.Lock()
        defer d.mu.Unlock()

        app, ok = d.apps[key]
        if ok {
                return app, nil
        }

        d.logger.Info("dynamically importing python app",
                zap.String("module", module),
                zap.String("working_dir", dir),
                zap.String("venv", venv),
        )

        app, err := d.factory(module, dir, venv)
        if err != nil {
                return nil, err
        }

        d.apps[key] = app

        if d.autoreload && dir != "" {
                d.startWatchingDir(dir, key)
        }

        return app, nil
}

func (d *DynamicApp) startWatchingDir(dir, key string) {
        absDir, err := filepath.Abs(dir)
        if err != nil {
                d.logger.Warn("autoreload: failed to resolve directory",
                        zap.String("dir", dir),
                        zap.Error(err),
                )
                return
        }

        if keys, ok := d.dirToKeys[absDir]; ok {
                for _, k := range keys {
                        if k == key {
                                return
                        }
                }
                d.dirToKeys[absDir] = append(keys, key)
                return
        }

        d.dirToKeys[absDir] = []string{key}
        watchDirRecursive(d.watcher, absDir, d.logger)
}

func (d *DynamicApp) watchForChanges() {
        var debounceTimer *time.Timer
        const debounceDuration = 500 * time.Millisecond

        pendingDirs := make(map[string]bool)
        var pendingMu sync.Mutex

        for {
                select {
                case event, ok := <-d.watcher.Events:
                        if !ok {
                                return
                        }
                        if !isPythonFileEvent(event) {
                                handleNewDirEvent(event, d.watcher)
                                continue
                        }

                        d.logger.Debug("python file changed (dynamic)",
                                zap.String("file", event.Name),
                                zap.String("op", event.Op.String()),
                        )

                        d.mu.RLock()
                        for absDir := range d.dirToKeys {
                                if strings.HasPrefix(event.Name, absDir+string(os.PathSeparator)) ||
                                        strings.HasPrefix(event.Name, absDir) {
                                        pendingMu.Lock()
                                        pendingDirs[absDir] = true
                                        pendingMu.Unlock()
                                }
                        }
                        d.mu.RUnlock()

                        if debounceTimer != nil {
                                debounceTimer.Stop()
                        }
                        debounceTimer = time.AfterFunc(debounceDuration, func() {
                                pendingMu.Lock()
                                dirs := make([]string, 0, len(pendingDirs))
                                for dir := range pendingDirs {
                                        dirs = append(dirs, dir)
                                }
                                pendingDirs = make(map[string]bool)
                                pendingMu.Unlock()

                                for _, dir := range dirs {
                                        d.reloadDir(dir)
                                }
                        })
                case err, ok := <-d.watcher.Errors:
                        if !ok {
                                return
                        }
                        d.logger.Error("autoreload watcher error", zap.Error(err))
                case <-d.stopCh:
                        if debounceTimer != nil {
                                debounceTimer.Stop()
                        }
                        return
                }
        }
}

// reloadDir evicts all apps associated with the given directory and
// cleans them up after a grace period.
func (d *DynamicApp) reloadDir(absDir string) {
        d.logger.Info("reloading dynamic python apps due to file changes",
                zap.String("working_dir", absDir),
        )

        d.mu.Lock()

        keys, ok := d.dirToKeys[absDir]
        if !ok {
                d.mu.Unlock()
                return
        }

        var oldApps []AppServer
        for _, key := range keys {
                if app, exists := d.apps[key]; exists {
                        oldApps = append(oldApps, app)
                        delete(d.apps, key)
                }
        }

        delete(d.dirToKeys, absDir)

        d.mu.Unlock()

        d.logger.Info("dynamic python apps evicted, will reimport on next request",
                zap.String("working_dir", absDir),
                zap.Int("apps_evicted", len(oldApps)),
        )

        if len(oldApps) > 0 {
                go func() {
                        time.Sleep(10 * time.Second)
                        for _, app := range oldApps {
                                if err := app.Cleanup(); err != nil {
                                        d.logger.Error("failed to cleanup old dynamic app",
                                                zap.Error(err),
                                        )
                                }
                        }
                }()
        }
}

// HandleRequest resolves placeholders from the request, gets or creates the
// appropriate app, and forwards the request.
func (d *DynamicApp) HandleRequest(w http.ResponseWriter, r *http.Request) error {
        key, module, dir, venv := d.resolve(r)
        app, err := d.getOrCreateApp(key, module, dir, venv)
        if err != nil {
                if d.autoreload && d.exitOnReloadFailure != nil {
                        d.logger.Error("failed to load python app (autoreload); terminating",
                                zap.String("module", module),
                                zap.String("working_dir", dir),
                                zap.Error(err),
                        )
                        d.exitOnReloadFailure(1)
                }
                return err
        }
        return app.HandleRequest(w, r)
}

// Cleanup frees all dynamically created apps and stops the autoreload watcher.
func (d *DynamicApp) Cleanup() error {
        d.mu.Lock()
        defer d.mu.Unlock()

        if d.autoreload && d.stopCh != nil {
                close(d.stopCh)
                d.watcher.Close()
        }

        var errs []error
        for key, app := range d.apps {
                if err := app.Cleanup(); err != nil {
                        errs = append(errs, err)
                }
                delete(d.apps, key)
        }
        return errors.Join(errs...)
}

//go:build linux

package caddysnake

import (
        "os/exec"
        "syscall"
)

func setSysProcAttr(cmd *exec.Cmd) {
        cmd.SysProcAttr = &syscall.SysProcAttr{
                Pdeathsig: syscall.SIGTERM,
        }
}