package main

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

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

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

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

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

	if sensor.Name() == "BME280" {
		sensors = append(sensors, haSensorDef{
			id:     "pressure",
			name:   "Pressure",
			unit:   "hPa",
			class:  "atmospheric_pressure",
			valTpl: "{{ 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)

	// Configure I2C0 bus
	i2c := machine.I2C0
	err := i2c.Configure(machine.I2CConfig{
		Frequency: 400 * machine.KHz,
		SDA:       machine.GPIO4,
		SCL:       machine.GPIO5,
	})
	if err != nil {
		panicErr("I2C configuration", err)
	}

	// Instantiate nil oled variable
	var oled *ssd1306.Device
	white := color.RGBA{R: 255, G: 255, B: 255, A: 255}

	var dummy [1]byte
	if err := i2c.ReadRegister(0x3C, 0x00, dummy[:]); err == nil {
		println("SSD1306 detected at 0x3C")
		dev := ssd1306.NewI2C(i2c)
		dev.Configure(ssd1306.Config{
			Address: 0x3C,
			Width:   128,
			Height:  64,
		})
		oled = dev

		oled.ClearBuffer()
		tinyfont.WriteLine(oled, &proggy.TinySZ8pt7b, 0, 16, "Booting...", white)
		oled.Display()
	} else {
		println("SSD1306 not found — running headless")
	}

	// Initialize the configured sensor
	sensor, err := InitSensor(i2c, SensorType)
	if err != nil {
		panicErr("Sensor init", err)
	}
	fmt.Printf("Started climate monitor using: %s\n", sensor.Name())

	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, sensor); err != nil {
		println("Warning: HA discovery failed:", err.Error())
	}

	for {
		t, p, h, hasPres, err := sensor.ReadValues()
		if err != nil {
			println(fmt.Sprintf("%s read error: %s", sensor.Name(), err.Error()))
		} else {
			var stateJSON string
			if hasPres {
				stateJSON = fmt.Sprintf(`{"temperature":%.2f,"humidity":%.2f,"pressure":%.2f}`, t, h, p)
			} else {
				stateJSON = fmt.Sprintf(`{"temperature":%.2f,"humidity":%.2f}`, t, h)
			}

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

			if oled != nil {
				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)

				if hasPres {
					tinyfont.WriteLine(oled, &proggy.TinySZ8pt7b, 0, 60, fmt.Sprintf("Pres: %.1f hPa", p), white)
				} else {
					tinyfont.WriteLine(oled, &proggy.TinySZ8pt7b, 0, 60, fmt.Sprintf("Sens: %s", sensor.Name()), white)
				}
				oled.Display()
			}
		}

		time.Sleep(UpdateInterval)
	}
}