// Package terminal gives the browser an interactive shell: each WebSocket
// connection to its handler spawns one shell inside a pseudo-terminal and
// pipes bytes both ways, which is exactly what xterm.js expects to drive.
//
// Wire protocol: the server sends raw terminal output as binary frames; the
// client sends JSON text frames — {"type":"input","data":"…"} for keystrokes
// and {"type":"resize","cols":N,"rows":N} when the viewport changes.
package terminal

import (
	"context"
	"encoding/json"
	"log/slog"
	"net/http"
	"os"
	"os/exec"

	"github.com/coder/websocket"
	"github.com/creack/pty"
)

// Service spawns one shell per WebSocket connection.
type Service struct {
	// Shell is the command to run; when empty, $SHELL then bash then sh.
	Shell []string
	// Cwd is the directory each shell starts in.
	Cwd string

	logger *slog.Logger
}

// New creates a terminal service starting shells in cwd.
//
// Example:
//
//	mux.Handle("GET /ws/terminal", terminal.New(cfg.Cwd, logger).WebSocketHandler())
func New(cwd string, logger *slog.Logger) *Service {
	if logger == nil {
		logger = slog.Default()
	}
	return &Service{Cwd: cwd, logger: logger}
}

// controlMessage is what the client sends over text frames.
type controlMessage struct {
	Type string `json:"type"`
	Data string `json:"data,omitempty"`
	Cols uint16 `json:"cols,omitempty"`
	Rows uint16 `json:"rows,omitempty"`
}

// WebSocketHandler upgrades the request and runs one shell session until the
// socket closes or the shell exits.
func (s *Service) WebSocketHandler() http.Handler {
	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		conn, err := websocket.Accept(w, r, &websocket.AcceptOptions{
			OriginPatterns: []string{"localhost:*", "127.0.0.1:*"},
		})
		if err != nil {
			s.logger.Warn("terminal websocket accept failed", "error", err)
			return
		}
		defer func() { _ = conn.CloseNow() }()
		s.serveShell(r.Context(), conn)
	})
}

// serveShell runs the shell and the two pumps for one connection.
func (s *Service) serveShell(ctx context.Context, conn *websocket.Conn) {
	cmd := exec.Command(s.shellCommand()[0], s.shellCommand()[1:]...)
	cmd.Dir = s.Cwd
	cmd.Env = append(os.Environ(), "TERM=xterm-256color")

	ptmx, err := pty.Start(cmd)
	if err != nil {
		s.logger.Error("pty start failed", "error", err)
		_ = conn.Close(websocket.StatusInternalError, "pty start failed")
		return
	}
	defer func() {
		_ = ptmx.Close()
		_ = cmd.Process.Kill()
		_ = cmd.Wait()
	}()

	// Output pump: pty → socket, as binary frames.
	outputDone := make(chan struct{})
	go func() {
		defer close(outputDone)
		buffer := make([]byte, 8192)
		for {
			n, readErr := ptmx.Read(buffer)
			if n > 0 {
				if writeErr := conn.Write(ctx, websocket.MessageBinary, buffer[:n]); writeErr != nil {
					return
				}
			}
			if readErr != nil {
				// The shell exited (or the pty closed): tell the client.
				_ = conn.Close(websocket.StatusNormalClosure, "shell exited")
				return
			}
		}
	}()

	// Input pump: socket → pty, until the client disconnects.
	s.readControlMessages(ctx, conn, ptmx)
	<-outputDone
}

// readControlMessages feeds client frames into the pty.
func (s *Service) readControlMessages(ctx context.Context, conn *websocket.Conn, ptmx *os.File) {
	for {
		kind, data, err := conn.Read(ctx)
		if err != nil {
			return
		}
		if kind != websocket.MessageText {
			continue
		}
		var msg controlMessage
		if err := json.Unmarshal(data, &msg); err != nil {
			continue // a malformed frame must not kill the session
		}
		switch msg.Type {
		case "input":
			_, _ = ptmx.WriteString(msg.Data)
		case "resize":
			if msg.Cols > 0 && msg.Rows > 0 {
				_ = pty.Setsize(ptmx, &pty.Winsize{Cols: msg.Cols, Rows: msg.Rows})
			}
		}
	}
}

// shellCommand picks the shell to run: the configured one, else $SHELL, else
// bash, else sh.
func (s *Service) shellCommand() []string {
	if len(s.Shell) > 0 {
		return s.Shell
	}
	if shell := os.Getenv("SHELL"); shell != "" {
		return []string{shell, "-l"}
	}
	if path, err := exec.LookPath("bash"); err == nil {
		return []string{path, "-l"}
	}
	return []string{"sh"}
}