#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
远程显示 Web 服务 - Flask 实现

用法: python remo_disp_server.py [装置IP]
启动后访问 http://localhost:8181
依赖: pip install flask
"""

import socket
import struct
import time
import sys
import os
import io
from typing import Optional, Tuple

from flask import Flask, request, jsonify, send_file, Response

# 协议常量
TAG_CLIENT, TAG_DEVICE = 0xAA, 0xBB
PORT = 7003
CMD_KEEPLIVE, CMD_INIT, CMD_KEY, CMD_LCDMEM = 0, 1, 2, 3

# 全局连接
_sock: Optional[socket.socket] = None
_init_info: Optional[dict] = None


def calc_crc(data: bytes) -> int:
    crc = 0
    for b in data:
        crc ^= b
    return crc & 0xFF


def build_frame(cmd: int, data: bytes) -> bytes:
    length = len(data)
    header = bytes([TAG_CLIENT, cmd & 0xFF, (length >> 8) & 0xFF, length & 0xFF])
    return header + data + bytes([calc_crc(data)])


def parse_frame(raw: bytes) -> Optional[Tuple[int, bytes]]:
    if len(raw) < 6 or raw[0] != TAG_DEVICE:
        return None
    length = (raw[2] << 8) | raw[3]
    if len(raw) < 5 + length + 1:
        return None
    data = raw[4:4 + length]
    if calc_crc(data) != raw[4 + length]:
        return None
    return (raw[1], data)


def connect(host: str, port: int = PORT) -> bool:
    global _sock, _init_info
    try:
        if _sock:
            _sock.close()
        _sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
        _sock.settimeout(5.0)
        _sock.connect((host, port))
        _sock.settimeout(2.0)
        _init_info = _request_init()
        return True
    except Exception:
        return False


def _send(cmd: int, data: bytes = b"") -> bool:
    if not _sock:
        return False
    try:
        _sock.sendall(build_frame(cmd, data))
        return True
    except Exception:
        return False


def _recv() -> Optional[Tuple[int, bytes]]:
    if not _sock:
        return None
    try:
        header = _sock.recv(5)
        if len(header) < 5:
            return None
        length = (header[2] << 8) | header[3]
        rest = _sock.recv(length + 1)
        if len(rest) < length + 1:
            return None
        return parse_frame(header + rest)
    except Exception:
        return None


def _request_init() -> Optional[dict]:
    if not _send(CMD_INIT):
        return None
    for _ in range(30):
        result = _recv()
        if result and result[0] == CMD_INIT:
            data = result[1]
            if len(data) >= 29:
                return {
                    "width": (data[0] << 8) | data[1],
                    "height": (data[2] << 8) | data[3],
                    "mem_size": (data[5] << 24) | (data[6] << 16) | (data[7] << 8) | data[8],
                }
        time.sleep(0.05)
    return None


def send_key(code: int) -> bool:
    return _send(CMD_KEY, bytes([code & 0xFF]))


def fetch_screen() -> Optional[bytes]:
    global _init_info
    info = _init_info or _request_init()
    if not info:
        return None
    mem_size = info["mem_size"]
    buffer = bytearray(mem_size)
    pos = 0
    while pos < mem_size:
        if not _send(CMD_LCDMEM, struct.pack(">I", pos)):
            return None
        for _ in range(40):
            result = _recv()
            if result and result[0] == CMD_LCDMEM:
                payload = result[1]
                if len(payload) >= 4:
                    start = struct.unpack(">I", payload[:4])[0]
                    chunk = payload[4:]
                    buffer[start : start + len(chunk)] = chunk
                    pos = start + len(chunk)
                    break
            time.sleep(0.02)
        else:
            return None
    return bytes(buffer)


def mono_to_png(data: bytes, width: int, height: int) -> Optional[bytes]:
    try:
        from PIL import Image
        img = Image.new("1", (width, height))
        pix = img.load()
        for y in range(height):
            for x in range(width):
                bi = (y * width + x) // 8
                bit = 7 - (x % 8)
                pix[x, y] = 255 if (bi < len(data) and (data[bi] >> bit) & 1) else 0
        buf = io.BytesIO()
        img.save(buf, format="PNG")
        return buf.getvalue()
    except ImportError:
        return None


app = Flask(__name__)


@app.route("/")
@app.route("/index.html")
def index():
    return send_file(
        os.path.join(os.path.dirname(__file__), "remo_disp_ui.html"),
        mimetype="text/html; charset=utf-8",
    )


@app.route("/connect")
def api_connect():
    host = request.args.get("host", "").strip()
    port = int(request.args.get("port", PORT))
    ok = connect(host, port)
    return jsonify(success=ok, error=None if ok else "connect failed")


@app.route("/key")
def api_key():
    code = int(request.args.get("code", 0))
    send_key(code)
    return jsonify(ok=True)


@app.route("/screen")
def api_screen():
    data = fetch_screen()
    if not data:
        return "", 500
    info = _init_info or {}
    w, h = info.get("width", 160), info.get("height", 160)
    png = mono_to_png(data, w, h)
    if png:
        return Response(png, mimetype="image/png")
    resp = Response(data, mimetype="application/octet-stream")
    resp.headers["X-Width"] = str(w)
    resp.headers["X-Height"] = str(h)
    return resp


def main():
    port = 8181
    print(f"远程显示服务: http://localhost:{port}")
    print("在浏览器中打开上述地址，点击「连接」输入装置 IP")
    if len(sys.argv) >= 2:
        connect(sys.argv[1])
        print(f"已预连接: {sys.argv[1]}")
    app.run(host="0.0.0.0", port=port, debug=False, threaded=True)


if __name__ == "__main__":
    main()