# CoreS3-Lite / UIFlow 2.0: PC上のOllamaへ質問する
import M5
import network
import requests2
import time

WIFI_SSID = "YOUR_WIFI_SSID"
WIFI_PASSWORD = "YOUR_WIFI_PASSWORD"
OLLAMA_URL = "http://192.168.1.100:11434/api/chat"
MODEL = "gemma3:1b"
PROMPT = "What is a sensor? Answer in simple English in at most 40 words."


def connect_wifi():
    wlan = network.WLAN(network.STA_IF)
    wlan.active(True)
    if not wlan.isconnected():
        wlan.connect(WIFI_SSID, WIFI_PASSWORD)
    started = time.ticks_ms()
    while not wlan.isconnected():
        M5.update()
        if time.ticks_diff(time.ticks_ms(), started) > 30000:
            raise RuntimeError("Wi-Fi timeout")
        time.sleep_ms(100)
    return wlan


def ask_ollama():
    response = None
    try:
        response = requests2.post(
            OLLAMA_URL,
            json={
                "model": MODEL,
                "messages": [{"role": "user", "content": PROMPT}],
                "stream": False,
                "options": {"num_predict": 100, "num_ctx": 2048}
            },
            headers={"Content-Type": "application/json"},
            timeout=120
        )
        if response.status_code != 200:
            print("Server response:", response.text)
            raise RuntimeError("HTTP " + str(response.status_code))
        data = response.json()
        if not isinstance(data, dict):
            raise RuntimeError("Invalid JSON response")
        if data.get("error"):
            raise RuntimeError(str(data["error"]))
        message = data.get("message")
        if not isinstance(message, dict):
            raise RuntimeError("Missing message")
        answer = message.get("content")
        if not isinstance(answer, str) or not answer.strip():
            raise RuntimeError("Empty answer")
        return answer
    finally:
        if response is not None:
            response.close()


def show(title, text, button="ASK"):
    lcd = M5.Lcd
    lcd.fillScreen(0x101820)
    lcd.setTextColor(0xFFFFFF, 0x101820)
    lcd.setTextSize(1)
    lcd.drawString(title, 8, 8)
    lcd.drawString("Full response: WebTerminal", 8, 25)
    # 標準フォントで表示できるASCIIに限定。全文はターミナルへ。
    safe = "".join(c if 32 <= ord(c) <= 126 else " " for c in text)
    lines = [safe[i:i + 48] for i in range(0, len(safe), 48)]
    for index, line in enumerate(lines[:8]):
        lcd.drawString(line, 8, 48 + index * 17)
    if len(lines) > 8:
        lcd.drawString("... see WebTerminal", 8, 187)
    lcd.fillRect(8, 207, 304, 27, 0x285880)
    lcd.setTextColor(0xFFFFFF, 0x285880)
    lcd.drawString(button, 16, 216)


def main():
    M5.begin()
    M5.Lcd.setRotation(1)
    if (M5.Lcd.width(), M5.Lcd.height()) != (320, 240):
        raise RuntimeError("Expected 320x240 display")
    show("Ollama client", "Tap ASK to connect and send the question.")
    was_down = False
    while True:
        M5.update()
        down = M5.Touch.getCount() > 0
        if down and not was_down and M5.Touch.getY() >= 207:
            try:
                show("Connecting...", "Connecting to Wi-Fi", "WAIT")
                connect_wifi()
                show("Thinking...", PROMPT, "WAIT")
                print("Question:", PROMPT)
                answer = ask_ollama()
                print("Answer:", answer)
                show("Answer", answer)
            except Exception as error:
                print("Error:", error)
                show("Error", str(error), "RETRY")
            # 押しっぱなしでは再送しない。次は指を離してからタップ。
        was_down = down
        time.sleep_ms(30)


if __name__ == "__main__":
    try:
        main()
    except KeyboardInterrupt:
        pass
    except Exception as error:
        import sys
        sys.print_exception(error)