import network, time, ujson
from machine import Pin, I2C
from umqtt.simple import MQTTClient
from i2c_lcd import I2cLcd   # ✅ LCD library

# ── WIFI ───────────────────────────────────────
SSID = "Wokwi-GUEST"
PASSWORD = ""

# ── MQTT ──────────────────────────────────────
BROKER = "192.168.162.64"
CLIENT_ID = "esp32-evm"

TOPIC_READY      = b"evm/ready"
TOPIC_VOTER      = b"evm/voter"
TOPIC_CANDIDATES = b"evm/candidates"
TOPIC_VOTE       = b"evm/vote"

# ── LCD DISPLAY (REPLACED OLED) ───────────────
i2c = I2C(0, scl=Pin(22), sda=Pin(21))
lcd = I2cLcd(i2c, 0x27, 4, 20)   # 20x4 LCD

def show(msg):
    lcd.clear()
    lines = msg.split("\n")
    for i, line in enumerate(lines[:4]):   # only 4 rows
        lcd.move_to(0, i)
        lcd.putstr(line[:20])              # max 20 chars

# ── KEYPAD ────────────────────────────────────
keys = [
    ['1','2','3','A'],
    ['4','5','6','B'],
    ['7','8','9','C'],
    ['*','0','#','D']
]

rows = [Pin(p, Pin.OUT) for p in (13,12,14,27)]
cols = [Pin(p, Pin.IN, Pin.PULL_UP) for p in (2,0,16,4)]

for r in rows:
    r.value(1)

def scan_key():
    for i, r in enumerate(rows):
        r.value(0)
        for j, c in enumerate(cols):
            if c.value() == 0:
                time.sleep_ms(40)  # debounce
                if c.value() == 0:
                    key = keys[i][j]
                    while c.value() == 0:
                        time.sleep_ms(5)
                    r.value(1)
                    return key
        r.value(1)
    return None

# ── STATE MACHINE ─────────────────────────────
WAIT_READY   = 0
ENTER_VOTER  = 1
ENTER_CONST  = 2
WAIT_CAND    = 3
SELECT_PARTY = 4

state = WAIT_READY

input_buffer = ""
voter_id = ""
constituency = ""
parties = []

# ── WIFI CONNECT ─────────────────────────────
def wifi_connect():
    wlan = network.WLAN(network.STA_IF)
    wlan.active(True)
    wlan.connect(SSID, PASSWORD)

    while not wlan.isconnected():
        time.sleep(0.5)

    print("WiFi Connected:", wlan.ifconfig())

# ── MQTT CALLBACK ────────────────────────────
def mqtt_callback(topic, msg):
    global state, parties

    print("MSG:", topic, msg)

    data = ujson.loads(msg)

    if topic == TOPIC_READY:
        state = ENTER_VOTER
        show("Enter Voter ID:")

    elif topic == TOPIC_CANDIDATES:
        parties = []

        for c in data["candidates"]:
            parties.append(c["party_name"])

        # LCD only 4 lines → show max 3 options
        text = "Select Party:\n"
        for i, p in enumerate(parties[:3]):
            text += "{}:{}\n".format(i+1, p[:15])

        show(text)
        state = SELECT_PARTY

# ── MQTT CONNECT ─────────────────────────────
def mqtt_connect():
    client = MQTTClient(CLIENT_ID, BROKER, port=1883)
    client.set_callback(mqtt_callback)

    while True:
        try:
            client.connect()
            break
        except:
            time.sleep(2)

    client.subscribe(TOPIC_READY)
    client.subscribe(TOPIC_CANDIDATES)

    return client

# ── SEND FUNCTIONS ───────────────────────────
def send_voter(client):
    data = {
        "voter_id": voter_id,
        "constituency": constituency
    }
    client.publish(TOPIC_VOTER, ujson.dumps(data))

def send_vote(client, party):
    data = {
        "voter_id": voter_id,
        "party_name": party,
        "constituency": constituency
    }
    client.publish(TOPIC_VOTE, ujson.dumps(data))

# ── KEY HANDLER ──────────────────────────────
def handle_key(k, client):
    global input_buffer, voter_id, constituency, state

    if k == '*':
        input_buffer = ""
        show("Cleared")
        return

    if k == '#':

        if state == ENTER_VOTER:
            voter_id = input_buffer
            input_buffer = ""
            state = ENTER_CONST
            show("Enter Const:")

        elif state == ENTER_CONST:
            constituency = input_buffer
            input_buffer = ""

            send_voter(client)
            show("Sending...")
            state = WAIT_CAND

        elif state == SELECT_PARTY:
            try:
                choice = int(input_buffer)

                if 1 <= choice <= len(parties):
                    send_vote(client, parties[choice-1])

                    show("Vote Sent!")
                    time.sleep(2)

                    state = ENTER_VOTER
                    show("Enter Voter ID:")

                else:
                    show("Invalid Option")

            except:
                show("Invalid Input")

            input_buffer = ""
        return

    # normal input
    input_buffer += k
    show(input_buffer)

# ── MAIN ─────────────────────────────────────
wifi_connect()
client = mqtt_connect()

show("Waiting server...")

while True:
    client.check_msg()

    key = scan_key()
    if key:
        print("KEY:", key)
        handle_key(key, client)

    time.sleep_ms(50)