package main

import (
	"encoding/json"
	"errors"
	"fmt"
	"net/http"
	"net/http/httputil"
	"net/url"
	"os"
	"strings"
	"time"

	"git.cin.su/m/Assistant/runtime/supervisor"
)

type server struct {
	lifeboat       *supervisor.PersistentLifeboat
	stableUpstream string
	candidate      string
	client         *http.Client
}

func main() {
	statePath := envOr("SUPERVISOR_STATE_PATH", "/var/lib/supervisor/release-state.json")
	stable := strings.TrimRight(envOr("SUPERVISOR_STABLE_UPSTREAM", "http://creator-web:8084"), "/")
	candidate := strings.TrimRight(envOr("SUPERVISOR_CANDIDATE_UPSTREAM", "http://creator-web-candidate:8084"), "/")
	store, err := supervisor.NewDurableStore(statePath)
	if err != nil {
		fatal(err)
	}
	lifeboat := supervisor.NewPersistentLifeboat(store)
	if _, err := lifeboat.Snapshot(); errors.Is(err, os.ErrNotExist) {
		baseline := supervisor.Release{ID: envOr("SUPERVISOR_BASELINE_RELEASE", "bootstrap"), Version: envOr("SUPERVISOR_BASELINE_RELEASE", "bootstrap")}
		if ok, startErr := lifeboat.EstablishBaseline(baseline); startErr != nil || !ok {
			fatal(fmt.Errorf("initialize baseline: %v", startErr))
		}
		if ok, observeErr := lifeboat.Record(supervisor.Observation{ReleaseID: baseline.ID, Ready: true, Reason: "supervisor baseline"}); observeErr != nil || !ok {
			fatal(fmt.Errorf("initialize baseline observation: %v", observeErr))
		}
		if ok, trafficErr := lifeboat.SetTraffic(supervisor.TrafficStable, stable, candidate); trafficErr != nil || !ok {
			fatal(fmt.Errorf("initialize traffic: %v", trafficErr))
		}
	} else if err != nil {
		fatal(err)
	}

	s := &server{lifeboat: lifeboat, stableUpstream: stable, candidate: candidate, client: &http.Client{Timeout: 30 * time.Second}}
	public := http.NewServeMux()
	public.HandleFunc("GET /healthz", s.health)
	public.HandleFunc("GET /readyz", s.ready)
	public.HandleFunc("/", s.proxy)
	control := http.NewServeMux()
	control.HandleFunc("GET /v1/state", s.state)
	control.HandleFunc("POST /v1/releases/start", s.start)
	control.HandleFunc("POST /v1/releases/start-set", s.startSet)
	control.HandleFunc("POST /v1/releases/baseline", s.baseline)
	control.HandleFunc("POST /v1/releases/baseline-set", s.baselineSet)
	control.HandleFunc("POST /v1/releases/observe", s.observe)
	control.HandleFunc("POST /v1/releases/observe-set", s.observeSet)
	control.HandleFunc("POST /v1/releases/promote", s.promote)
	control.HandleFunc("POST /v1/releases/rollback", s.rollback)
	control.HandleFunc("POST /v1/traffic/canary", s.canary)
	control.HandleFunc("POST /v1/traffic/promote", s.promoteTraffic)
	control.HandleFunc("POST /v1/traffic/rollback", s.rollbackTraffic)

	go func() {
		if err := http.ListenAndServe(envOr("SUPERVISOR_CONTROL_ADDR", ":8095"), control); err != nil {
			fatal(err)
		}
	}()
	if err := http.ListenAndServe(envOr("LISTEN_ADDR", ":8080"), public); err != nil {
		fatal(err)
	}
}

func (s *server) health(w http.ResponseWriter, _ *http.Request) {
	writeJSON(w, map[string]string{"runtime": "supervisor", "status": "ok"})
}

func (s *server) ready(w http.ResponseWriter, r *http.Request) {
	state, err := s.lifeboat.Snapshot()
	if err != nil || state.StableUpstream == "" {
		http.Error(w, "supervisor state unavailable", http.StatusServiceUnavailable)
		return
	}
	if state.StableReleaseSet != nil && !s.checkReleaseSet(r, state.StableReleaseSet) {
		http.Error(w, "stable release set unavailable", http.StatusServiceUnavailable)
		return
	}
	request, err := http.NewRequestWithContext(r.Context(), http.MethodGet, state.StableUpstream+"/healthz", nil)
	if err != nil {
		http.Error(w, "stable upstream unavailable", http.StatusServiceUnavailable)
		return
	}
	response, err := s.client.Do(request)
	if err != nil || response.StatusCode/100 != 2 {
		if response != nil {
			response.Body.Close()
		}
		http.Error(w, "stable upstream unavailable", http.StatusServiceUnavailable)
		return
	}
	response.Body.Close()
	writeJSON(w, map[string]string{"runtime": "supervisor", "status": "ready"})
}

func (s *server) checkReleaseSet(r *http.Request, releaseSet *supervisor.ReleaseSet) bool {
	if releaseSet == nil || !releaseSet.Validate() {
		return false
	}
	for _, name := range releaseSet.ComponentNames() {
		component := releaseSet.Components[name]
		request, err := http.NewRequestWithContext(r.Context(), http.MethodGet, component.ReadyURL, nil)
		if err != nil {
			return false
		}
		response, err := s.client.Do(request)
		if err != nil || response.StatusCode/100 != 2 {
			if response != nil {
				response.Body.Close()
			}
			return false
		}
		response.Body.Close()
	}
	return true
}

func (s *server) proxy(w http.ResponseWriter, r *http.Request) {
	state, err := s.lifeboat.Snapshot()
	if err != nil {
		http.Error(w, "supervisor state unavailable", http.StatusServiceUnavailable)
		return
	}
	target := state.StableUpstream
	selected := "stable"
	if state.TrafficMode == supervisor.TrafficCandidate || (state.TrafficMode == supervisor.TrafficCanaryHeader && strings.EqualFold(r.Header.Get("X-Assistant-Canary"), "true")) {
		target = state.CandidateUpstream
		selected = "candidate"
	}
	parsed, err := url.Parse(target)
	if err != nil || parsed.Scheme == "" || parsed.Host == "" {
		http.Error(w, "invalid upstream", http.StatusServiceUnavailable)
		return
	}
	proxy := httputil.NewSingleHostReverseProxy(parsed)
	proxy.ModifyResponse = func(response *http.Response) error {
		response.Header.Set("X-Assistant-Traffic-Target", selected)
		return nil
	}
	proxy.ErrorHandler = func(writer http.ResponseWriter, _ *http.Request, _ error) {
		http.Error(writer, "upstream unavailable", http.StatusBadGateway)
	}
	proxy.ServeHTTP(w, r)
}

func (s *server) state(w http.ResponseWriter, _ *http.Request) {
	state, err := s.lifeboat.Snapshot()
	if err != nil {
		http.Error(w, "state unavailable", 503)
		return
	}
	writeJSON(w, state)
}
func (s *server) start(w http.ResponseWriter, r *http.Request) {
	var release supervisor.Release
	if !decode(w, r, &release) {
		return
	}
	ok, err := s.lifeboat.Start(release)
	respond(w, ok, err)
}
func (s *server) startSet(w http.ResponseWriter, r *http.Request) {
	var releaseSet supervisor.ReleaseSet
	if !decode(w, r, &releaseSet) {
		return
	}
	ok, err := s.lifeboat.StartSet(releaseSet)
	respond(w, ok, err)
}
func (s *server) baseline(w http.ResponseWriter, r *http.Request) {
	var release supervisor.Release
	if !decode(w, r, &release) {
		return
	}
	ok, err := s.lifeboat.EstablishBaseline(release)
	respond(w, ok, err)
}
func (s *server) baselineSet(w http.ResponseWriter, r *http.Request) {
	var releaseSet supervisor.ReleaseSet
	if !decode(w, r, &releaseSet) {
		return
	}
	ok, err := s.lifeboat.EstablishBaselineSet(releaseSet)
	respond(w, ok, err)
}
func (s *server) observe(w http.ResponseWriter, r *http.Request) {
	var observation supervisor.Observation
	if !decode(w, r, &observation) {
		return
	}
	ok, err := s.lifeboat.Record(observation)
	respond(w, ok, err)
}
func (s *server) observeSet(w http.ResponseWriter, r *http.Request) {
	var observation supervisor.ReleaseSetObservation
	if !decode(w, r, &observation) {
		return
	}
	ok, err := s.lifeboat.RecordSet(observation)
	respond(w, ok, err)
}
func (s *server) promote(w http.ResponseWriter, _ *http.Request) {
	state, snapshotErr := s.lifeboat.Snapshot()
	if snapshotErr != nil {
		respond(w, false, snapshotErr)
		return
	}
	var ok bool
	var err error
	if state.CandidateReleaseSet != nil {
		ok, err = s.lifeboat.PromoteSet()
	} else {
		ok, err = s.lifeboat.Promote()
	}
	respond(w, ok, err)
}
func (s *server) rollback(w http.ResponseWriter, _ *http.Request) {
	state, snapshotErr := s.lifeboat.Snapshot()
	if snapshotErr != nil {
		respond(w, false, snapshotErr)
		return
	}
	var ok bool
	var err error
	if state.CandidateReleaseSet != nil || state.RollbackTargetReleaseSet != nil {
		ok, err = s.lifeboat.RollbackSet()
	} else {
		ok, err = s.lifeboat.Rollback()
	}
	respond(w, ok, err)
}
func (s *server) canary(w http.ResponseWriter, _ *http.Request) {
	state, snapshotErr := s.lifeboat.Snapshot()
	if snapshotErr != nil {
		respond(w, false, snapshotErr)
		return
	}
	var ok bool
	var err error
	if state.CandidateReleaseSet != nil {
		ok, err = s.lifeboat.SetTrafficForSets(supervisor.TrafficCanaryHeader)
	} else {
		ok, err = s.lifeboat.SetTraffic(supervisor.TrafficCanaryHeader, s.stableUpstream, s.candidate)
	}
	respond(w, ok, err)
}
func (s *server) promoteTraffic(w http.ResponseWriter, _ *http.Request) {
	ok, err := s.lifeboat.Promote()
	respond(w, ok, err)
}
func (s *server) rollbackTraffic(w http.ResponseWriter, _ *http.Request) {
	ok, err := s.lifeboat.RollbackTraffic()
	respond(w, ok, err)
}

func decode(w http.ResponseWriter, r *http.Request, value any) bool {
	decoder := json.NewDecoder(http.MaxBytesReader(w, r.Body, 16<<10))
	decoder.DisallowUnknownFields()
	if err := decoder.Decode(value); err != nil {
		http.Error(w, "invalid control request", 400)
		return false
	}
	return true
}
func respond(w http.ResponseWriter, ok bool, err error) {
	if err != nil {
		http.Error(w, "control state unavailable", 503)
		return
	}
	if !ok {
		http.Error(w, "control transition rejected", 409)
		return
	}
	writeJSON(w, map[string]bool{"ok": true})
}
func writeJSON(w http.ResponseWriter, value any) {
	w.Header().Set("Content-Type", "application/json")
	_ = json.NewEncoder(w).Encode(value)
}
func envOr(name, fallback string) string {
	if value := strings.TrimSpace(os.Getenv(name)); value != "" {
		return value
	}
	return fallback
}
func fatal(err error) { fmt.Fprintln(os.Stderr, err); os.Exit(1) }