package main

import (
	"fmt"
	"log"
	"time"

	"google.golang.org/protobuf/proto"

	uartframe "powerpod/gotool/uart"
	"powerpod/gotool/pb"
)

const (
	otaHostChunkSize  = 200
	otaFlashBlockSize = 4096
	otaPrepareTimeout = 120 * time.Second
	otaDefaultTimeout = 15 * time.Second
)

const (
	otaStPreparing = 1
	otaStReady     = 2
	otaStBlockAck  = 3
	otaStSuccess   = 4
	otaStFailed    = 5
)

// OTAProgress is pushed to the dashboard during web uploads.
type OTAProgress struct {
	Type    string `json:"type"` // always "ota_progress"
	Phase   string `json:"phase"` // preparing, ready, uploading, done, error
	Percent int    `json:"percent"`
	Message string `json:"message"`
	Bytes   uint32 `json:"bytes_written,omitempty"`
	Slot    uint32 `json:"target_slot,omitempty"`
}

type otaProgressFn func(OTAProgress)

func runOTAUpload(m *managedSerial, firmware []byte, onProgress otaProgressFn) error {
	m.mu.Lock()
	defer m.mu.Unlock()
	err := runOTAOnPortUnlocked(m, firmware, onProgress)
	if err != nil {
		m.invalidateLocked(err)
	}
	return err
}

func runOTAOnPortUnlocked(m *managedSerial, firmware []byte, onProgress otaProgressFn) error {
	if len(firmware) == 0 {
		return fmt.Errorf("empty firmware")
	}
	notify := func(phase string, percent int, msg string, extra ...OTAProgress) {
		if onProgress == nil {
			return
		}
		p := OTAProgress{Type: "ota_progress", Phase: phase, Percent: percent, Message: msg}
		if len(extra) > 0 {
			p.Bytes = extra[0].Bytes
			p.Slot = extra[0].Slot
		}
		onProgress(p)
	}

	if m.sp == nil {
		if err := m.openLocked(); err != nil {
			notify("error", 0, err.Error())
			return err
		}
	}

	sp := m.sp
	if err := sp.port.SetReadTimeout(otaPrepareTimeout); err != nil {
		notify("error", 0, err.Error())
		return err
	}
	defer sp.port.SetReadTimeout(readTimeout)

	notify("preparing", 0, fmt.Sprintf("OTA start (%d bytes)…", len(firmware)))

	if err := writeUartMessage(sp, &pb.UartMessage{
		Type: pb.MessageType_OTA_START,
		Payload: &pb.UartMessage_OtaStart{
			OtaStart: &pb.OtaStartPayload{TotalSize: uint32(len(firmware))},
		},
	}, false); err != nil {
		notify("error", 0, err.Error())
		return err
	}

	ready, err := waitOtaStatus(sp, otaStReady, otaPrepareTimeout, func(msg string) {
		notify("preparing", 2, msg)
	})
	if err != nil {
		notify("error", 0, err.Error())
		return err
	}
	notify("ready", 5, fmt.Sprintf("Ziel-Slot %d bereit", ready.GetTargetSlot()))

	if err := sp.port.SetReadTimeout(otaDefaultTimeout); err != nil {
		notify("error", 0, err.Error())
		return err
	}

	var seq uint32
	for offset := 0; offset < len(firmware); {
		bytesInBlock := 0
		for bytesInBlock < otaFlashBlockSize && offset < len(firmware) {
			n := otaHostChunkSize
			room := otaFlashBlockSize - bytesInBlock
			if n > room {
				n = room
			}
			if offset+n > len(firmware) {
				n = len(firmware) - offset
			}
			chunk := firmware[offset : offset+n]

			if err := writeUartMessage(sp, &pb.UartMessage{
				Type: pb.MessageType_OTA_PAYLOAD,
				Payload: &pb.UartMessage_OtaPayload{
					OtaPayload: &pb.OtaPayload{Seq: seq, Data: chunk},
				},
			}, false); err != nil {
				notify("error", 0, err.Error())
				return err
			}
			seq++
			offset += n
			bytesInBlock += n

			pct := 5 + (offset * 90 / len(firmware))
			notify("uploading", pct, fmt.Sprintf("%d / %d bytes", offset, len(firmware)))
		}

		if bytesInBlock == otaFlashBlockSize {
			st, err := waitOtaStatus(sp, otaStBlockAck, otaDefaultTimeout, nil)
			if err != nil {
				notify("error", 0, err.Error())
				return err
			}
			pct := 5 + (offset * 90 / len(firmware))
			notify("uploading", pct, fmt.Sprintf("Block geschrieben (%d bytes in flash)", st.GetBytesWritten()),
				OTAProgress{Bytes: st.GetBytesWritten()})
		}
	}

	if err := writeUartMessage(sp, &pb.UartMessage{
		Type: pb.MessageType_OTA_END,
		Payload: &pb.UartMessage_OtaEnd{
			OtaEnd: &pb.OtaEndPayload{},
		},
	}, false); err != nil {
		notify("error", 0, err.Error())
		return err
	}

	st, err := readOtaStatus(sp)
	if err != nil {
		notify("error", 0, err.Error())
		return err
	}
	if st.GetStatus() != otaStSuccess {
		err := fmt.Errorf("OTA failed: status=%d error=%d", st.GetStatus(), st.GetError())
		notify("error", 0, err.Error())
		return err
	}

	notify("done", 100, fmt.Sprintf("Erfolg — %d bytes auf Slot %d (Neustart)", st.GetBytesWritten(), st.GetTargetSlot()),
		OTAProgress{Bytes: st.GetBytesWritten(), Slot: st.GetTargetSlot()})
	return nil
}

func writeUartMessage(sp *serialPort, msg *pb.UartMessage, logFrame bool) error {
	frame, err := encodeUartMessage(msg)
	if err != nil {
		return err
	}
	if logFrame {
		log.Printf("sending %s (%d frame bytes)", msg.Type, len(frame))
	}
	_, err = sp.port.Write(frame)
	return err
}

func waitOtaStatus(sp *serialPort, want uint32, timeout time.Duration, onPreparing func(string)) (*pb.OtaStatusPayload, error) {
	deadline := time.Now().Add(timeout)
	for {
		if time.Now().After(deadline) {
			return nil, fmt.Errorf("timeout waiting for OTA status %d", want)
		}
		if err := sp.port.SetReadTimeout(time.Until(deadline)); err != nil {
			return nil, err
		}
		st, err := readOtaStatus(sp)
		if err != nil {
			return nil, err
		}
		switch st.GetStatus() {
		case want:
			return st, nil
		case otaStPreparing:
			if onPreparing != nil {
				onPreparing("Partition wird vorbereitet (~30s)…")
			}
		case otaStFailed:
			return nil, fmt.Errorf("OTA failed (error=%d)", st.GetError())
		}
	}
}

func readOtaStatus(sp *serialPort) (*pb.OtaStatusPayload, error) {
	payload, err := uartframe.ReadFrame(sp.port, nil)
	if err != nil {
		return nil, fmt.Errorf("read response: %w", err)
	}
	msg, err := decodeUartPayload(payload)
	if err != nil {
		return nil, err
	}
	if msg.GetType() != pb.MessageType_OTA_STATUS {
		return nil, fmt.Errorf("unexpected response type %v", msg.GetType())
	}
	st := msg.GetOtaStatus()
	if st == nil {
		return nil, fmt.Errorf("missing ota_status")
	}
	return st, nil
}

func encodeUartMessage(msg *pb.UartMessage) ([]byte, error) {
	body, err := proto.Marshal(msg)
	if err != nil {
		return nil, err
	}
	payload := append([]byte{byte(msg.Type)}, body...)
	return uartframe.EncodeFrame(payload)
}

func decodeUartPayload(payload []byte) (*pb.UartMessage, error) {
	if len(payload) == 0 {
		return nil, fmt.Errorf("empty response")
	}
	var msg pb.UartMessage
	if err := proto.Unmarshal(payload[1:], &msg); err != nil {
		return nil, err
	}
	msg.Type = pb.MessageType(payload[0])
	return &msg, nil
}