// ha-watchdog: list Home Assistant apps, pick which to watch, auto-restart them.
//
// Backend selection:
//   - SUPERVISOR_TOKEN set  -> Supervisor REST API (use when running as an app)
//   - otherwise             -> shells out to the `ha` CLI
//
// Usage:
//   ha-watchdog            interactive menu
//   ha-watchdog -daemon    headless: monitor the saved watchlist until killed
//   ha-watchdog -version   print version and build date
//
// Version info is injected at build time (see Makefile):
//   -ldflags "-X main.version=1.1.0 -X main.buildDate=2026-11-09T12:00:00Z"
package main

import (
	"bufio"
	"bytes"
	"context"
	"encoding/json"
	"errors"
	"flag"
	"fmt"
	"io"
	"log"
	"net/http"
	"os"
	"os/exec"
	"os/signal"
	"regexp"
	"runtime"
	"runtime/debug"
	"sort"
	"strconv"
	"strings"
	"sync"
	"syscall"
	"time"
)

// ---------- Version info ----------

var (
	version   = "1.1.0"
	buildDate = "unknown" // set via -ldflags; falls back to VCS time if available
)

func versionString() string {
	bd := buildDate
	if bd == "unknown" {
		if info, ok := debug.ReadBuildInfo(); ok {
			for _, s := range info.Settings {
				if s.Key == "vcs.time" && s.Value != "" {
					bd = s.Value
				}
			}
		}
	}
	return fmt.Sprintf("ha-watchdog %s (built %s, %s %s/%s)",
		version, bd, runtime.Version(), runtime.GOOS, runtime.GOARCH)
}

// ---------- Types ----------

type Addon struct {
	Slug            string `json:"slug"`
	Name            string `json:"name"`
	State           string `json:"state"`
	Version         string `json:"version"`
	VersionLatest   string `json:"version_latest"`
	UpdateAvailable bool   `json:"update_available"`
	Available       bool   `json:"available"`
	Detached        bool   `json:"detached"`
}

type envelope struct {
	Result  string `json:"result"`
	Message string `json:"message"`
	Data    struct {
		Addons []Addon `json:"addons"`
	} `json:"data"`
}

type Backend interface {
	List() ([]Addon, error)
	Start(slug string) error
	// Notify raises a Home Assistant persistent notification.
	// Backends that can't do this return an error; callers just log it.
	Notify(title, message string) error
}

// ---------- Supervisor API backend ----------

type apiBackend struct {
	base, token string
	client      *http.Client
}

func (a *apiBackend) do(method, path string, payload []byte) ([]byte, error) {
	var body io.Reader
	if payload != nil {
		body = bytes.NewReader(payload)
	}
	req, err := http.NewRequest(method, a.base+path, body)
	if err != nil {
		return nil, err
	}
	req.Header.Set("Authorization", "Bearer "+a.token)
	if payload != nil {
		req.Header.Set("Content-Type", "application/json")
	}
	resp, err := a.client.Do(req)
	if err != nil {
		return nil, err
	}
	defer resp.Body.Close()
	out, _ := io.ReadAll(resp.Body)
	if resp.StatusCode >= 300 {
		return out, fmt.Errorf("supervisor returned %s: %s", resp.Status, strings.TrimSpace(string(out)))
	}
	return out, nil
}

func (a *apiBackend) List() ([]Addon, error) {
	out, err := a.do("GET", "/addons", nil)
	if err != nil {
		return nil, err
	}
	return parseList(out)
}

func (a *apiBackend) Start(slug string) error {
	_, err := a.do("POST", "/addons/"+slug+"/start", nil)
	return err
}

// Notify calls the HA Core API through the Supervisor proxy.
// Requires `homeassistant_api: true` in the app's config.yaml.
func (a *apiBackend) Notify(title, message string) error {
	payload, _ := json.Marshal(map[string]string{
		"title":           title,
		"message":         message,
		"notification_id": "ha_watchdog",
	})
	_, err := a.do("POST", "/core/api/services/persistent_notification/create", payload)
	return err
}

// ---------- `ha` CLI backend ----------

type cliBackend struct{}

func (cliBackend) List() ([]Addon, error) {
	out, err := exec.Command("ha", "addons", "--raw-json").Output()
	if err != nil {
		return nil, fmt.Errorf("ha addons: %w", err)
	}
	return parseList(out)
}

func (cliBackend) Start(slug string) error {
	out, err := exec.Command("ha", "addons", "start", slug).CombinedOutput()
	if err != nil {
		return fmt.Errorf("ha addons start %s: %w: %s", slug, err, strings.TrimSpace(string(out)))
	}
	return nil
}

func (cliBackend) Notify(title, message string) error {
	return errors.New("notifications need the Supervisor API (run as an HA app with homeassistant_api: true)")
}

func parseList(b []byte) ([]Addon, error) {
	var env envelope
	if err := json.Unmarshal(b, &env); err != nil {
		return nil, err
	}
	if env.Result != "ok" {
		return nil, errors.New("supervisor error: " + env.Message)
	}
	sort.Slice(env.Data.Addons, func(i, j int) bool {
		return strings.ToLower(env.Data.Addons[i].Name) < strings.ToLower(env.Data.Addons[j].Name)
	})
	return env.Data.Addons, nil
}

// ---------- SSH (remote) backend ----------
//
// Runs the `ha` CLI on the Home Assistant host over SSH, so the watchdog can
// live on a separate monitoring machine. Notifications (optional) go through
// the normal Home Assistant REST API using a long-lived access token.

var slugRe = regexp.MustCompile(`^[A-Za-z0-9_-]+$`)

type sshBackend struct {
	target  string // user@host
	port    int
	key     string // optional identity file
	haURL   string // optional, e.g. http://homeassistant.local:8123
	haToken string // optional long-lived access token
	client  *http.Client
}

func (s *sshBackend) run(remoteCmd string) ([]byte, error) {
	args := []string{"-o", "BatchMode=yes", "-o", "ConnectTimeout=10", "-p", strconv.Itoa(s.port)}
	if s.key != "" {
		args = append(args, "-i", s.key)
	}
	args = append(args, s.target, remoteCmd)

	ctx, cancel := context.WithTimeout(context.Background(), 90*time.Second)
	defer cancel()
	cmd := exec.CommandContext(ctx, "ssh", args...)
	var stderr bytes.Buffer
	cmd.Stderr = &stderr
	out, err := cmd.Output()
	if err != nil {
		return nil, fmt.Errorf("ssh %s %q: %w: %s", s.target, remoteCmd, err, strings.TrimSpace(stderr.String()))
	}
	return out, nil
}

func (s *sshBackend) List() ([]Addon, error) {
	out, err := s.run("ha addons --raw-json")
	if err != nil {
		return nil, err
	}
	return parseList(out)
}

func (s *sshBackend) Start(slug string) error {
	if !slugRe.MatchString(slug) {
		return fmt.Errorf("refusing to run with unexpected slug %q", slug)
	}
	_, err := s.run("ha addons start " + slug)
	return err
}

func (s *sshBackend) Notify(title, message string) error {
	if s.haURL == "" || s.haToken == "" {
		return errors.New("remote notifications need -ha-url and -ha-token")
	}
	payload, _ := json.Marshal(map[string]string{
		"title":           title,
		"message":         message,
		"notification_id": "ha_watchdog",
	})
	req, err := http.NewRequest("POST", strings.TrimRight(s.haURL, "/")+"/api/services/persistent_notification/create", bytes.NewReader(payload))
	if err != nil {
		return err
	}
	req.Header.Set("Authorization", "Bearer "+s.haToken)
	req.Header.Set("Content-Type", "application/json")
	resp, err := s.client.Do(req)
	if err != nil {
		return err
	}
	defer resp.Body.Close()
	if resp.StatusCode >= 300 {
		body, _ := io.ReadAll(resp.Body)
		return fmt.Errorf("home assistant returned %s: %s", resp.Status, strings.TrimSpace(string(body)))
	}
	return nil
}

type backendConfig struct {
	sshTarget string
	sshPort   int
	sshKey    string
	haURL     string
	haToken   string
}

// newBackend picks a backend and returns it with a short description for logs.
func newBackend(c backendConfig) (Backend, string) {
	if c.sshTarget != "" {
		return &sshBackend{
			target: c.sshTarget, port: c.sshPort, key: c.sshKey,
			haURL: c.haURL, haToken: c.haToken,
			client: &http.Client{Timeout: 30 * time.Second},
		}, "remote over SSH (" + c.sshTarget + ")"
	}
	if tok := os.Getenv("SUPERVISOR_TOKEN"); tok != "" {
		return &apiBackend{base: "http://supervisor", token: tok, client: &http.Client{Timeout: 60 * time.Second}},
			"Supervisor API (local)"
	}
	return cliBackend{}, "ha CLI (local)"
}

// ---------- Watchlist (persisted) ----------

type Watchlist struct {
	mu    sync.Mutex
	path  string
	slugs map[string]bool
}

func loadWatchlist(path string) *Watchlist {
	w := &Watchlist{path: path, slugs: map[string]bool{}}
	data, err := os.ReadFile(path)
	if err != nil {
		return w
	}
	var list []string
	if json.Unmarshal(data, &list) == nil {
		for _, s := range list {
			w.slugs[s] = true
		}
	}
	return w
}

func (w *Watchlist) saveLocked() {
	list := make([]string, 0, len(w.slugs))
	for s := range w.slugs {
		list = append(list, s)
	}
	sort.Strings(list)
	data, _ := json.MarshalIndent(list, "", "  ")
	if err := os.WriteFile(w.path, data, 0o644); err != nil {
		log.Printf("warning: could not save watchlist: %v", err)
	}
}

func (w *Watchlist) Add(slug string) {
	w.mu.Lock()
	defer w.mu.Unlock()
	w.slugs[slug] = true
	w.saveLocked()
}

func (w *Watchlist) Remove(slug string) {
	w.mu.Lock()
	defer w.mu.Unlock()
	delete(w.slugs, slug)
	w.saveLocked()
}

func (w *Watchlist) Has(slug string) bool {
	w.mu.Lock()
	defer w.mu.Unlock()
	return w.slugs[slug]
}

func (w *Watchlist) Slugs() []string {
	w.mu.Lock()
	defer w.mu.Unlock()
	list := make([]string, 0, len(w.slugs))
	for s := range w.slugs {
		list = append(list, s)
	}
	sort.Strings(list)
	return list
}

// ---------- Monitor ----------

type Monitor struct {
	be         Backend
	wl         *Watchlist
	interval   time.Duration
	cooldown   time.Duration // minimum gap between start attempts for one app
	maxRetries int           // give up after this many attempts within retryWindow
	window     time.Duration

	mu       sync.Mutex
	attempts map[string][]time.Time // recent start attempts per slug
	gaveUp   map[string]bool        // apps we've stopped retrying
	skipped  map[string]bool        // apps we've already logged as unstartable
}

func newMonitor(be Backend, wl *Watchlist, interval, cooldown time.Duration, maxRetries int, window time.Duration) *Monitor {
	return &Monitor{
		be: be, wl: wl, interval: interval, cooldown: cooldown,
		maxRetries: maxRetries, window: window,
		attempts: map[string][]time.Time{},
		gaveUp:   map[string]bool{},
		skipped:  map[string]bool{},
	}
}

func (m *Monitor) notify(title, msg string) {
	if err := m.be.Notify(title, msg); err != nil {
		log.Printf("notify: %v", err)
	}
}

// recentAttempts drops expired attempts and returns the remaining list. Caller holds m.mu.
func (m *Monitor) recentAttempts(slug string) []time.Time {
	cutoff := time.Now().Add(-m.window)
	var keep []time.Time
	for _, t := range m.attempts[slug] {
		if t.After(cutoff) {
			keep = append(keep, t)
		}
	}
	m.attempts[slug] = keep
	return keep
}

func (m *Monitor) check() {
	addons, err := m.be.List()
	if err != nil {
		log.Printf("monitor: list failed: %v", err)
		return
	}
	state := map[string]Addon{}
	for _, a := range addons {
		state[a.Slug] = a
	}

	for _, slug := range m.wl.Slugs() {
		a, ok := state[slug]
		if !ok {
			m.mu.Lock()
			first := !m.skipped[slug]
			m.skipped[slug] = true
			m.mu.Unlock()
			if first {
				log.Printf("monitor: %s is on the watchlist but not installed, skipping", slug)
			}
			continue
		}

		// Can't start these, so don't try.
		if !a.Available || a.Detached {
			m.mu.Lock()
			first := !m.skipped[slug]
			m.skipped[slug] = true
			m.mu.Unlock()
			if first {
				log.Printf("monitor: %s (%s) is unavailable or detached, skipping", a.Name, slug)
			}
			continue
		}

		m.mu.Lock()
		delete(m.skipped, slug)

		if a.State == "started" {
			// Recovered. Allow a future give-up notification, but keep the attempt
			// history so a flapping app still trips the cap quickly.
			delete(m.gaveUp, slug)
			m.mu.Unlock()
			continue
		}
		if m.gaveUp[slug] {
			m.mu.Unlock()
			continue
		}

		recent := m.recentAttempts(slug)
		if len(recent) >= m.maxRetries {
			m.gaveUp[slug] = true
			m.mu.Unlock()
			msg := fmt.Sprintf("%s (%s) is %q and failed to stay up after %d start attempts in %s. "+
				"Watchdog has stopped retrying. Check the app's logs.",
				a.Name, slug, a.State, m.maxRetries, m.window)
			log.Printf("monitor: GIVING UP on %s: %s", slug, msg)
			m.notify("Watchdog gave up on "+a.Name, msg)
			continue
		}
		if len(recent) > 0 && time.Since(recent[len(recent)-1]) < m.cooldown {
			m.mu.Unlock()
			continue
		}
		m.attempts[slug] = append(recent, time.Now())
		n := len(m.attempts[slug])
		m.mu.Unlock()

		log.Printf("monitor: %s (%s) is %q, starting it (attempt %d/%d)", a.Name, slug, a.State, n, m.maxRetries)
		if err := m.be.Start(slug); err != nil {
			log.Printf("monitor: start %s failed: %v", slug, err)
		} else {
			m.notify("Watchdog restarted "+a.Name,
				fmt.Sprintf("%s was %q and has been started again (attempt %d/%d).", a.Name, a.State, n, m.maxRetries))
		}
	}
}

// Reset clears give-up state and attempt history for all apps.
func (m *Monitor) Reset() {
	m.mu.Lock()
	defer m.mu.Unlock()
	m.attempts = map[string][]time.Time{}
	m.gaveUp = map[string]bool{}
}

// Status returns slugs the monitor has given up on.
func (m *Monitor) GaveUp() []string {
	m.mu.Lock()
	defer m.mu.Unlock()
	var out []string
	for s := range m.gaveUp {
		out = append(out, s)
	}
	sort.Strings(out)
	return out
}

func (m *Monitor) Run(ctx context.Context) {
	m.check()
	t := time.NewTicker(m.interval)
	defer t.Stop()
	for {
		select {
		case <-ctx.Done():
			return
		case <-t.C:
			m.check()
		}
	}
}

// ---------- Interactive UI ----------

func listApps(be Backend, wl *Watchlist) []Addon {
	addons, err := be.List()
	if err != nil {
		fmt.Println("Error:", err)
		return nil
	}
	fmt.Printf("\n%-4s %-30s %-9s %-16s %-16s %-7s %s\n", "#", "Name", "State", "Version", "Update", "Watched", "Notes")
	for i, a := range addons {
		upd := ""
		if a.UpdateAvailable {
			upd = "-> " + a.VersionLatest
		}
		watched := ""
		if wl.Has(a.Slug) {
			watched = "yes"
		}
		var notes []string
		if !a.Available {
			notes = append(notes, "unavailable")
		}
		if a.Detached {
			notes = append(notes, "detached")
		}
		fmt.Printf("%-4d %-30.30s %-9s %-16.16s %-16.16s %-7s %s\n",
			i+1, a.Name, a.State, a.Version, upd, watched, strings.Join(notes, ","))
	}
	return addons
}

func parseNumbers(s string, max int) []int {
	var out []int
	for _, f := range strings.FieldsFunc(s, func(r rune) bool { return r == ',' || r == ' ' }) {
		n, err := strconv.Atoi(f)
		if err != nil || n < 1 || n > max {
			fmt.Printf("Ignoring %q\n", f)
			continue
		}
		out = append(out, n-1)
	}
	return out
}

func interactive(be Backend, wl *Watchlist, m *Monitor) {
	in := bufio.NewScanner(os.Stdin)
	prompt := func(msg string) string {
		fmt.Print(msg)
		if !in.Scan() {
			os.Exit(0)
		}
		return strings.TrimSpace(in.Text())
	}

	var cancel context.CancelFunc
	running := func() bool { return cancel != nil }

	for {
		status := "OFF"
		if running() {
			status = "ON"
		}
		fmt.Printf("\n=== HA App Watchdog %s | monitor: %s | %d watched ===\n", version, status, len(wl.Slugs()))
		if g := m.GaveUp(); len(g) > 0 {
			fmt.Printf("!! Gave up on: %s (use option 6 to reset)\n", strings.Join(g, ", "))
		}
		fmt.Println("1) List installed apps")
		fmt.Println("2) Add apps to watchlist")
		fmt.Println("3) Remove apps from watchlist")
		fmt.Println("4) Show watchlist")
		fmt.Println("5) Start/stop monitoring")
		fmt.Println("6) Reset retry counters / give-up state")
		fmt.Println("v) Version info")
		fmt.Println("q) Quit")

		switch prompt("> ") {
		case "1":
			listApps(be, wl)

		case "2":
			addons := listApps(be, wl)
			if addons == nil {
				continue
			}
			for _, i := range parseNumbers(prompt("Numbers to add (e.g. 1,4,7): "), len(addons)) {
				wl.Add(addons[i].Slug)
				fmt.Println("Added", addons[i].Name)
			}

		case "3":
			slugs := wl.Slugs()
			if len(slugs) == 0 {
				fmt.Println("Watchlist is empty.")
				continue
			}
			for i, s := range slugs {
				fmt.Printf("%d) %s\n", i+1, s)
			}
			for _, i := range parseNumbers(prompt("Numbers to remove: "), len(slugs)) {
				wl.Remove(slugs[i])
				fmt.Println("Removed", slugs[i])
			}

		case "4":
			slugs := wl.Slugs()
			if len(slugs) == 0 {
				fmt.Println("Watchlist is empty.")
			}
			for i, s := range slugs {
				fmt.Printf("%d) %s\n", i+1, s)
			}

		case "5":
			if running() {
				cancel()
				cancel = nil
				fmt.Println("Monitoring stopped.")
			} else {
				var ctx context.Context
				ctx, cancel = context.WithCancel(context.Background())
				go m.Run(ctx)
				fmt.Printf("Monitoring started (every %s, max %d starts per %s). Logs print here.\n",
					m.interval, m.maxRetries, m.window)
			}

		case "6":
			m.Reset()
			fmt.Println("Retry counters and give-up state cleared.")

		case "v", "V":
			fmt.Println(versionString())

		case "q", "Q":
			if running() {
				cancel()
			}
			return
		}
	}
}

// ---------- main ----------

func defaultWatchlistPath() string {
	if p := os.Getenv("WATCHDOG_FILE"); p != "" {
		return p
	}
	if st, err := os.Stat("/data"); err == nil && st.IsDir() {
		return "/data/watchlist.json" // persistent storage inside an HA app
	}
	return "watchlist.json"
}

func main() {
	daemon := flag.Bool("daemon", false, "run headless, monitoring the saved watchlist")
	showVersion := flag.Bool("version", false, "print version and build date, then exit")
	interval := flag.Duration("interval", 30*time.Second, "how often to check app state")
	cooldown := flag.Duration("cooldown", 60*time.Second, "minimum time between start attempts for one app")
	maxRetries := flag.Int("max-retries", 5, "give up on an app after this many start attempts within -retry-window")
	window := flag.Duration("retry-window", 10*time.Minute, "time window for counting start attempts")
	file := flag.String("file", defaultWatchlistPath(), "watchlist file path")
	sshTarget := flag.String("ssh", "", "monitor a remote Home Assistant host over SSH, e.g. root@homeassistant.local")
	sshPort := flag.Int("ssh-port", 22, "SSH port on the Home Assistant host")
	sshKey := flag.String("ssh-key", "", "SSH identity file (private key) to use")
	haURL := flag.String("ha-url", "", "Home Assistant URL for remote notifications, e.g. http://homeassistant.local:8123")
	haToken := flag.String("ha-token", os.Getenv("HA_TOKEN"), "long-lived access token for remote notifications (or set HA_TOKEN)")
	flag.Parse()

	if *showVersion {
		fmt.Println(versionString())
		return
	}

	be, backendDesc := newBackend(backendConfig{
		sshTarget: *sshTarget, sshPort: *sshPort, sshKey: *sshKey,
		haURL: *haURL, haToken: *haToken,
	})
	wl := loadWatchlist(*file)
	m := newMonitor(be, wl, *interval, *cooldown, *maxRetries, *window)

	if *daemon {
		ctx, stop := signal.NotifyContext(context.Background(), os.Interrupt, syscall.SIGTERM)
		defer stop()
		log.Printf("%s", versionString())
		log.Printf("backend: %s", backendDesc)
		log.Printf("daemon mode: watching %v every %s (max %d starts per %s)",
			wl.Slugs(), *interval, *maxRetries, *window)
		m.Run(ctx)
		return
	}
	fmt.Println("Backend:", backendDesc)
	interactive(be, wl, m)
}