EMS/main.go

package main

import (
	"context"
	"flag"
	"fmt"
	"log/slog"
	"net/http"
	"os"
	"os/signal"
	"path/filepath"
	"syscall"
	"time"

	"github.com/prometheus/client_golang/prometheus"
	"github.com/prometheus/client_golang/prometheus/promhttp"
	"github.com/tb/ems/internal/actuator"
	"github.com/tb/ems/internal/collector"
	"github.com/tb/ems/internal/config"
	"github.com/tb/ems/internal/engine"
	"github.com/tb/ems/internal/forecast"
	"github.com/tb/ems/internal/metrics"
	"github.com/tb/ems/internal/status"
	"github.com/tb/ems/internal/viessmann"
)

func main() {
	configPath := flag.String("config", "configs/ems-config.yaml", "Path to config file")
	dryRun := flag.Bool("dry-run", false, "Run without executing actions (log only)")
	testViessmann := flag.Bool("test-viessmann", false, "Test Viessmann API: read setpoint, round-trip write, +1°C delta, then restore")
	flag.Parse()

	// Load configuration
	cfg, err := config.Load(*configPath)
	if err != nil {
		fmt.Fprintf(os.Stderr, "failed to load config: %v\n", err)
		os.Exit(1)
	}

	// Setup structured logging
	logLevel := slog.LevelInfo
	switch cfg.EMS.LogLevel {
	case "debug":
		logLevel = slog.LevelDebug
	case "warn":
		logLevel = slog.LevelWarn
	case "error":
		logLevel = slog.LevelError
	}
	logger := slog.New(slog.NewJSONHandler(os.Stdout, &slog.HandlerOptions{
		Level: logLevel,
	}))

	// Viessmann client — needed for both normal operation and --test-viessmann
	var vc *viessmann.Client
	if cfg.Viessmann.InstallationID != "" && cfg.Viessmann.ClientID != "" {
		var err error
		vc, err = viessmann.NewClient(cfg.Viessmann, logger)
		if err != nil {
			logger.Warn("Viessmann client init failed, WW boost disabled", "error", err)
		} else {
			logger.Info("Viessmann client initialized")
		}
	}

	// --test-viessmann: run the API test and exit — do not start the control loop
	if *testViessmann {
		if vc == nil {
			fmt.Fprintln(os.Stderr, "ERROR: Viessmann credentials not configured or token failed to load")
			os.Exit(1)
		}
		os.Exit(runViessmannTest(vc))
	}

	logger.Info("starting EMS",
		"config", *configPath,
		"dry_run", *dryRun,
		"poll_interval", cfg.EMS.PollInterval,
		"listen_addr", cfg.EMS.ListenAddr,
	)

	// Initialize components
	coll := collector.NewCollector(cfg, logger)
	eng := engine.NewEngine(cfg, logger)
	fc := forecast.NewClient(cfg.Forecast, cfg.Strategic, logger)

	act := actuator.NewActuator(cfg, vc, logger)

	// Startup: attempt state recovery from Shelly read-back
	recoverCtx, recoverCancel := context.WithTimeout(context.Background(), 10*time.Second)
	defer recoverCancel()
	if shouldRecover(cfg.EMS.StateFile, cfg.EMS.RecoveryTimeoutParsed(), logger) {
		if states, err := act.ReadAllStates(recoverCtx); err != nil {
			logger.Warn("state recovery failed, starting with all consumers off", "error", err)
		} else {
			eng.RecoverState(states)
		}
	} else {
		logger.Info("state recovery skipped (downtime exceeded threshold or first run)")
	}

	// Prometheus metrics
	reg := prometheus.NewRegistry()
	m := metrics.NewMetrics(reg)

	// Status store (shared between HTTP handler and control loop)
	wwConfigured := cfg.Viessmann.InstallationID != ""
	statusStore := status.NewStore(*dryRun, wwConfigured, cfg.Strategic)

	// Metrics HTTP server
	mux := http.NewServeMux()
	mux.Handle("/metrics", promhttp.HandlerFor(reg, promhttp.HandlerOpts{}))
	mux.HandleFunc("/health", func(w http.ResponseWriter, r *http.Request) {
		w.WriteHeader(http.StatusOK)
		w.Write([]byte("ok"))
	})
	mux.HandleFunc("/override", overrideHandler(act, eng, logger))
	mux.HandleFunc("/", statusStore.Handler())

	srv := &http.Server{
		Addr:    cfg.EMS.ListenAddr,
		Handler: mux,
	}

	go func() {
		logger.Info("metrics server listening", "addr", cfg.EMS.ListenAddr)
		if err := srv.ListenAndServe(); err != http.ErrServerClosed {
			logger.Error("metrics server error", "error", err)
		}
	}()

	// Graceful shutdown
	ctx, cancel := context.WithCancel(context.Background())
	defer cancel()

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

	// Main control loop
	ticker := time.NewTicker(cfg.EMS.PollIntervalParsed())
	defer ticker.Stop()

	logger.Info("EMS control loop started")

	// Run once immediately
	runCycle(ctx, coll, eng, act, fc, m, statusStore, cfg, cfg.EMS.StateFile, logger, *dryRun)

	for {
		select {
		case <-ticker.C:
			runCycle(ctx, coll, eng, act, fc, m, statusStore, cfg, cfg.EMS.StateFile, logger, *dryRun)

		case sig := <-sigCh:
			logger.Info("received signal, shutting down", "signal", sig)
			cancel()

			shutdownCtx, shutdownCancel := context.WithTimeout(context.Background(), 5*time.Second)
			defer shutdownCancel()
			srv.Shutdown(shutdownCtx)
			return
		}
	}
}

func runCycle(
	ctx context.Context,
	coll *collector.Collector,
	eng *engine.Engine,
	act *actuator.Actuator,
	fc *forecast.Client,
	m *metrics.Metrics,
	store *status.Store,
	cfg *config.Config,
	stateFile string,
	logger *slog.Logger,
	dryRun bool,
) {
	now := time.Now()

	// Step 1: Collect current state from Prometheus
	state, err := coll.Collect(ctx)
	if err != nil {
		logger.Error("collection failed", "error", err)
		store.Update(collector.SystemState{}, nil, nil, nil, err)
		return
	}

	// Write heartbeat so next startup can assess downtime
	writeHeartbeat(stateFile, logger)

	// Step 1b: Read Shelly states and sync to engine (detects manual overrides).
	// Partial results are fine — unreachable devices are logged inside ReadAllStates.
	if shellyStates, err := act.ReadAllStates(ctx); err != nil {
		logger.Warn("all Shelly devices unreachable, skipping override detection", "error", err)
	} else {
		eng.SyncHardwareState(shellyStates, now, cfg.EMS.OverrideTimeoutParsed())
		store.SyncDeviceStates(shellyStates)
	}

	// Update metrics with current state
	m.GridPowerW.Set(state.GridPowerW)
	m.BatterySOC.Set(state.BatterySOC)
	m.PVProductionW.Set(state.PVProductionW)

	// Track energy flow (rough estimation based on 2-min intervals)
	intervalHours := 2.0 / 60.0
	if state.GridPowerW > 0 {
		m.GridImportKWh.Add(state.GridPowerW / 1000.0 * intervalHours)
	} else {
		m.GridExportKWh.Add(-state.GridPowerW / 1000.0 * intervalHours)
	}

	// Step 2: Fetch forecast (cached; only hits the API once per day)
	var fcResult *forecast.Result
	if r, err := fc.Today(ctx); err == nil && r.TotalKWh > 0 {
		fcResult = &r
	}

	// Step 3: Run decision engine
	wwBoostC := computeWWBoost(fcResult, cfg)
	decisionStart := time.Now()
	actions := eng.Decide(state, now, wwBoostC)
	m.DecisionDuration.Observe(time.Since(decisionStart).Seconds())

	// Update consumer state metrics
	m.UpdateConsumerStates(eng.ConsumerStates())

	// Update status page
	store.Update(state, actions, eng.Overrides(), fcResult, nil)
	store.SyncConsumerStates(eng.ConsumerStates())

	if len(actions) == 0 {
		logger.Debug("no actions this cycle",
			"grid_w", state.GridPowerW,
			"soc", state.BatterySOC,
		)
		return
	}

	// Step 4: Execute actions
	m.RecordActions(actions)

	if dryRun {
		for _, a := range actions {
			logger.Info("[DRY RUN] would execute",
				"consumer", a.Consumer,
				"turn_on", a.TurnOn,
				"reason", a.Reason,
			)
		}
		return
	}

	if err := act.Execute(ctx, actions); err != nil {
		logger.Error("execution failed", "error", err)
	}
}

// shouldRecover checks the heartbeat file to decide whether to read back
// Shelly states on startup. Returns false if the file is missing (first run)
// or older than the recovery timeout.
func shouldRecover(stateFile string, timeout time.Duration, logger *slog.Logger) bool {
	if stateFile == "" {
		return false
	}
	info, err := os.Stat(stateFile)
	if err != nil {
		return false // first run or file was cleaned up
	}
	age := time.Since(info.ModTime())
	if age > timeout {
		logger.Info("heartbeat too old, skipping state recovery",
			"age", age.Round(time.Second),
			"timeout", timeout,
		)
		return false
	}
	return true
}

// overrideHandler handles manual on/off requests from the status page.
// Switches the hardware immediately and informs the engine directly so
// the override lockout is applied without waiting for the next cycle.
func overrideHandler(act *actuator.Actuator, eng *engine.Engine, logger *slog.Logger) http.HandlerFunc {
	consumerByKey := map[string]engine.Consumer{
		"sg_ready":  engine.ConsumerSGReady,
		"wallbox_a": engine.ConsumerWallboxA,
		"wallbox_b": engine.ConsumerWallboxB,
	}

	return func(w http.ResponseWriter, r *http.Request) {
		if r.Method != http.MethodPost {
			http.Redirect(w, r, "/", http.StatusSeeOther)
			return
		}

		consumerKey := r.FormValue("consumer")
		stateVal := r.FormValue("state")
		durationStr := r.FormValue("duration")

		consumer, ok := consumerByKey[consumerKey]
		if !ok {
			http.Error(w, "unknown consumer", http.StatusBadRequest)
			return
		}

		duration, err := time.ParseDuration(durationStr)
		if err != nil || duration <= 0 {
			duration = time.Hour // safe default
		}

		turnOn := stateVal == "on"

		ctx, cancel := context.WithTimeout(r.Context(), 5*time.Second)
		defer cancel()

		action := engine.Action{
			Consumer: consumer,
			TurnOn:   turnOn,
			Reason:   fmt.Sprintf("manual override via web UI (%s)", duration),
		}
		if err := act.Execute(ctx, []engine.Action{action}); err != nil {
			logger.Error("web UI override failed", "consumer", consumerKey, "state", stateVal, "error", err)
			http.Error(w, "switch failed — check logs", http.StatusInternalServerError)
			return
		}

		// Inform engine immediately so override lockout is active before next cycle
		eng.ApplyOverride(consumer, turnOn, duration)

		logger.Info("web UI override executed",
			"consumer", consumerKey,
			"state", stateVal,
			"duration", duration,
		)
		http.Redirect(w, r, "/", http.StatusSeeOther)
	}
}

// computeWWBoost derives the WW temperature boost in °C from the daily PV forecast.
// Returns 0 if the forecast is unavailable or below the mid threshold.
func computeWWBoost(fc *forecast.Result, cfg *config.Config) float64 {
	if fc == nil || fc.TotalKWh == 0 {
		return 0
	}
	switch {
	case fc.TotalKWh >= cfg.Strategic.ForecastHighKWh:
		return cfg.Strategic.WWBoostHighC
	case fc.TotalKWh >= cfg.Strategic.ForecastMidKWh:
		return cfg.Strategic.WWBoostMidC
	default:
		return 0
	}
}

// runViessmannTest exercises the Viessmann API read/write path without starting
// the EMS control loop. It performs three stages:
//  1. Read current DHW setpoint (confirms auth + endpoint)
//  2. Write the same value back (confirms write path with zero net change)
//  3. Write setpoint +1°C, read back to confirm, then restore original
//
// Returns 0 on success, 1 on any failure.
func runViessmannTest(vc *viessmann.Client) int {
	ctx, cancel := context.WithTimeout(context.Background(), 120*time.Second)
	defer cancel()

	pass := func(format string, args ...any) { fmt.Printf("  PASS  "+format+"\n", args...) }
	fail := func(format string, args ...any) { fmt.Fprintf(os.Stderr, "  FAIL  "+format+"\n", args...) }
	step := func(format string, args ...any) { fmt.Printf("\n[test] "+format+"\n", args...) }

	fmt.Println("=== Viessmann API test ===")

	// Stage 1: read current setpoint
	step("Stage 1 — read current DHW setpoint")
	original, err := vc.GetDHWTemperature(ctx)
	if err != nil {
		fail("GetDHWTemperature: %v", err)
		return 1
	}
	pass("current setpoint = %.1f°C", original)

	// Stage 2: write same value back (round-trip, no observable effect)
	step("Stage 2 — write same value back (%.1f°C → %.1f°C, no net change)", original, original)
	if err := vc.SetDHWTemperature(ctx, original); err != nil {
		fail("SetDHWTemperature(%.1f): %v", original, err)
		return 1
	}
	readback, err := vc.GetDHWTemperature(ctx)
	if err != nil {
		fail("read-back after round-trip: %v", err)
		return 1
	}
	if readback != original {
		fail("round-trip mismatch: wrote %.1f, read back %.1f", original, readback)
		return 1
	}
	pass("round-trip confirmed (%.1f°C)", readback)

	// Stage 3: +1°C delta — heat pump won't notice (5°C hysteresis), but API change is visible.
	// The Viessmann cloud queues the write to the gateway; the GET endpoint returns the last
	// *confirmed* device value, so we wait up to 30s for the change to propagate.
	delta := original + 1
	step("Stage 3 — +1°C delta test (%.1f°C → %.1f°C)", original, delta)
	if err := vc.SetDHWTemperature(ctx, delta); err != nil {
		fail("SetDHWTemperature(%.1f): %v", delta, err)
		_ = vc.SetDHWTemperature(ctx, original)
		return 1
	}
	fmt.Println("  ...waiting 30s for cloud→gateway propagation...")
	time.Sleep(30 * time.Second)
	readback, err = vc.GetDHWTemperature(ctx)
	if err != nil {
		fail("read-back after +1°C: %v", err)
		_ = vc.SetDHWTemperature(ctx, original)
		return 1
	}
	if readback != delta {
		fail("+1°C mismatch: wrote %.1f, read back %.1f after 30s — gateway may be offline or propagation takes longer", delta, readback)
		_ = vc.SetDHWTemperature(ctx, original)
		return 1
	}
	pass("+1°C confirmed (%.1f°C) — check Viessmann app now if you want visual confirmation", readback)

	// Restore original
	step("Restore — writing original setpoint back (%.1f°C)", original)
	if err := vc.SetDHWTemperature(ctx, original); err != nil {
		fail("restore SetDHWTemperature(%.1f): %v", original, err)
		return 1
	}
	fmt.Println("  ...waiting 30s for cloud→gateway propagation...")
	time.Sleep(30 * time.Second)
	readback, err = vc.GetDHWTemperature(ctx)
	if err != nil {
		fail("read-back after restore: %v", err)
		return 1
	}
	if readback != original {
		fail("restore mismatch: wrote %.1f, read back %.1f", original, readback)
		return 1
	}
	pass("restored to %.1f°C", readback)

	fmt.Println("\n=== ALL STAGES PASSED — Viessmann API read/write confirmed ===")
	return 0
}

// writeHeartbeat updates the heartbeat file timestamp each cycle.
func writeHeartbeat(stateFile string, logger *slog.Logger) {
	if stateFile == "" {
		return
	}
	if err := os.MkdirAll(filepath.Dir(stateFile), 0755); err != nil {
		logger.Warn("could not create heartbeat directory", "error", err)
		return
	}
	if err := os.WriteFile(stateFile, []byte(time.Now().Format(time.RFC3339)), 0644); err != nil {
		logger.Warn("could not write heartbeat file", "error", err)
	}
}