package main

import (
	"encoding/json"
	"fmt"
	"log"
	"net/http"
	"os"
	"sync"
	"time"
)

// ---------------------------------------------------------------------------
// FSM states
// ---------------------------------------------------------------------------

type State int

const (
	StateReady State = iota
	StateDraining
	StateDrainedWaitLB
	StateTerminating
	StateFailed
)

func (s State) String() string {
	switch s {
	case StateReady:
		return "READY"
	case StateDraining:
		return "DRAINING"
	case StateDrainedWaitLB:
		return "DRAINED_WAIT_LB"
	case StateTerminating:
		return "TERMINATING"
	case StateFailed:
		return "FAILED"
	default:
		return "UNKNOWN"
	}
}

// ---------------------------------------------------------------------------
// Transition history (ring buffer, max 100 entries)
// ---------------------------------------------------------------------------

type Transition struct {
	At     time.Time `json:"at"`
	From   string    `json:"from"`
	To     string    `json:"to"`
	Reason string    `json:"reason"`
}

const maxTransitions = 100

type ringBuffer struct {
	buf  [maxTransitions]Transition
	head int // next write index
	size int
}

func (r *ringBuffer) push(t Transition) {
	r.buf[r.head] = t
	r.head = (r.head + 1) % maxTransitions
	if r.size < maxTransitions {
		r.size++
	}
}

// snapshot returns entries in insertion order.
func (r *ringBuffer) snapshot() []Transition {
	out := make([]Transition, r.size)
	start := (r.head - r.size + maxTransitions) % maxTransitions
	for i := 0; i < r.size; i++ {
		out[i] = r.buf[(start+i)%maxTransitions]
	}
	return out
}

// ---------------------------------------------------------------------------
// FSM
// ---------------------------------------------------------------------------

type FSM struct {
	mu          sync.RWMutex
	state       State
	since       time.Time
	transitions ringBuffer
}

func newFSM() *FSM {
	return &FSM{
		state: StateReady,
		since: time.Now().UTC(),
	}
}

// advance attempts to move from expectedFrom → to. Returns 409 if wrong state.
func (f *FSM) advance(expectedFrom, to State, reason string) (State, bool) {
	f.mu.Lock()
	defer f.mu.Unlock()

	if f.state != expectedFrom {
		return f.state, false
	}
	f.doTransition(to, reason)
	return f.state, true
}

// fail moves any state → FAILED.
func (f *FSM) fail(reason string) State {
	f.mu.Lock()
	defer f.mu.Unlock()
	f.doTransition(StateFailed, reason)
	return f.state
}

// must be called with lock held.
func (f *FSM) doTransition(to State, reason string) {
	from := f.state
	f.state = to
	f.since = time.Now().UTC()
	t := Transition{At: f.since, From: from.String(), To: to.String(), Reason: reason}
	f.transitions.push(t)
	logTransition(t)
}

func (f *FSM) snapshot() (State, time.Time, []Transition) {
	f.mu.RLock()
	defer f.mu.RUnlock()
	return f.state, f.since, f.transitions.snapshot()
}

func (f *FSM) current() State {
	f.mu.RLock()
	defer f.mu.RUnlock()
	return f.state
}

// ---------------------------------------------------------------------------
// Logging helpers
// ---------------------------------------------------------------------------

func logTransition(t Transition) {
	fmt.Printf("ts=%s event=transition from=%s to=%s reason=%s\n",
		t.At.Format(time.RFC3339), t.From, t.To, t.Reason)
}

func logRequest(method, path, remote string, status int, dur time.Duration) {
	fmt.Printf("ts=%s level=info msg=request method=%s path=%s remote=%s status=%d duration_ms=%d\n",
		time.Now().UTC().Format(time.RFC3339), method, path, remote, status, dur.Milliseconds())
}

// ---------------------------------------------------------------------------
// HTTP helpers
// ---------------------------------------------------------------------------

type responseRecorder struct {
	http.ResponseWriter
	status int
}

func (r *responseRecorder) WriteHeader(code int) {
	r.status = code
	r.ResponseWriter.WriteHeader(code)
}

func logging(next http.HandlerFunc) http.HandlerFunc {
	return func(w http.ResponseWriter, r *http.Request) {
		start := time.Now()
		rec := &responseRecorder{ResponseWriter: w, status: http.StatusOK}
		next(rec, r)
		logRequest(r.Method, r.URL.Path, r.RemoteAddr, rec.status, time.Since(start))
	}
}

func writeJSON(w http.ResponseWriter, status int, v any) {
	w.Header().Set("Content-Type", "application/json")
	w.WriteHeader(status)
	_ = json.NewEncoder(w).Encode(v)
}

func methodGuard(method string, next http.HandlerFunc) http.HandlerFunc {
	return func(w http.ResponseWriter, r *http.Request) {
		if r.Method != method {
			http.Error(w, "method not allowed", http.StatusMethodNotAllowed)
			return
		}
		next(w, r)
	}
}

// ---------------------------------------------------------------------------
// Handlers
// ---------------------------------------------------------------------------

func handleRoot(fsm *FSM) http.HandlerFunc {
	return func(w http.ResponseWriter, r *http.Request) {
		if r.URL.Path != "/" {
			http.NotFound(w, r)
			return
		}
		fmt.Fprintf(w, "service-a-hc %s\n", fsm.current())
	}
}

// GET /health_check.html — HAProxy backend health
// READY, DRAINING → 200 "OK"
// DRAINED_WAIT_LB, TERMINATING, FAILED → 503 "DRAIN"
func handleHAProxyHealth(fsm *FSM) http.HandlerFunc {
	return func(w http.ResponseWriter, r *http.Request) {
		s := fsm.current()
		switch s {
		case StateReady, StateDraining:
			w.WriteHeader(http.StatusOK)
			fmt.Fprint(w, "OK")
		default:
			w.WriteHeader(http.StatusServiceUnavailable)
			fmt.Fprint(w, "DRAIN")
		}
	}
}

// GET /health — K8s readiness
// READY, DRAINING, DRAINED_WAIT_LB → 200
// TERMINATING, FAILED → 503
func handleReadiness(fsm *FSM) http.HandlerFunc {
	return func(w http.ResponseWriter, r *http.Request) {
		s := fsm.current()
		body := map[string]string{"state": s.String()}
		switch s {
		case StateReady, StateDraining, StateDrainedWaitLB:
			writeJSON(w, http.StatusOK, body)
		default:
			writeJSON(w, http.StatusServiceUnavailable, body)
		}
	}
}

// GET /live — K8s liveness
// All states except FAILED → 200
// FAILED → 500
func handleLiveness(fsm *FSM) http.HandlerFunc {
	return func(w http.ResponseWriter, r *http.Request) {
		s := fsm.current()
		body := map[string]string{"state": s.String()}
		if s == StateFailed {
			writeJSON(w, http.StatusInternalServerError, body)
		} else {
			writeJSON(w, http.StatusOK, body)
		}
	}
}

// GET /drain/status — debug full history
func handleDrainStatus(fsm *FSM) http.HandlerFunc {
	return func(w http.ResponseWriter, r *http.Request) {
		s, since, transitions := fsm.snapshot()
		writeJSON(w, http.StatusOK, map[string]any{
			"state":       s.String(),
			"since":       since.Format(time.RFC3339),
			"transitions": transitions,
		})
	}
}

// POST /drain/start — READY → DRAINING
func handleDrainStart(fsm *FSM) http.HandlerFunc {
	return func(w http.ResponseWriter, r *http.Request) {
		newState, ok := fsm.advance(StateReady, StateDraining, "preStop")
		if !ok {
			writeJSON(w, http.StatusConflict, map[string]string{
				"error": "not in READY state",
				"state": newState.String(),
			})
			return
		}
		writeJSON(w, http.StatusOK, map[string]string{"state": newState.String()})
	}
}

// POST /drain/lb-fail — DRAINING → DRAINED_WAIT_LB
func handleDrainLBFail(fsm *FSM) http.HandlerFunc {
	return func(w http.ResponseWriter, r *http.Request) {
		newState, ok := fsm.advance(StateDraining, StateDrainedWaitLB, "lb-fail")
		if !ok {
			writeJSON(w, http.StatusConflict, map[string]string{
				"error": "not in DRAINING state",
				"state": newState.String(),
			})
			return
		}
		writeJSON(w, http.StatusOK, map[string]string{"state": newState.String()})
	}
}

// POST /disable-readiness — DRAINED_WAIT_LB → TERMINATING
func handleDisableReadiness(fsm *FSM) http.HandlerFunc {
	return func(w http.ResponseWriter, r *http.Request) {
		newState, ok := fsm.advance(StateDrainedWaitLB, StateTerminating, "disable-readiness")
		if !ok {
			writeJSON(w, http.StatusConflict, map[string]string{
				"error": "not in DRAINED_WAIT_LB state",
				"state": newState.String(),
			})
			return
		}
		writeJSON(w, http.StatusOK, map[string]string{"state": newState.String()})
	}
}

// POST /fail — any → FAILED
func handleFail(fsm *FSM) http.HandlerFunc {
	return func(w http.ResponseWriter, r *http.Request) {
		s := fsm.fail("manual")
		writeJSON(w, http.StatusOK, map[string]string{"state": s.String()})
	}
}

// ---------------------------------------------------------------------------
// Main
// ---------------------------------------------------------------------------

func main() {
	addr := ":18180"
	if v := os.Getenv("HC_ADDR"); v != "" {
		addr = v
	}

	fsm := newFSM()

	mux := http.NewServeMux()

	mux.HandleFunc("/", logging(methodGuard(http.MethodGet, handleRoot(fsm))))
	mux.HandleFunc("/health_check.html", logging(methodGuard(http.MethodGet, handleHAProxyHealth(fsm))))
	mux.HandleFunc("/health", logging(methodGuard(http.MethodGet, handleReadiness(fsm))))
	mux.HandleFunc("/live", logging(methodGuard(http.MethodGet, handleLiveness(fsm))))
	mux.HandleFunc("/drain/status", logging(methodGuard(http.MethodGet, handleDrainStatus(fsm))))
	mux.HandleFunc("/drain/start", logging(methodGuard(http.MethodPost, handleDrainStart(fsm))))
	mux.HandleFunc("/drain/lb-fail", logging(methodGuard(http.MethodPost, handleDrainLBFail(fsm))))
	mux.HandleFunc("/disable-readiness", logging(methodGuard(http.MethodPost, handleDisableReadiness(fsm))))
	mux.HandleFunc("/fail", logging(methodGuard(http.MethodPost, handleFail(fsm))))

	fmt.Printf("ts=%s level=info msg=starting addr=%s state=%s\n",
		time.Now().UTC().Format(time.RFC3339), addr, fsm.current())

	srv := &http.Server{
		Addr:         addr,
		Handler:      mux,
		ReadTimeout:  5 * time.Second,
		WriteTimeout: 10 * time.Second,
		IdleTimeout:  60 * time.Second,
		ErrorLog:     log.New(os.Stderr, "", 0),
	}

	if err := srv.ListenAndServe(); err != nil {
		fmt.Fprintf(os.Stderr, "ts=%s level=error msg=%v\n",
			time.Now().UTC().Format(time.RFC3339), err)
		os.Exit(1)
	}
}