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)")
	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,
	}))

	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)

	// Viessmann client — optional, only if credentials are configured
	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")
		}
	}

	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, 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.
// For Shelly consumers, it switches the hardware directly; the existing
// SyncHardwareState mechanism detects the change next cycle and applies
// the override lockout automatically.
func overrideHandler(act *actuator.Actuator, 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")

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

		turnOn := stateVal == "on"

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

		action := engine.Action{
			Consumer: consumer,
			TurnOn:   turnOn,
			Reason:   "manual override via web UI",
		}
		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
		}

		logger.Info("web UI override executed", "consumer", consumerKey, "state", stateVal)
		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
	}
}

// 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)
	}
}