package main

import (
	"bytes"
	"context"
	"encoding/binary"
	"fmt"
	"io"
	"machine"
	"strconv"
	"time"
	"image/color"

	"tinygo.org/x/drivers/ssd1306"
	"tinygo.org/x/tinyfont"
	"tinygo.org/x/tinyfont/proggy"
)

var StateTopic = fmt.Sprintf("homeassistant/sensor/%s/state", NodeID)

// -------------------------------------------------------------
// Minimal MQTT 3.1.1 Client
// -------------------------------------------------------------
type SimpleMQTT struct {
	rw io.ReadWriter
}

func NewSimpleMQTT(rw io.ReadWriter) *SimpleMQTT {
	return &SimpleMQTT{rw: rw}
}

func (m *SimpleMQTT) Connect(clientID, username, password string) error {
	flags := byte(0x02) // Clean session
	if username != "" {
		flags |= 0x80
	}
	if password != "" {
		flags |= 0x40
	}

	// Variable header: MQTT (proto name), Level 4 (MQTT 3.1.1), flags, KeepAlive 60s
	varHeader := []byte{0x00, 0x04, 'M', 'Q', 'T', 'T', 0x04, flags, 0x00, 0x3C}

	payload := encodeString(clientID)
	if username != "" {
		payload = append(payload, encodeString(username)...)
	}
	if password != "" {
		payload = append(payload, encodeString(password)...)
	}

	body := append(varHeader, payload...)
	packet := append([]byte{0x10}, encodeLength(len(body))...)
	packet = append(packet, body...)

	if _, err := m.rw.Write(packet); err != nil {
		return err
	}

	resp := make([]byte, 4)
	if _, err := io.ReadFull(m.rw, resp); err != nil {
		return err
	}
	if resp[0] != 0x20 || resp[3] != 0x00 {
		return fmt.Errorf("MQTT connection refused: code %d", resp[3])
	}
	return nil
}

func (m *SimpleMQTT) Publish(topic string, payload []byte, retain bool) error {
	cmd := byte(0x30)
	if retain {
		cmd |= 0x01
	}

	body := append(encodeString(topic), payload...)
	packet := append([]byte{cmd}, encodeLength(len(body))...)
	packet = append(packet, body...)

	_, err := m.rw.Write(packet)
	return err
}

func encodeString(s string) []byte {
	b := []byte(s)
	length := uint16(len(b))
	return append([]byte{byte(length >> 8), byte(length & 0xFF)}, b...)
}

func encodeLength(length int) []byte {
	var encoded []byte
	for {
		digit := byte(length % 128)
		length /= 128
		if length > 0 {
			digit |= 0x80
		}
		encoded = append(encoded, digit)
		if length == 0 {
			break
		}
	}
	return encoded
}

// -------------------------------------------------------------
// BME280 Driver
// -------------------------------------------------------------
const (
	BME280Addr  = 0x77
	regCalib00  = 0x88
	regCalib26  = 0xE1
	regReset    = 0xE0
	regCtrlHum  = 0xF2
	regCtrlMeas = 0xF4
	regConfig   = 0xF5
	regData     = 0xF7
)

type BME280Calib struct {
	digT1 uint16
	digT2 int16
	digT3 int16
	digP1 uint16
	digP2 int16
	digP3 int16
	digP4 int16
	digP5 int16
	digP6 int16
	digP7 int16
	digP8 int16
	digP9 int16
	digH1 uint8
	digH2 int16
	digH3 uint8
	digH4 int16
	digH5 int16
	digH6 int8
}

type BME280 struct {
	bus   *machine.I2C
	addr  uint8
	calib BME280Calib
}

func NewBME280(bus *machine.I2C, addr uint8) (*BME280, error) {
	b := &BME280{bus: bus, addr: addr}
	if err := bus.WriteRegister(addr, regReset, []byte{0xB6}); err != nil {
		return nil, err
	}
	time.Sleep(100 * time.Millisecond)

	if err := b.readCalibration(); err != nil {
		return nil, err
	}
	if err := bus.WriteRegister(addr, regCtrlHum, []byte{0x01}); err != nil {
		return nil, err
	}
	if err := bus.WriteRegister(addr, regCtrlMeas, []byte{0x27}); err != nil {
		return nil, err
	}
	if err := bus.WriteRegister(addr, regConfig, []byte{0xA0}); err != nil {
		return nil, err
	}
	return b, nil
}

func (b *BME280) readCalibration() error {
	var buf24 [24]byte
	if err := b.bus.ReadRegister(b.addr, regCalib00, buf24[:]); err != nil {
		return err
	}
	r := bytes.NewReader(buf24[:])
	binary.Read(r, binary.LittleEndian, &b.calib.digT1)
	binary.Read(r, binary.LittleEndian, &b.calib.digT2)
	binary.Read(r, binary.LittleEndian, &b.calib.digT3)
	binary.Read(r, binary.LittleEndian, &b.calib.digP1)
	binary.Read(r, binary.LittleEndian, &b.calib.digP2)
	binary.Read(r, binary.LittleEndian, &b.calib.digP3)
	binary.Read(r, binary.LittleEndian, &b.calib.digP4)
	binary.Read(r, binary.LittleEndian, &b.calib.digP5)
	binary.Read(r, binary.LittleEndian, &b.calib.digP6)
	binary.Read(r, binary.LittleEndian, &b.calib.digP7)
	binary.Read(r, binary.LittleEndian, &b.calib.digP8)
	binary.Read(r, binary.LittleEndian, &b.calib.digP9)

	var h1 [1]byte
	if err := b.bus.ReadRegister(b.addr, 0xA1, h1[:]); err != nil {
		return err
	}
	b.calib.digH1 = h1[0]

	var hBuf [7]byte
	if err := b.bus.ReadRegister(b.addr, regCalib26, hBuf[:]); err != nil {
		return err
	}
	b.calib.digH2 = int16(binary.LittleEndian.Uint16(hBuf[0:2]))
	b.calib.digH3 = hBuf[2]
	b.calib.digH4 = (int16(hBuf[3]) << 4) | (int16(hBuf[4]) & 0x0F)
	b.calib.digH5 = (int16(hBuf[5]) << 4) | (int16(hBuf[4]) >> 4)
	b.calib.digH6 = int8(hBuf[6])
	return nil
}

func (b *BME280) ReadValues() (float64, float64, float64, error) {
	var raw [8]byte
	if err := b.bus.ReadRegister(b.addr, regData, raw[:]); err != nil {
		return 0, 0, 0, err
	}

	rawP := (int32(raw[0]) << 12) | (int32(raw[1]) << 4) | (int32(raw[2]) >> 4)
	rawT := (int32(raw[3]) << 12) | (int32(raw[4]) << 4) | (int32(raw[5]) >> 4)
	rawH := (int32(raw[6]) << 8) | int32(raw[7])

	var1 := (((rawT >> 3) - (int32(b.calib.digT1) << 1)) * int32(b.calib.digT2)) >> 11
	var2 := (((((rawT >> 4) - int32(b.calib.digT1)) * ((rawT >> 4) - int32(b.calib.digT1))) >> 12) * int32(b.calib.digT3)) >> 14
	tFine := var1 + var2
	temp := float64((tFine*5+128)>>8) / 100.0

	pVar1 := int64(tFine) - 128000
	pVar2 := pVar1*pVar1*int64(b.calib.digP6) + ((pVar1 * int64(b.calib.digP5)) << 17) + (int64(b.calib.digP4) << 35)
	pVar1 = ((pVar1 * pVar1 * int64(b.calib.digP3)) >> 8) + ((pVar1 * int64(b.calib.digP2)) << 12)
	pVar1 = (((int64(1) << 47) + pVar1) * int64(b.calib.digP1)) >> 33

	var pres float64
	if pVar1 != 0 {
		p := int64(1048576 - rawP)
		p = (((p << 31) - pVar2) * 3125) / pVar1
		pVar1 = (int64(b.calib.digP9) * (p >> 13) * (p >> 13)) >> 25
		pVar2 = (int64(b.calib.digP8) * p) >> 19
		p = ((p + pVar1 + pVar2) >> 8) + (int64(b.calib.digP7) << 4)
		pres = (float64(p) / 256.0) / 100.0
	}

	hVar := tFine - 76800
	hVar = (((((rawH << 14) - (int32(b.calib.digH4) << 20) - (int32(b.calib.digH5) * hVar)) + 16384) >> 15) *
		(((((((hVar * int32(b.calib.digH6)) >> 10) * (((hVar * int32(b.calib.digH3)) >> 11) + 32768)) >> 10) + 2097152)*
			int32(b.calib.digH2) + 8192) >> 14))
	hVar = hVar - (((((hVar >> 15) * (hVar >> 15)) >> 7) * int32(b.calib.digH1)) >> 4)
	if hVar < 0 {
		hVar = 0
	} else if hVar > 419430400 {
		hVar = 419430400
	}
	hum := float64(hVar>>12) / 1024.0

	return temp, pres, hum, nil
}

// -------------------------------------------------------------
// Home Assistant Auto-Discovery
// -------------------------------------------------------------
type haSensorDef struct {
	id     string
	name   string
	unit   string
	class  string
	valTpl string
}

func registerHADiscovery(mqtt *SimpleMQTT) error {
	deviceJSON := fmt.Sprintf(`"device":{"identifiers":["%s"],"name":"%s","model":"Pico 2 W","manufacturer":"Raspberry Pi"}`,
		NodeID, NodeName)

	sensors := []haSensorDef{
		{"temperature", "Temperature", "°C", "temperature", "{{ value_json.temperature }}"},
		{"humidity", "Humidity", "%", "humidity", "{{ value_json.humidity }}"},
		{"pressure", "Pressure", "hPa", "atmospheric_pressure", "{{ value_json.pressure }}"},
	}

	for _, s := range sensors {
		topic := fmt.Sprintf("homeassistant/sensor/%s/%s/config", NodeID, s.id)
		payload := fmt.Sprintf(`{"name":"%s","has_entity_name":true,"unique_id":"%s_%s","device_class":"%s","state_class":"measurement","unit_of_measurement":"%s","state_topic":"%s","value_template":"%s",%s}`,
			s.name, NodeID, s.id, s.class, s.unit, StateTopic, s.valTpl, deviceJSON)

		if err := mqtt.Publish(topic, []byte(payload), true); err != nil {
			return err
		}
		println("Registered HA discovery:", s.id)
	}
	return nil
}

// -------------------------------------------------------------
// Panic Handler & Entry Point
// -------------------------------------------------------------
func panicErr(msg string, err error) {
	if err != nil {
		println("FATAL:", msg, "-", err.Error())
	}
	for {
		time.Sleep(500 * time.Millisecond)
	}
}

func main() {
	time.Sleep(2 * time.Second)

	// Bump frequency to 400kHz for snappy display flushes
	i2c := machine.I2C0
	err := i2c.Configure(machine.I2CConfig{
		Frequency: 400 * machine.KHz,
		SDA:       machine.GPIO16,
		SCL:       machine.GPIO17,
	})
	if err != nil {
		panicErr("I2C configuration", err)
	}

	// Initialize OLED (shares the same I2C bus at 0x3C)
	oled := ssd1306.NewI2C(i2c)
	oled.Configure(ssd1306.Config{
		Address: 0x3C, // Check 0x3D if 0x3C doesn't respond
		Width:   128,
		Height:  64,
	})
	oled.ClearBuffer()
	white := color.RGBA{R: 255, G: 255, B: 255, A: 255}
	tinyfont.WriteLine(oled, &proggy.TinySZ8pt7b, 0, 16, "Booting...", white)
	oled.Display()

	bme, err := NewBME280(i2c, BME280Addr)
	if err != nil {
		panicErr("BME280 init", err)
	}

	netStack, err := InitNetwork(WifiSSID, WifiPass)
	if err != nil {
		panicErr("network init", err)
	}

	portNum, _ := strconv.Atoi(MQTTPort)
	ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
	defer cancel()

	conn, err := netStack.DialTCP(ctx, MQTTHost, uint16(portNum))
	if err != nil {
		panicErr("TCP dial", err)
	}

	mqtt := NewSimpleMQTT(conn)
	println("Connecting MQTT session...")
	if err := mqtt.Connect(ClientID, MQTTUser, MQTTPass); err != nil {
		panicErr("MQTT connect", err)
	}
	println("MQTT connected.")

	if err := registerHADiscovery(mqtt); err != nil {
		println("Warning: HA discovery failed:", err.Error())
	}

	for {
		t, p, h, err := bme.ReadValues()
		if err != nil {
			println("BME280 read error:", err.Error())
		} else {
			// Publish MQTT
			stateJSON := fmt.Sprintf(`{"temperature":%.2f,"humidity":%.2f,"pressure":%.2f}`, t, h, p)
			if err := mqtt.Publish(StateTopic, []byte(stateJSON), false); err != nil {
				println("MQTT publish error:", err.Error())
			} else {
				println(fmt.Sprintf("[%s] Published: %s", NodeName, stateJSON))
			}

			// Render to OLED
			oled.ClearBuffer()
			tinyfont.WriteLine(oled, &proggy.TinySZ8pt7b, 0, 14, fmt.Sprintf("Node: %s", NodeName), white)
			tinyfont.WriteLine(oled, &proggy.TinySZ8pt7b, 0, 30, fmt.Sprintf("Temp: %.1f C", t), white)
			tinyfont.WriteLine(oled, &proggy.TinySZ8pt7b, 0, 46, fmt.Sprintf("Hum:  %.1f %%", h), white)
			tinyfont.WriteLine(oled, &proggy.TinySZ8pt7b, 0, 60, fmt.Sprintf("Pres: %.1f hPa", p), white)
			oled.Display()
		}

		time.Sleep(UpdateInterval)
	}
}