#!/usr/bin/env python3
"""
LXST Whisplay Client - Group voice chat for Raspberry Pi Zero 2W
with PiSugar Whisplay HAT.

Connects to lxst-relay voice channels via Reticulum/LXST.
Uses the Whisplay HAT for LCD display, button (PTT), RGB LED,
and WM8960 audio I/O.
"""

import RNS
import os
import sys
import time
import threading
import argparse
import numpy as np
import signal as signal_module

from LXST import Pipeline
from LXST.Network import SignallingReceiver, Packetizer, LinkSource
from LXST.Sinks import LineSink
from LXST.Sources import LineSource
from LXST.Codecs import Opus
from LXST.Primitives.Telephony import Signalling

script_dir = os.path.dirname(os.path.abspath(__file__))
whisplay_runtime = os.path.abspath(os.path.join(script_dir, "..", "Whisplay", "runtime"))
if whisplay_runtime not in sys.path:
    sys.path.insert(0, whisplay_runtime)

from lib.whisplay_client import create_whisplay_hardware

HAS_PIL = False
try:
    from PIL import Image, ImageDraw, ImageFont
    HAS_PIL = True
except ImportError:
    pass

APP_NAME = "lxst"
PRIMITIVE_NAME = "telephony"

STATE_IDLE = 0
STATE_CONNECTING = 1
STATE_CONNECTED = 2
STATE_TRANSMITTING = 3
STATE_FAILED = 4

STATUS_STR = {
    STATE_IDLE: "Disconnected",
    STATE_CONNECTING: "Connecting...",
    STATE_CONNECTED: "Connected",
    STATE_TRANSMITTING: "TX",
    STATE_FAILED: "Failed",
}

ACCENT_IDLE = (60, 150, 255)
ACCENT_CONNECTING = (255, 180, 0)
ACCENT_CONNECTED = (0, 200, 80)
ACCENT_TX = (255, 50, 50)
ACCENT_FAILED = (200, 50, 50)

RGB_IDLE = (0, 0, 100)
RGB_CONNECTING = (100, 80, 0)
RGB_CONNECTED = (0, 100, 0)
RGB_TX = (255, 0, 0)
RGB_FAILED = (100, 0, 0)

WHISPLAY_APP_ID = "lxst-whisplay-client"
DISPLAY_NAME = "LXST Voice"

class LXSTWhisplayClient(SignallingReceiver):
    def __init__(self, channel_hash, identity_path=None, rns_config=None, verbosity=0):
        super().__init__()
        self.channel_hash = channel_hash
        self.state = STATE_IDLE
        self.link = None
        self.link_source = None
        self.packetizer = None
        self.line_source = None
        self.line_sink = None
        self.tx_pipeline = None
        self.participant_count = 0
        self._exit = False
        self._cancel_connect = False
        self._lock = threading.Lock()
        self._press_time = 0
        self._press_state = STATE_IDLE
        self._error_msg = ""

        self._init_rns(identity_path, rns_config, verbosity)
        self._init_whisplay()
        self._init_fonts()

    def _init_rns(self, identity_path, rns_config, verbosity):
        if identity_path and os.path.isfile(identity_path):
            self.identity = RNS.Identity.from_file(identity_path)
            if self.identity is None:
                RNS.log("Could not load identity, generating new", RNS.LOG_ERROR)
                self.identity = RNS.Identity()
                if identity_path:
                    self.identity.to_file(identity_path)
        else:
            RNS.log("No identity found, generating new", RNS.LOG_NOTICE)
            self.identity = RNS.Identity()
            if identity_path:
                self.identity.to_file(identity_path)

        self.reticulum = RNS.Reticulum(configdir=rns_config, loglevel=3 + verbosity)

        self.destination = RNS.Destination(
            self.identity,
            RNS.Destination.IN,
            RNS.Destination.SINGLE,
            APP_NAME,
            PRIMITIVE_NAME,
        )
        self.destination.set_proof_strategy(RNS.Destination.PROVE_NONE)

        RNS.log(
            f"LXST Client identity: {RNS.prettyhexrep(self.identity.hash)}",
            RNS.LOG_NOTICE,
        )

    def _init_whisplay(self):
        self.board = create_whisplay_hardware(
            app_id=WHISPLAY_APP_ID,
            display_name=DISPLAY_NAME,
            icon="LX",
            use_daemon_default_log=True,
        )
        self.board.set_backlight(70)
        self.board.on_button_press(self._on_button_press)
        self.board.on_button_release(self._on_button_release)
        if hasattr(self.board, "on_exit_request"):
            self.board.on_exit_request(self._on_exit_request)
        if hasattr(self.board, "on_focus_revoked"):
            self.board.on_focus_revoked(self._on_focus_revoked)

        if not HAS_PIL:
            self.board.fill_screen(0)

    def _init_fonts(self):
        self._fonts = {}
        if not HAS_PIL:
            return
        for size, bold in [(22, True), (18, False), (14, False), (12, False)]:
            self._fonts[(size, bold)] = self._load_font(size, bold)

    def _load_font(self, size, bold=False):
        candidates = []
        if bold:
            candidates.extend([
                "/usr/share/fonts/truetype/dejavu/DejaVuSans-Bold.ttf",
            ])
        candidates.extend([
            "/usr/share/fonts/truetype/dejavu/DejaVuSans.ttf",
            "/usr/share/fonts/truetype/freefont/FreeSans.ttf",
        ])
        for path in candidates:
            if os.path.exists(path):
                try:
                    return ImageFont.truetype(path, size)
                except Exception:
                    continue
        return ImageFont.load_default()

    def _wrap_text(self, draw, text, font, max_width):
        if not text:
            return [""]
        words = text.split()
        if not words:
            return [text]
        lines = []
        current = words[0]
        for word in words[1:]:
            cand = f"{current} {word}"
            bbox = draw.textbbox((0, 0), cand, font=font)
            if (bbox[2] - bbox[0]) <= max_width:
                current = cand
            else:
                lines.append(current)
                current = word
        lines.append(current)
        return lines

    def _rgb565_bytes(self, image):
        rgb = image.convert("RGB")
        w, h = rgb.width, rgb.height
        buf = bytearray(w * h * 2)
        i = 0
        for y in range(h):
            for x in range(w):
                r, g, b = rgb.getpixel((x, y))
                v = ((r & 0xF8) << 8) | ((g & 0xFC) << 3) | (b >> 3)
                buf[i] = (v >> 8) & 0xFF
                buf[i + 1] = v & 0xFF
                i += 2
        return bytes(buf)

    def _render_frame(self, accent=None):
        if not HAS_PIL:
            return None
        w = self.board.LCD_WIDTH
        h = self.board.LCD_HEIGHT
        accent = accent or ACCENT_IDLE
        bg = (11, 16, 24)

        img = Image.new("RGB", (w, h), bg)
        draw = ImageDraw.Draw(img)

        font_title = self._fonts.get((22, True), ImageFont.load_default())
        font_body = self._fonts.get((18, False), ImageFont.load_default())
        font_small = self._fonts.get((14, False), ImageFont.load_default())
        font_tiny = self._fonts.get((12, False), ImageFont.load_default())

        draw.rounded_rectangle(
            (12, 16, w - 12, h - 16), radius=18,
            fill=(20, 28, 40), outline=(40, 55, 72), width=2,
        )

        draw.rounded_rectangle(
            (20, 22, w - 20, 58), radius=12, fill=accent,
        )

        state_str = STATUS_STR.get(self.state, "Unknown")
        r_accent, g_accent, b_accent = accent
        luminance = 0.299 * r_accent + 0.587 * g_accent + 0.114 * b_accent
        text_color = (18, 24, 32) if luminance >= 170 else (255, 255, 255)

        draw.text((30, 31), state_str, fill=text_color, font=font_small)

        cy = 74
        draw.text((24, cy), "LXST Voice", fill=(255, 255, 255), font=font_title)
        cy += 32

        ch = self.channel_hash
        display_ch = ch[:16] + "..." if len(ch) > 19 else ch
        draw.text((24, cy), display_ch, fill=(180, 190, 200), font=font_body)
        cy += 26

        if self.state in (STATE_CONNECTED, STATE_TRANSMITTING):
            pc = self.participant_count
            draw.text(
                (24, cy), f"Participants: {pc}",
                fill=(150, 200, 150), font=font_body,
            )
            cy += 28

        cy += 8

        if self.state == STATE_IDLE:
            for line in self._wrap_text(draw, "Press button to connect", font_body, w - 56):
                draw.text((24, cy), line, fill=(160, 180, 210), font=font_body)
                cy += 24
        elif self.state == STATE_CONNECTING:
            for line in self._wrap_text(draw, "Connecting to channel...", font_body, w - 56):
                draw.text((24, cy), line, fill=(255, 200, 100), font=font_body)
                cy += 24
            for line in self._wrap_text(draw, "Press to cancel", font_small, w - 56):
                draw.text((24, cy), line, fill=(200, 170, 100), font=font_small)
                cy += 20
        elif self.state == STATE_CONNECTED:
            for line in self._wrap_text(draw, "Hold button to talk", font_body, w - 56):
                draw.text((24, cy), line, fill=(160, 220, 180), font=font_body)
                cy += 24
            for line in self._wrap_text(draw, "Long press 3s to leave", font_small, w - 56):
                draw.text((24, cy), line, fill=(150, 180, 160), font=font_small)
                cy += 20
        elif self.state == STATE_TRANSMITTING:
            draw.rounded_rectangle(
                (28, cy, w - 28, cy + 36), radius=8,
                fill=(220, 40, 40), outline=(255, 80, 80), width=2,
            )
            tx_label = "TRANSMITTING"
            bbox = draw.textbbox((0, 0), tx_label, font=font_title)
            tw = bbox[2] - bbox[0]
            draw.text(
                ((w - tw) // 2, cy + 6), tx_label,
                fill=(255, 255, 255), font=font_title,
            )
            cy += 48
            for line in self._wrap_text(draw, "Release to stop", font_small, w - 56):
                draw.text((24, cy), line, fill=(255, 150, 150), font=font_small)
                cy += 20
        elif self.state == STATE_FAILED:
            for line in self._wrap_text(draw, self._error_msg, font_body, w - 56):
                draw.text((24, cy), line, fill=(255, 150, 150), font=font_body)
                cy += 24
            for line in self._wrap_text(draw, "Press to retry", font_small, w - 56):
                draw.text((24, cy), line, fill=(200, 150, 150), font=font_small)
                cy += 20

        y_pos = h - 40
        draw.rounded_rectangle(
            (20, y_pos, w - 20, h - 22), radius=12,
            fill=(28, 37, 50),
        )
        if self.state == STATE_CONNECTED:
            draw.text(
                (28, y_pos + 8), "RX", fill=(100, 200, 100), font=font_tiny,
            )
            bar_w = w - 80
            bar_x = 60
            draw.rounded_rectangle(
                (bar_x, y_pos + 10, bar_x + bar_w, y_pos + 22),
                radius=4, fill=(40, 55, 70),
            )
            draw.rounded_rectangle(
                (bar_x, y_pos + 10, bar_x + int(bar_w * 0.6), y_pos + 22),
                radius=4, fill=(0, 200, 80),
            )
        elif self.state == STATE_TRANSMITTING:
            draw.text(
                (28, y_pos + 8), "TX", fill=(255, 100, 100), font=font_tiny,
            )
            bar_w = w - 80
            bar_x = 60
            draw.rounded_rectangle(
                (bar_x, y_pos + 10, bar_x + bar_w, y_pos + 22),
                radius=4, fill=(40, 55, 70),
            )
            draw.rounded_rectangle(
                (bar_x, y_pos + 10, bar_x + bar_w, y_pos + 22),
                radius=4, fill=(255, 50, 50),
            )
        elif self.state == STATE_IDLE:
            draw.text(
                (28, y_pos + 8), "Standby", fill=(150, 170, 190), font=font_tiny,
            )

        return self._rgb565_bytes(img)

    def _update_display(self):
        if self.state == STATE_IDLE:
            accent = ACCENT_IDLE
        elif self.state == STATE_CONNECTING:
            accent = ACCENT_CONNECTING
        elif self.state == STATE_CONNECTED:
            accent = ACCENT_CONNECTED
        elif self.state == STATE_TRANSMITTING:
            accent = ACCENT_TX
        elif self.state == STATE_FAILED:
            accent = ACCENT_FAILED
        else:
            accent = ACCENT_IDLE

        frame = self._render_frame(accent=accent)
        if frame:
            self.board.draw_image(
                0, 0, self.board.LCD_WIDTH, self.board.LCD_HEIGHT, frame,
            )

    def _set_state(self, state, error_msg=""):
        with self._lock:
            self.state = state
            self._error_msg = error_msg
            self._update_display()

    def _fail(self, msg):
        self._cleanup_all()
        self.board.set_rgb(*RGB_FAILED)
        self._set_state(STATE_FAILED, msg)
        RNS.log(msg, RNS.LOG_ERROR)
        time.sleep(4)
        with self._lock:
            if self.state == STATE_FAILED:
                self.state = STATE_IDLE
                self._update_display()
                self.board.set_rgb(*RGB_IDLE)

    def connect(self):
        with self._lock:
            if self.state not in (STATE_IDLE, STATE_FAILED):
                return
            self.state = STATE_CONNECTING
            self._update_display()

        self.board.set_rgb(*RGB_CONNECTING)

        try:
            channel_id = RNS.Identity.from_hex(self.channel_hash)
        except Exception as e:
            self._fail(f"Invalid channel hash: {e}")
            return

        dest = RNS.Destination(
            channel_id,
            RNS.Destination.OUT,
            RNS.Destination.SINGLE,
            APP_NAME,
            PRIMITIVE_NAME,
        )

        self.link = RNS.Link(dest)
        self.link.set_link_closed_callback(self._on_link_closed)

        deadline = time.time() + 20
        while self.link.status != RNS.Link.ACTIVE and time.time() < deadline:
            time.sleep(0.05)
            if self._cancel_connect:
                self._cancel_connect = False
                self._cleanup_link()
                self._set_state(STATE_IDLE)
                self.board.set_rgb(*RGB_IDLE)
                return

        if self.link.status != RNS.Link.ACTIVE:
            self._fail("Connection timeout - channel unreachable")
            return

        try:
            self.line_sink = LineSink()
            self.link_source = LinkSource(self.link, self, sink=self.line_sink)
            self.link_source.start()

            self.tx_codec = Opus(profile=Opus.PROFILE_VOICE_MEDIUM)
            self.packetizer = Packetizer(self.link)
            self.line_source = LineSource(codec=self.tx_codec, sink=self.packetizer)
            self.tx_pipeline = Pipeline(
                source=self.line_source,
                codec=self.tx_codec,
                sink=self.packetizer,
            )
        except Exception as e:
            self._fail(f"Audio setup failed: {e}")
            return

        with self._lock:
            self.state = STATE_CONNECTED
            self._update_display()
        self.board.set_rgb(*RGB_CONNECTED)

        RNS.log(
            f"Connected to channel {RNS.prettyhexrep(channel_id.hash)}",
            RNS.LOG_NOTICE,
        )

    def _cleanup_link(self):
        if self.link:
            try:
                if self.link.status == RNS.Link.ACTIVE:
                    self.link.teardown()
            except Exception:
                pass
            self.link = None

    def _cleanup_pipelines(self):
        if self.tx_pipeline and self.tx_pipeline.running:
            self.tx_pipeline.stop()
        self.tx_pipeline = None
        if self.link_source:
            self.link_source.stop()
            self.link_source = None
        if self.line_sink:
            self.line_sink.stop()
            self.line_sink = None
        if self.line_source:
            self.line_source = None
        if self.packetizer:
            self.packetizer.stop()
            self.packetizer = None

    def _cleanup_all(self):
        self._cleanup_pipelines()
        self._cleanup_link()

    def disconnect(self):
        with self._lock:
            was_connected = self.state not in (STATE_IDLE,)
            self._cleanup_all()
            self.state = STATE_IDLE
            self.participant_count = 0
            if was_connected:
                self._update_display()
            self.board.set_rgb(*RGB_IDLE)

        RNS.log("Disconnected from channel", RNS.LOG_NOTICE)

    def _on_link_closed(self, link):
        self.disconnect()

    def _ptt_start(self):
        with self._lock:
            if self.state != STATE_CONNECTED:
                return
            if self.tx_pipeline and not self.tx_pipeline.running:
                self.tx_pipeline.start()
                self.state = STATE_TRANSMITTING
                self.board.set_rgb(*RGB_TX)
                self._update_display()

    def _ptt_stop(self):
        with self._lock:
            if self.state != STATE_TRANSMITTING:
                return
            if self.tx_pipeline and self.tx_pipeline.running:
                self.tx_pipeline.stop()
            self.state = STATE_CONNECTED
            self.board.set_rgb(*RGB_CONNECTED)
            self._update_display()

    def _on_button_press(self):
        self._press_time = time.time()
        with self._lock:
            self._press_state = self.state
            state = self.state

        if state == STATE_CONNECTED:
            self._ptt_start()
        elif state == STATE_IDLE or state == STATE_FAILED:
            threading.Thread(target=self.connect, daemon=True).start()
        elif state == STATE_CONNECTING:
            self._cancel_connect = True

    def _on_button_release(self):
        held = time.time() - self._press_time
        if self._press_state == STATE_CONNECTED and held >= 3.0:
            threading.Thread(target=self.disconnect, daemon=True).start()
            return
        self._ptt_stop()

    def _on_exit_request(self, payload=None):
        self._exit = True
        self.disconnect()

    def _on_focus_revoked(self, payload=None):
        self.disconnect()

    def signalling_received(self, signals, source):
        pass

    def run(self):
        self.board.set_rgb(*RGB_IDLE)
        self._update_display()

        signal_module.signal(signal_module.SIGINT, self._sigint_handler)
        signal_module.signal(signal_module.SIGTERM, self._sigterm_handler)

        RNS.log(
            f"LXST Whisplay Client ready on channel hash: {self.channel_hash}",
            RNS.LOG_NOTICE,
        )
        RNS.log("Press button to connect", RNS.LOG_NOTICE)

        while not self._exit:
            time.sleep(0.1)

        self.cleanup()

    def _sigint_handler(self, sig, frame):
        self._exit = True

    def _sigterm_handler(self, sig, frame):
        self._exit = True

    def cleanup(self):
        self._exit = True
        self._cancel_connect = True
        self.disconnect()
        try:
            self.board.set_rgb(0, 0, 0)
            self.board.cleanup()
        except Exception:
            pass


def main():
    parser = argparse.ArgumentParser(
        description="LXST Whisplay Client - Group voice chat for Pi Zero 2W",
    )
    parser.add_argument(
        "channel",
        help="Channel destination hash (hex string from lxst-relay bot)",
    )
    parser.add_argument(
        "--identity", "-i", default=None,
        help="Path to persistent identity file",
    )
    parser.add_argument(
        "--rns-config", default=None,
        help="Path to alternative Reticulum config directory",
    )
    parser.add_argument(
        "-v", "--verbose", action="count", default=0,
    )
    args = parser.parse_args()

    client = LXSTWhisplayClient(
        channel_hash=args.channel,
        identity_path=args.identity,
        rns_config=args.rns_config,
        verbosity=args.verbose,
    )

    try:
        client.run()
    except KeyboardInterrupt:
        client.cleanup()
        print("")


if __name__ == "__main__":
    main()