package main

import (
	"context"
	"encoding/json"
	"errors"
	"fmt"
	"os"
	"path/filepath"
	"strings"
	"time"
)

const (
	firmwareBootStaged        = "STAGED"
	firmwareBootRebootPending = "REBOOT_PENDING"
	firmwareBootConfirmed     = "CONFIRMED"
	firmwareBootRolledBack    = "ROLLED_BACK"
)

type simulatedFirmwareBootState struct {
	Phase              string    `json:"phase"`
	CommandID          string    `json:"command_id"`
	ReleaseID          string    `json:"release_id"`
	TargetVersion      string    `json:"target_version"`
	PreviousVersion    string    `json:"previous_version"`
	Digest             string    `json:"sha256"`
	ActiveSlot         string    `json:"active_slot"`
	PreviousSlot       string    `json:"previous_slot"`
	TargetSlot         string    `json:"target_slot"`
	FailedSlot         string    `json:"failed_slot,omitempty"`
	BootAttemptID      string    `json:"boot_attempt_id,omitempty"`
	TerminalResultCode string    `json:"terminal_result_code,omitempty"`
	HealthDeadline     time.Time `json:"health_deadline,omitempty"`
	UpdatedAt          time.Time `json:"updated_at"`
}

func simulatedABFirmwareEnabled() (bool, error) {
	value := strings.ToLower(strings.TrimSpace(os.Getenv("SIM_FIRMWARE_AB_UPDATE")))
	if value == "" || value == "false" {
		return false, nil
	}
	if value != "true" {
		return false, errors.New("SIM_FIRMWARE_AB_UPDATE must be true or false")
	}
	statePath := strings.TrimSpace(os.Getenv("SIM_FIRMWARE_STATE_FILE"))
	if statePath == "" || !filepath.IsAbs(statePath) {
		return false, errors.New("SIM_FIRMWARE_STATE_FILE must be an absolute path when A/B firmware simulation is enabled")
	}
	evidencePath := strings.TrimSpace(os.Getenv("SIM_FIRMWARE_EVIDENCE_FILE"))
	if evidencePath != "" && !filepath.IsAbs(evidencePath) {
		return false, errors.New("SIM_FIRMWARE_EVIDENCE_FILE must be an absolute path")
	}
	return true, nil
}

func simulatedFirmwareHealthWindow() (time.Duration, error) {
	value := strings.TrimSpace(os.Getenv("SIM_FIRMWARE_BOOT_HEALTH_WINDOW"))
	if value == "" {
		return 2 * time.Minute, nil
	}
	duration, err := time.ParseDuration(value)
	if err != nil || duration < time.Second || duration > 30*time.Minute {
		return 0, errors.New("SIM_FIRMWARE_BOOT_HEALTH_WINDOW must be between 1s and 30m")
	}
	return duration, nil
}

func simulatedFirmwareInitialSlot() (string, error) {
	value := strings.ToUpper(env("SIM_FIRMWARE_ACTIVE_SLOT", "A"))
	if value != "A" && value != "B" {
		return "", errors.New("SIM_FIRMWARE_ACTIVE_SLOT must be A or B")
	}
	return value, nil
}

func inactiveFirmwareSlot(active string) string {
	if active == "A" {
		return "B"
	}
	return "A"
}

func loadSimulatedFirmwareBootState() (*simulatedFirmwareBootState, error) {
	path := strings.TrimSpace(os.Getenv("SIM_FIRMWARE_STATE_FILE"))
	content, err := os.ReadFile(path)
	if errors.Is(err, os.ErrNotExist) {
		return nil, nil
	}
	if err != nil {
		return nil, fmt.Errorf("read simulated firmware state: %w", err)
	}
	var state simulatedFirmwareBootState
	decoder := json.NewDecoder(strings.NewReader(string(content)))
	decoder.DisallowUnknownFields()
	if err := decoder.Decode(&state); err != nil {
		return nil, fmt.Errorf("decode simulated firmware state: %w", err)
	}
	if state.ActiveSlot != "A" && state.ActiveSlot != "B" {
		return nil, errors.New("simulated firmware state has an invalid active slot")
	}
	switch state.Phase {
	case firmwareBootStaged, firmwareBootRebootPending, firmwareBootConfirmed, firmwareBootRolledBack:
	default:
		return nil, errors.New("simulated firmware state has an invalid phase")
	}
	return &state, nil
}

func saveSimulatedFirmwareBootState(state simulatedFirmwareBootState) error {
	path := strings.TrimSpace(os.Getenv("SIM_FIRMWARE_STATE_FILE"))
	if !filepath.IsAbs(path) {
		return errors.New("SIM_FIRMWARE_STATE_FILE must be an absolute path")
	}
	state.UpdatedAt = time.Now().UTC()
	content, err := json.MarshalIndent(state, "", "  ")
	if err != nil {
		return err
	}
	file, err := os.OpenFile(path, os.O_CREATE|os.O_WRONLY|os.O_TRUNC, 0o600)
	if err != nil {
		return fmt.Errorf("open simulated firmware state: %w", err)
	}
	if _, err = file.Write(content); err == nil {
		err = file.Sync()
	}
	closeErr := file.Close()
	if err != nil {
		return fmt.Errorf("write simulated firmware state: %w", err)
	}
	if closeErr != nil {
		return fmt.Errorf("close simulated firmware state: %w", closeErr)
	}
	return nil
}

func (s *simulator) resumeABFirmwareCommand(ctx context.Context, command simulatedDeviceCommand, manifest simulatedFirmwareManifest, fault string, delay time.Duration) (bool, firmwareExecutionOutcome, error) {
	state, err := loadSimulatedFirmwareBootState()
	if err != nil || state == nil {
		return false, firmwareFailed, err
	}
	if state.Phase == firmwareBootConfirmed || state.Phase == firmwareBootRolledBack {
		if state.CommandID != command.ID {
			return false, firmwareCompleted, nil
		}
		if state.Phase == firmwareBootConfirmed {
			err := s.ackDeviceCommand(ctx, command.ID, "SUCCEEDED", map[string]any{
				"phase": "confirmed", "progress_percent": 100, "result_code": "SIMULATED_AB_BOOT_CONFIRMED", "installed_version": manifest.Version,
				"verified_sha256": manifest.Digest, "verified_signing_key_id": manifest.SigningKeyID, "active_slot": state.ActiveSlot, "previous_slot": state.PreviousSlot, "boot_attempt_id": state.BootAttemptID,
			})
			return true, firmwareCompleted, err
		}
		resultCode := state.TerminalResultCode
		if resultCode == "" {
			resultCode = "BOOT_HEALTH_FAILED_ROLLED_BACK"
		}
		err := s.ackDeviceCommand(ctx, command.ID, "FAILED", map[string]any{
			"phase": "rolled_back", "progress_percent": 95, "result_code": resultCode, "rollback_applied": true,
			"active_slot": state.ActiveSlot, "failed_slot": state.FailedSlot, "rollback_version": state.TargetVersion, "boot_attempt_id": state.BootAttemptID,
		})
		return true, firmwareFailed, err
	}
	if state.CommandID != command.ID || state.ReleaseID != manifest.ReleaseID || state.TargetVersion != manifest.Version || state.Digest != manifest.Digest {
		return true, firmwareFailed, s.ackDeviceCommand(ctx, command.ID, "FAILED", map[string]any{"phase": "rolled_back", "progress_percent": 0, "result_code": "LOCAL_UPDATE_STATE_CONFLICT", "rollback_applied": false})
	}
	if state.Phase == firmwareBootStaged {
		outcome, err := s.armSimulatedFirmwareBoot(ctx, command.ID, *state)
		return true, outcome, err
	}
	bootState := *state
	bootState.ActiveSlot = state.TargetSlot
	if err := recordSimulatedFirmwareEvidence(command, manifest, "boot_verifying", bootState, "BOOT_HEALTH_CHECK_STARTED"); err != nil {
		return true, firmwareFailed, err
	}
	if err := s.ackDeviceCommand(ctx, command.ID, "ACKNOWLEDGED", map[string]any{
		"phase": "boot_verifying", "progress_percent": 95, "result_code": "BOOT_HEALTH_CHECK_STARTED",
		"active_slot": state.TargetSlot, "previous_slot": state.PreviousSlot, "boot_attempt_id": state.BootAttemptID, "boot_health_deadline": state.HealthDeadline,
	}); err != nil {
		return true, firmwareFailed, err
	}
	if fault == "power_loss_during_boot" {
		if err := recordSimulatedFirmwareEvidence(command, manifest, "power_loss_during_boot", *state, "SIMULATED_POWER_LOSS"); err != nil {
			return true, firmwareFailed, err
		}
		return true, firmwareInterrupted, nil
	}
	rollbackCode := ""
	if fault == "boot_health_failed" {
		rollbackCode = "BOOT_HEALTH_FAILED_ROLLED_BACK"
	} else if !state.HealthDeadline.IsZero() && time.Now().After(state.HealthDeadline) {
		rollbackCode = "BOOT_HEALTH_TIMEOUT_ROLLED_BACK"
	}
	if rollbackCode != "" {
		failedSlot := state.TargetSlot
		state.Phase = firmwareBootRolledBack
		state.ActiveSlot = state.PreviousSlot
		state.FailedSlot = failedSlot
		state.TerminalResultCode = rollbackCode
		state.TargetSlot = ""
		state.TargetVersion = state.PreviousVersion
		if err := saveSimulatedFirmwareBootState(*state); err != nil {
			return true, firmwareFailed, err
		}
		if err := recordSimulatedFirmwareEvidence(command, manifest, "rolled_back", *state, rollbackCode); err != nil {
			return true, firmwareFailed, err
		}
		if err := s.ackDeviceCommand(ctx, command.ID, "FAILED", map[string]any{
			"phase": "rolled_back", "progress_percent": 95, "result_code": rollbackCode, "rollback_applied": true,
			"active_slot": state.ActiveSlot, "failed_slot": failedSlot, "rollback_version": state.PreviousVersion, "boot_attempt_id": state.BootAttemptID,
		}); err != nil {
			return true, firmwareFailed, err
		}
		return true, firmwareFailed, nil
	}
	if err := waitForContext(ctx, delay); err != nil {
		return true, firmwareInterrupted, err
	}
	confirmedSlot := state.TargetSlot
	state.Phase = firmwareBootConfirmed
	state.ActiveSlot = confirmedSlot
	state.TerminalResultCode = "SIMULATED_AB_BOOT_CONFIRMED"
	state.TargetSlot = ""
	state.PreviousVersion = state.TargetVersion
	if err := saveSimulatedFirmwareBootState(*state); err != nil {
		return true, firmwareFailed, err
	}
	if err := recordSimulatedFirmwareEvidence(command, manifest, "confirmed", *state, "SIMULATED_AB_BOOT_CONFIRMED"); err != nil {
		return true, firmwareFailed, err
	}
	if err := s.ackDeviceCommand(ctx, command.ID, "SUCCEEDED", map[string]any{
		"phase": "confirmed", "progress_percent": 100, "result_code": "SIMULATED_AB_BOOT_CONFIRMED", "installed_version": manifest.Version,
		"verified_sha256": manifest.Digest, "verified_signing_key_id": manifest.SigningKeyID, "active_slot": confirmedSlot, "previous_slot": state.PreviousSlot, "boot_attempt_id": state.BootAttemptID,
	}); err != nil {
		return true, firmwareFailed, err
	}
	return true, firmwareCompleted, nil
}

func (s *simulator) stageABFirmwareCommand(ctx context.Context, command simulatedDeviceCommand, manifest simulatedFirmwareManifest, fault string) (firmwareExecutionOutcome, error) {
	current, err := loadSimulatedFirmwareBootState()
	if err != nil {
		return firmwareFailed, err
	}
	activeSlot, err := simulatedFirmwareInitialSlot()
	if err != nil {
		return firmwareFailed, err
	}
	previousVersion := env("SIM_FIRMWARE_CURRENT_VERSION", "0.2.0-sim")
	if current != nil && (current.Phase == firmwareBootConfirmed || current.Phase == firmwareBootRolledBack) {
		activeSlot = current.ActiveSlot
		previousVersion = current.TargetVersion
	}
	state := simulatedFirmwareBootState{
		Phase: firmwareBootStaged, CommandID: command.ID, ReleaseID: manifest.ReleaseID, TargetVersion: manifest.Version,
		PreviousVersion: previousVersion, Digest: manifest.Digest, ActiveSlot: activeSlot, PreviousSlot: activeSlot, TargetSlot: inactiveFirmwareSlot(activeSlot),
	}
	if err := saveSimulatedFirmwareBootState(state); err != nil {
		return firmwareFailed, err
	}
	if err := recordSimulatedFirmwareEvidence(command, manifest, "staging", state, "INACTIVE_SLOT_STAGED"); err != nil {
		return firmwareFailed, err
	}
	if err := s.ackDeviceCommand(ctx, command.ID, "ACKNOWLEDGED", map[string]any{
		"phase": "staging", "progress_percent": 85, "result_code": "INACTIVE_SLOT_STAGED", "active_slot": state.ActiveSlot, "target_slot": state.TargetSlot,
	}); err != nil {
		return firmwareFailed, err
	}
	if fault == "power_loss_after_stage" {
		if err := recordSimulatedFirmwareEvidence(command, manifest, "power_loss_after_stage", state, "SIMULATED_POWER_LOSS"); err != nil {
			return firmwareFailed, err
		}
		return firmwareInterrupted, nil
	}
	return s.armSimulatedFirmwareBoot(ctx, command.ID, state)
}

func (s *simulator) armSimulatedFirmwareBoot(ctx context.Context, commandID string, state simulatedFirmwareBootState) (firmwareExecutionOutcome, error) {
	window, err := simulatedFirmwareHealthWindow()
	if err != nil {
		return firmwareFailed, err
	}
	state.Phase = firmwareBootRebootPending
	state.BootAttemptID = nonce()
	state.HealthDeadline = time.Now().UTC().Add(window)
	if err := saveSimulatedFirmwareBootState(state); err != nil {
		return firmwareFailed, err
	}
	if err := recordSimulatedFirmwareEvidence(simulatedDeviceCommand{ID: commandID}, simulatedFirmwareManifest{ReleaseID: state.ReleaseID, Version: state.TargetVersion, Digest: state.Digest}, "reboot_pending", state, "REBOOT_ARMED"); err != nil {
		return firmwareFailed, err
	}
	if err := s.ackDeviceCommand(ctx, commandID, "ACKNOWLEDGED", map[string]any{
		"phase": "reboot_pending", "progress_percent": 90, "result_code": "REBOOT_ARMED", "active_slot": state.ActiveSlot,
		"target_slot": state.TargetSlot, "boot_attempt_id": state.BootAttemptID, "boot_health_deadline": state.HealthDeadline,
	}); err != nil {
		return firmwareFailed, err
	}
	return firmwareRebootRequired, nil
}