// Simple HTTP server that responds 200 with a test message including client IP.
//
// Behavior:
//   - GET returns HTTP/1.1 200 with a message including the requester IP
//   - HEAD returns the same headers as GET but no body
//   - Displays incoming X-* headers when present
//   - If User-Agent contains "curl" or "wget" (case-insensitive), the server
//     responds with Content-Type: text/plain; otherwise Content-Type: text/html.
//   - --pem specifies a PEM bundle file (cert chain + private key). When provided,
//     the server listens on both HTTP (--port) and HTTPS (--tls-port).
//     Without --pem the server listens on plain HTTP only.
//   - --look controls which HTML variant is returned for non-CLI agents:
//   - basic     - plain HTML with no external references
//   - nice      - includes Google Font "Noto Sans" (default)
//   - bootstrap - Bootstrap 5 layout
//   - tailwind  - Tailwind CSS via @tailwindcss/browser@latest
//   - The URL query parameter "look" (e.g. /?look=nice) overrides --look for
//     that request only. Values are case-insensitive and must be basic,nice,bootstrap,tailwind.
//
// Usage:
//
//	go run ok-server.go                          # binds 0.0.0.0:8080, nice look
//	go run ok-server.go --look basic
//	go run ok-server.go -b 127.0.0.1 -p 8080 --look tailwind
//	go run ok-server.go --pem /path/to/bundle.pem  # HTTP :8080 + HTTPS :8443
//	go run ok-server.go --pem /path/to/bundle.pem --tls-port 9443
package main

import (
	"context"
	"crypto/tls"
	"flag"
	"fmt"
	"html"
	"net"
	"net/http"
	"os"
	"os/signal"
	"sort"
	"strconv"
	"strings"
	"syscall"
	"time"
)

const plainTemplate = "Hello, This is a test HTTP server.\n\nYour request came from {ip}.\n\n{proxy_headers_block}Have a nice day!\n"

const htmlBasic = `<!doctype html>
<html lang="en">
  <head>
    <meta charset="utf-8">
    <title>Test HTTP server</title>
  </head>
  <body>
    <h1>Hello,</h1>
    <p>This is a test HTTP server.</p>
    <p>Your request came from {ip_html}.</p>
    {proxy_headers_html}
    <p>Have a nice day!</p>
  </body>
</html>
`

const htmlNice = `<!doctype html>
<html lang="en">
  <head>
    <meta charset="utf-8">
    <meta name="viewport" content="width=device-width, initial-scale=1">
    <title>Test HTTP server</title>
    <!-- Google Font: Noto Sans -->
    <link href="https://fonts.googleapis.com/css2?family=Noto+Sans:wght@400;700&display=swap" rel="stylesheet">
    <style>
      body { font-family: "Noto Sans", system-ui, -apple-system, "Segoe UI", Roboto, "Helvetica Neue", Arial, sans-serif; padding: 2rem; background: #f6f7fb; }
      .card { background: white; padding: 1.5rem; border-radius: 8px; box-shadow: 0 2px 8px rgba(0,0,0,0.06); max-width: 800px; margin: 2rem auto; }
      h1 { margin-top: 0; }
    </style>
  </head>
  <body>
    <div class="card">
      <h1>Hello,</h1>
      <p>This is a test HTTP server.</p>
      <p>Your request came from {ip_html}.</p>
      {proxy_headers_html}
      <p>Have a nice day!</p>
    </div>
  </body>
</html>
`

const htmlBootstrap = `<!doctype html>
<html lang="en">
  <head>
    <meta charset="utf-8">
    <meta name="viewport" content="width=device-width, initial-scale=1">
    <title>Test HTTP server</title>
    <link href="https://cdn.jsdelivr.net/npm/bootstrap@5.3.2/dist/css/bootstrap.min.css" rel="stylesheet" crossorigin="anonymous">
  </head>
  <body class="bg-light">
    <nav class="navbar navbar-expand-lg navbar-dark bg-primary mb-4">
      <div class="container">
        <a class="navbar-brand" href="#">Test HTTP server</a>
      </div>
    </nav>

    <main class="container">
      <div class="row justify-content-center">
        <div class="col-md-8">
          <div class="card shadow-sm">
            <div class="card-body">
              <h1 class="card-title">Hello,</h1>
              <p class="card-text">This is a test HTTP server.</p>
              <p class="card-text">Your request came from {ip_html}.</p>
              {proxy_headers_html}
              <hr>
              <p class="mb-0">Have a nice day!</p>
            </div>
          </div>
        </div>
      </div>
    </main>

    <footer class="text-center mt-4 mb-4">
      <small class="text-muted">Simple status page for firewall/testing</small>
    </footer>
  </body>
</html>
`

const htmlTailwind = `<!doctype html>
<html lang="en">
  <head>
    <meta charset="utf-8">
    <meta name="viewport" content="width=device-width, initial-scale=1">
    <title>Test HTTP server</title>
    <script src="https://cdn.jsdelivr.net/npm/@tailwindcss/browser@latest"></script>
  </head>
  <body class="min-h-screen bg-slate-100 text-slate-900">
    <div class="mx-auto max-w-3xl px-4 py-10 sm:py-16">
      <div class="overflow-hidden rounded-2xl border border-slate-200 bg-white shadow-lg">
        <div class="bg-slate-900 px-6 py-4 text-slate-100">
          <h1 class="text-xl font-semibold">Test HTTP server</h1>
          <p class="mt-1 text-sm text-slate-300">Simple status page for firewall/testing</p>
        </div>
        <main class="space-y-4 px-6 py-6">
          <p class="text-2xl font-bold">Hello,</p>
          <p>This is a test HTTP server.</p>
          <p>Your request came from {ip_html}.</p>
          {proxy_headers_html}
          <p class="pt-2">Have a nice day!</p>
        </main>
      </div>
    </div>
  </body>
</html>
`

var lookTemplates = map[string]string{
	"basic":     htmlBasic,
	"nice":      htmlNice,
	"bootstrap": htmlBootstrap,
	"tailwind":  htmlTailwind,
}

func getClientIP(r *http.Request) string {
	host, _, err := net.SplitHostPort(r.RemoteAddr)
	if err != nil {
		return r.RemoteAddr
	}
	return host
}

func isCliAgent(ua string) bool {
	ua = strings.ToLower(ua)
	return strings.Contains(ua, "curl") || strings.Contains(ua, "wget")
}

func getRequestLook(r *http.Request) string {
	look := strings.ToLower(strings.TrimSpace(r.URL.Query().Get("look")))
	if _, ok := lookTemplates[look]; ok {
		return look
	}
	return ""
}

func collectXHeaders(r *http.Request) [][2]string {
	var result [][2]string
	for name, values := range r.Header {
		if strings.HasPrefix(strings.ToLower(name), "x-") {
			result = append(result, [2]string{name, strings.Join(values, ", ")})
		}
	}
	sort.Slice(result, func(i, j int) bool {
		return strings.ToLower(result[i][0]) < strings.ToLower(result[j][0])
	})
	return result
}

func proxyMarkup(xHeaders [][2]string) (string, string) {
	if len(xHeaders) == 0 {
		return "", ""
	}
	lines := make([]string, len(xHeaders))
	for i, h := range xHeaders {
		lines[i] = fmt.Sprintf("  - %s: %s", h[0], h[1])
	}
	block := "Reverse proxy condition: detected via X-* headers.\nX-* headers:\n" +
		strings.Join(lines, "\n") + "\n\n"

	htmlItems := make([]string, len(xHeaders))
	for i, h := range xHeaders {
		htmlItems[i] = fmt.Sprintf("<li><code>%s</code>: %s</li>",
			html.EscapeString(h[0]), html.EscapeString(h[1]))
	}
	htmlStr := "<p>Reverse proxy condition: <strong>detected via X-* headers</strong>.</p>" +
		"<p>X-* headers:</p><ul>" + strings.Join(htmlItems, "") + "</ul>"

	return block, htmlStr
}

func makeBodyAndType(r *http.Request, clientIP, defaultLook string) ([]byte, string) {
	ua := r.Header.Get("User-Agent")
	xHeaders := collectXHeaders(r)
	proxyHeadersBlock, proxyHeadersHTML := proxyMarkup(xHeaders)

	var ipText, ipHTML string
	xff := strings.TrimSpace(strings.SplitN(r.Header.Get("X-Forwarded-For"), ",", 2)[0])
	if xff != "" {
		realIP := xff
		if strings.HasPrefix(realIP, "[") {
			if idx := strings.Index(realIP, "]"); idx != -1 {
				realIP = realIP[1:idx]
			} else {
				realIP = realIP[1:]
			}
		} else if strings.Count(realIP, ":") == 1 {
			realIP = strings.SplitN(realIP, ":", 2)[0]
		}
		ipText = fmt.Sprintf("%s (forwarded by proxy %s)", realIP, clientIP)
		ipHTML = fmt.Sprintf("<strong>%s</strong> (forwarded by proxy <strong>%s</strong>)",
			html.EscapeString(realIP), html.EscapeString(clientIP))
	} else {
		ipText = clientIP
		ipHTML = fmt.Sprintf("<strong>%s</strong>", html.EscapeString(clientIP))
	}

	if isCliAgent(ua) {
		body := strings.NewReplacer(
			"{ip}", ipText,
			"{proxy_headers_block}", proxyHeadersBlock,
		).Replace(plainTemplate)
		return []byte(body), "text/plain; charset=utf-8"
	}

	look := getRequestLook(r)
	if look == "" {
		look = defaultLook
	}
	body := strings.NewReplacer(
		"{ip_html}", ipHTML,
		"{proxy_headers_html}", proxyHeadersHTML,
	).Replace(lookTemplates[look])
	return []byte(body), "text/html; charset=utf-8"
}

func makeHandler(defaultLook string) http.Handler {
	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		clientIP := getClientIP(r)
		now := time.Now().UTC().Format("02/Jan/2006:15:04:05 +0000")
		fmt.Fprintf(os.Stderr, "%s - - [%s] \"%s %s\"\n", clientIP, now, r.Method, r.URL.RequestURI())

		if r.Method != http.MethodGet && r.Method != http.MethodHead {
			w.Header().Set("Content-Length", "0")
			w.WriteHeader(http.StatusMethodNotAllowed)
			return
		}

		body, contentType := makeBodyAndType(r, clientIP, defaultLook)
		w.Header().Set("Content-Type", contentType)
		w.Header().Set("Content-Length", strconv.Itoa(len(body)))
		w.WriteHeader(http.StatusOK)
		if r.Method != http.MethodHead {
			w.Write(body) //nolint:errcheck
		}
	})
}

func main() {
	fs := flag.NewFlagSet("ok-server", flag.ContinueOnError)
	fs.Usage = func() {
		fmt.Println("Usage:")
		fmt.Println("  go run ok-server.go [-b|--bind ADDRESS] [-p|--port PORT] [--tls-port PORT] [--pem PEMFILE] [--look basic|nice|bootstrap|tailwind]")
		fmt.Println()
		fmt.Println("Defaults:")
		fmt.Println("  --bind     0.0.0.0")
		fmt.Println("  --port     8080")
		fmt.Println("  --tls-port 8443")
		fmt.Println("  --pem      (none, plain HTTP only)")
		fmt.Println("  --look     nice")
	}

	var bind, look, pemFile string
	var port, tlsPort int
	fs.StringVar(&bind, "b", "0.0.0.0", "address to bind to")
	fs.StringVar(&bind, "bind", "0.0.0.0", "address to bind to")
	fs.IntVar(&port, "p", 8080, "HTTP port")
	fs.IntVar(&port, "port", 8080, "HTTP port")
	fs.IntVar(&tlsPort, "tls-port", 8443, "HTTPS port (when --pem is provided)")
	fs.StringVar(&pemFile, "pem", "", "PEM bundle (cert+key); enables HTTPS alongside HTTP")
	fs.StringVar(&look, "look", "nice", "default HTML look: basic|nice|bootstrap|tailwind")

	if err := fs.Parse(os.Args[1:]); err != nil {
		if err == flag.ErrHelp {
			os.Exit(0)
		}
		fmt.Fprintln(os.Stderr, err)
		os.Exit(2)
	}

	if _, ok := lookTemplates[look]; !ok {
		fmt.Fprintf(os.Stderr, "Invalid look: %q (expected one of: basic, nice, bootstrap, tailwind)\n", look)
		os.Exit(2)
	}

	handler := makeHandler(look)

	httpSrv := &http.Server{
		Addr:    fmt.Sprintf("%s:%d", bind, port),
		Handler: handler,
	}

	var httpsSrv *http.Server
	if pemFile != "" {
		cert, err := tls.LoadX509KeyPair(pemFile, pemFile)
		if err != nil {
			fmt.Fprintf(os.Stderr, "Cannot load PEM file %q: %v\n", pemFile, err)
			os.Exit(1)
		}
		httpsSrv = &http.Server{
			Addr:    fmt.Sprintf("%s:%d", bind, tlsPort),
			Handler: handler,
			TLSConfig: &tls.Config{
				Certificates: []tls.Certificate{cert},
			},
		}
	}

	go func() {
		if err := httpSrv.ListenAndServe(); err != nil && err != http.ErrServerClosed {
			fmt.Fprintf(os.Stderr, "HTTP server error: %v\n", err)
			os.Exit(1)
		}
	}()
	fmt.Printf("Serving on http://%s:%d (default look=%s)\n", bind, port, look)

	if httpsSrv != nil {
		go func() {
			if err := httpsSrv.ListenAndServeTLS("", ""); err != nil && err != http.ErrServerClosed {
				fmt.Fprintf(os.Stderr, "HTTPS server error: %v\n", err)
				os.Exit(1)
			}
		}()
		fmt.Printf("Serving on https://%s:%d (default look=%s)\n", bind, tlsPort, look)
	}

	quit := make(chan os.Signal, 1)
	signal.Notify(quit, syscall.SIGINT, syscall.SIGTERM)
	<-quit

	fmt.Println("\nShutting down server")
	ctx := context.Background()
	httpSrv.Shutdown(ctx)
	if httpsSrv != nil {
		httpsSrv.Shutdown(ctx)
	}
}
