require 'i2c'
require 'cyw43'
require 'socket'

# -------------------------------------------------------------
# Configuration Loader
# -------------------------------------------------------------
begin
  require './config'
rescue LoadError
  puts "ERROR: config.rb not found."
  puts "Create config.rb on the device with your Wi-Fi and MQTT credentials."
  loop { sleep 2 }
end

STATE_TOPIC = "homeassistant/sensor/#{Config::NODE_ID}/state"

# -------------------------------------------------------------
# Minimal MQTT 3.1.1 Client
# -------------------------------------------------------------
class SimpleMQTT
  def initialize(host, port = 1883)
    @host = host
    @port = port
    @sock = nil
  end

  def connect(client_id, username = nil, password = nil)
    @sock = TCPSocket.new(@host, @port)

    # Flags: Clean session (bit 1) + optional credentials
    flags = 0x02
    flags |= 0x80 if username
    flags |= 0x40 if password

    # Variable header: Protocol Name (MQTT) + Level (4) + Flags + KeepAlive (60s)
    var_header = [0x00, 0x04, 0x4D, 0x51, 0x54, 0x54, 0x04, flags, 0x00, 0x3C].pack('C*')

    # Payload
    payload = encode_str(client_id)
    payload << encode_str(username) if username
    payload << encode_str(password) if password

    body = var_header + payload
    packet = [0x10].pack('C') + encode_len(body.bytesize) + body
    @sock.write(packet)

    # Read CONNACK (4 bytes: 0x20, 0x02, ack_flags, return_code)
    resp = @sock.read(4)
    if resp.nil? || resp.getbyte(0) != 0x20 || resp.getbyte(3) != 0x00
      code = resp ? resp.getbyte(3) : 'nil'
      raise "MQTT connection failed: code #{code}"
    end
  end

  def publish(topic, payload, retain = false)
    cmd = 0x30 | (retain ? 0x01 : 0x00)
    body = encode_str(topic) + payload.to_s
    packet = [cmd].pack('C') + encode_len(body.bytesize) + body
    @sock.write(packet)
  end

  def disconnect
    @sock.write([0xE0, 0x00].pack('C*')) if @sock
    @sock.close if @sock
  rescue
    nil
  end

  private

  def encode_str(str)
    [str.bytesize].pack('n') + str
  end

  def encode_len(len)
    encoded = []
    loop do
      digit = len % 128
      len /= 128
      digit |= 0x80 if len > 0
      encoded << digit
      break if len == 0
    end
    encoded.pack('C*')
  end
end

# -------------------------------------------------------------
# BME280 Driver (I2C)
# -------------------------------------------------------------
class BME280
  DEFAULT_ADDR = 0x76

  def initialize(i2c, addr = DEFAULT_ADDR)
    @i2c = i2c
    @addr = addr

    # Reset sensor
    @i2c.write(@addr, [0xE0, 0xB6].pack('C*'))
    sleep 0.1

    read_calibration_data

    # Humidity oversampling x1
    @i2c.write(@addr, [0xF2, 0x01].pack('C*'))
    # Pressure oversampling x1, Temp oversampling x1, Normal mode
    @i2c.write(@addr, [0xF4, 0x27].pack('C*'))
    # Standby 1000ms, filter off
    @i2c.write(@addr, [0xF5, 0xA0].pack('C*'))
  end

  def read_calibration_data
    calib = @i2c.read_from_mem(@addr, 0x88, 24)
    c = calib.unpack('S<s<s<S<s<s<S<s<s<')
    @dig_t1, @dig_t2, @dig_t3 = c[0], c[1], c[2]
    @dig_p1, @dig_p2, @dig_p3, @dig_p4, @dig_p5, @dig_p6 = c[3], c[4], c[5], c[6], c[7], c[8]

    @dig_h1 = @i2c.read_from_mem(@addr, 0xA1, 1).unpack1('C')
    h = @i2c.read_from_mem(@addr, 0xE1, 7).unpack('s<CBb<b<c')
    @dig_h2, @dig_h3 = h[0], h[1]
    e4 = @i2c.read_from_mem(@addr, 0xE4, 2).unpack('CC')
    @dig_h4 = (e4[0] << 4) | (e4[1] & 0x0F)
    @dig_h5 = (e4[1] >> 4) | (@i2c.read_from_mem(@addr, 0xE5, 1).unpack1('C') << 4)
    @dig_h6 = @i2c.read_from_mem(@addr, 0xE7, 1).unpack1('c')
  end

  def read_values
    raw = @i2c.read_from_mem(@addr, 0xF7, 8).unpack('C*')
    raw_p = ((raw[0] << 16) | (raw[1] << 8) | raw[2]) >> 4
    raw_t = ((raw[3] << 16) | (raw[4] << 8) | raw[5]) >> 4
    raw_h = (raw[6] << 8) | raw[7]

    # Temperature calculation (Celsius)
    var1 = (((raw_t >> 3) - (@dig_t1 << 1)) * @dig_t2) >> 11
    var2 = (((((raw_t >> 4) - @dig_t1) * ((raw_t >> 4) - @dig_t1)) >> 12) * @dig_t3) >> 14
    t_fine = var1 + var2
    temp = ((t_fine * 5 + 128) >> 8) / 100.0

    # Pressure calculation (hPa)
    p_var1 = t_fine - 128000
    p_var2 = p_var1 * p_var1 * @dig_p6 + ((p_var1 * @dig_p5) << 17) + (@dig_p4 << 35)
    p_var1 = ((p_var1 * p_var1 * @dig_p3) >> 8) + ((p_var1 * @dig_p2) << 12)
    p_var1 = (((1 << 47) + p_var1) * @dig_p1) >> 33

    pressure = if p_var1 == 0
                 0
               else
                 p = 1048576 - raw_p
                 p = (((p << 31) - p_var2) * 3125) / p_var1
                 p_var1 = (@dig_p9 * (p >> 13) * (p >> 13)) >> 25
                 p_var2 = (@dig_p8 * p) >> 19
                 ((p + p_var1 + p_var2) >> 8) + (@dig_p7 << 4)
                 (p / 256.0) / 100.0
               end

    # Humidity calculation (%)
    h_var = t_fine - 76800
    h_var = (((((raw_h << 14) - (@dig_h4 << 20) - (@dig_h5 * h_var)) + 16384) >> 15) *
             (((((((h_var * @dig_h6) >> 10) * (((h_var * @dig_h3) >> 11) + 32768)) >> 10) + 2097152) *
               @dig_h2 + 8192) >> 14))
    h_var = h_var - (((((h_var >> 15) * (h_var >> 15)) >> 7) * @dig_h1) >> 4)
    h_var = 0 if h_var < 0
    h_var = 419430400 if h_var > 419430400
    humidity = (h_var >> 12) / 1024.0

    {
      temperature: temp.round(2),
      pressure: pressure.round(2),
      humidity: humidity.round(2)
    }
  end
end

# -------------------------------------------------------------
# Home Assistant MQTT Discovery
# -------------------------------------------------------------
def register_ha_discovery(mqtt)
  device_info = "\"device\":{\"identifiers\":[\"#{Config::NODE_ID}\"],\"name\":\"Pico 2W Weather\",\"model\":\"Pico 2 W\",\"manufacturer\":\"Raspberry Pi\"}"

  configs = [
    {
      id: "temperature",
      name: "Temperature",
      unit: "°C",
      cls: "temperature",
      val_tpl: "{{ value_json.temperature }}"
    },
    {
      id: "humidity",
      name: "Humidity",
      unit: "%",
      cls: "humidity",
      val_tpl: "{{ value_json.humidity }}"
    },
    {
      id: "pressure",
      name: "Pressure",
      unit: "hPa",
      cls: "atmospheric_pressure",
      val_tpl: "{{ value_json.pressure }}"
    }
  ]

  configs.each do |c|
    topic = "homeassistant/sensor/#{Config::NODE_ID}/#{c[:id]}/config"
    payload = "{\"name\":\"#{c[:name]}\",\"state_topic\":\"#{STATE_TOPIC}\",\"unit_of_measurement\":\"#{c[:unit]}\",\"device_class\":\"#{c[:cls]}\",\"value_template\":\"#{c[:val_tpl]}\",\"unique_id\":\"#{Config::NODE_ID}_#{c[:id]}\",#{device_info}}"
    mqtt.publish(topic, payload, true)
  end
end

# -------------------------------------------------------------
# Hardware & Network Init
# -------------------------------------------------------------
puts "Connecting to Wi-Fi..."
CYW43.init
CYW43.enable_sta_mode
CYW43.connect_ap(Config::WIFI_SSID, Config::WIFI_PASS)

while !CYW43.tcpip_link_status_up?
  sleep 1
  print "."
end
puts "\nConnected! IP: #{CYW43.ip_address}"

# GP4 = SDA, GP5 = SCL
i2c = I2C.new(unit: :RP2040_I2C0, sda_pin: 4, scl_pin: 5, frequency: 100_000)
sensor = BME280.new(i2c)

mqtt = SimpleMQTT.new(Config::MQTT_BROKER, Config::MQTT_PORT)
puts "Connecting to MQTT broker at #{Config::MQTT_BROKER}..."
mqtt.connect(Config::CLIENT_ID, Config::MQTT_USER, Config::MQTT_PASS)
puts "Connected! Registering Home Assistant discovery..."

register_ha_discovery(mqtt)

# -------------------------------------------------------------
# Main Loop
# -------------------------------------------------------------
loop do
  data = sensor.read_values
  puts "Temp: #{data[:temperature]} °C | Hum: #{data[:humidity]} % | Press: #{data[:pressure]} hPa"

  payload = "{\"temperature\":#{data[:temperature]},\"humidity\":#{data[:humidity]},\"pressure\":#{data[:pressure]}}"
  mqtt.publish(STATE_TOPIC, payload)

  sleep 30
rescue => e
  puts "Error in telemetry loop: #{e.message}. Attempting reconnect..."
  sleep 5
  begin
    mqtt.disconnect
    mqtt.connect(Config::CLIENT_ID, Config::MQTT_USER, Config::MQTT_PASS)
  rescue
    nil
  end
end