#!/usr/bin/env python3
"""
JESIM tsunami GUI -> Scratch/TurboWarp WebSocket relay.

v5 no-axis-flip relay:
- Pass grid_x/grid_y, mapped_grid_x/mapped_grid_y, display_grid_x/display_grid_y
  through exactly as C exports them.
- Do NOT invert y. The C side now binds mapped_grid_x/y to the selected
  existing simulator station id (station_index[sim_station_no-1]).
- Keeps data-only filtering and reset messages, but does not rewrite coordinates.
"""
import asyncio
import base64
import hashlib
import json
import struct
import time
from pathlib import Path
from typing import Any

HOST = "127.0.0.1"
PORT = 8765
HERE = Path(__file__).resolve().parent
OUTBOX = HERE / "jesim_ws_outbox.jsonl"
GRID_BIN = HERE / "tsunami_data.bin"
CLEAR_OUTBOX_ON_START = True
SEND_RESET_EVENT_ON_OUTBOX_TRUNCATE = True
GUID = "258EAFA5-E914-47DA-95CA-C5AB0DC85B11"

ALLOWED_KINDS = {
    "warning_check",
    "observation_arrival",
}

clients = set()
tail_pos = 0
start_time = time.time()
sent_count = 0
skipped_count = 0
reset_count = 0


def read_grid_size() -> tuple[int, int]:
    """Read width/height from tsunami_data.bin header. Fallback is the bundled map size."""
    try:
        data = GRID_BIN.read_bytes()[:28]
        if len(data) >= 28 and data[:14] == b"TSUNAMI_SIM_V1":
            width, height, _station_count = struct.unpack("<III", data[16:28])
            if 0 < width < 10000 and 0 < height < 10000:
                return int(width), int(height)
    except Exception as e:
        print(f"[relay] could not read grid size from {GRID_BIN.name}: {e}")
    return 897, 540


GRID_WIDTH, GRID_HEIGHT = read_grid_size()


def transform_outgoing_json(text: str) -> str:
    """Relay payloads without coordinate conversion.

    We still parse JSON so we can add a small debug flag, but all coordinate
    fields remain byte-for-byte numeric values from tsunami_gui_win32.c.
    """
    try:
        obj = json.loads(text)
    except Exception:
        return text

    if obj.get("type") != "jesim.tsunami" or obj.get("kind") not in ALLOWED_KINDS:
        return text

    obj["relay_coordinate_transform"] = "none"
    obj["grid_y_inverted_by_relay"] = False
    obj.setdefault("grid_width", GRID_WIDTH)
    obj.setdefault("grid_height", GRID_HEIGHT)
    obj.setdefault("grid_coordinate_system", "tssim_export_raw_no_axis_flip")
    return json.dumps(obj, ensure_ascii=False, separators=(",", ":"))


def make_reset_message() -> str:
    global reset_count
    reset_count += 1
    obj = {
        "schema": 1,
        "type": "jesim.tsunami",
        "kind": "sim_reset",
        "clear_observations": True,
        "clear_markers": True,
        "reset_count": reset_count,
        "grid_width": GRID_WIDTH,
        "grid_height": GRID_HEIGHT,
        "grid_coordinate_system": "tssim_export_raw_no_axis_flip",
        "reason": "outbox_truncated_or_simulation_reset",
    }
    return json.dumps(obj, ensure_ascii=False, separators=(",", ":"))


def ws_accept(key: str) -> str:
    digest = hashlib.sha1((key.strip() + GUID).encode("ascii")).digest()
    return base64.b64encode(digest).decode("ascii")


def ws_frame(text: str) -> bytes:
    data = text.encode("utf-8")
    n = len(data)
    if n < 126:
        header = bytes([0x81, n])
    elif n < 65536:
        header = bytes([0x81, 126]) + n.to_bytes(2, "big")
    else:
        header = bytes([0x81, 127]) + n.to_bytes(8, "big")
    return header + data


def parse_kind(text: str) -> str:
    try:
        return str(json.loads(text).get("kind", "unknown"))
    except Exception:
        return "raw"


def is_real_data_message(text: str) -> bool:
    try:
        obj = json.loads(text)
    except Exception:
        return False
    if obj.get("type") != "jesim.tsunami":
        return False
    return obj.get("kind") in ALLOWED_KINDS


def clear_outbox_on_startup():
    global tail_pos
    if CLEAR_OUTBOX_ON_START:
        try:
            OUTBOX.write_text("", encoding="utf-8")
            tail_pos = 0
            print(f"[relay] cleared old outbox at startup: {OUTBOX.resolve()}")
            return
        except Exception as e:
            print("[relay] could not clear outbox:", e)
    if OUTBOX.exists():
        tail_pos = OUTBOX.stat().st_size


async def broadcast(text: str):
    global sent_count
    if not clients:
        return

    dead = []
    frame = ws_frame(text)
    data_len = len(text.encode("utf-8"))
    for w in list(clients):
        try:
            w.write(frame)
            await w.drain()
        except Exception:
            dead.append(w)

    for w in dead:
        clients.discard(w)

    sent_count += 1
    print(f"[relay] sent {data_len} bytes to {len(clients)} client(s); kind={parse_kind(text)}")


async def handle_client(reader: asyncio.StreamReader, writer: asyncio.StreamWriter):
    try:
        raw = await asyncio.wait_for(reader.readuntil(b"\r\n\r\n"), timeout=5)
        headers = raw.decode("latin1", errors="ignore").split("\r\n")
        key = ""
        for h in headers:
            if h.lower().startswith("sec-websocket-key:"):
                key = h.split(":", 1)[1].strip()
                break

        if not key:
            writer.close()
            await writer.wait_closed()
            return

        response = (
            "HTTP/1.1 101 Switching Protocols\r\n"
            "Upgrade: websocket\r\n"
            "Connection: Upgrade\r\n"
            f"Sec-WebSocket-Accept: {ws_accept(key)}\r\n"
            "\r\n"
        )
        writer.write(response.encode("ascii"))
        await writer.drain()

        clients.add(writer)
        print(f"[relay] client connected; clients={len(clients)}")
        print("[relay] WebSocket connected. Sending real tsunami data with display-y conversion.")

        while not reader.at_eof():
            await asyncio.sleep(1)

    except Exception as e:
        print("[relay] client closed:", e)

    finally:
        clients.discard(writer)
        try:
            writer.close()
            await writer.wait_closed()
        except Exception:
            pass
        print(f"[relay] client disconnected; clients={len(clients)}")


async def tail_outbox():
    global tail_pos, skipped_count

    while True:
        if OUTBOX.exists():
            try:
                size = OUTBOX.stat().st_size
                if size < tail_pos:
                    tail_pos = 0
                    print("[relay] outbox truncated/reset detected")
                    if SEND_RESET_EVENT_ON_OUTBOX_TRUNCATE:
                        await broadcast(make_reset_message())

                # If no client is connected, keep tail_pos fixed so real data
                # written while Scratch is reconnecting can still be delivered.
                if clients:
                    with OUTBOX.open("r", encoding="utf-8", errors="replace") as f:
                        f.seek(tail_pos)
                        for line in f:
                            line = line.strip()
                            if not line:
                                continue
                            if is_real_data_message(line):
                                await broadcast(transform_outgoing_json(line))
                            else:
                                skipped_count += 1
                                print(f"[relay] skipped non-data message; kind={parse_kind(line)}")
                        tail_pos = f.tell()

            except Exception as e:
                print("[relay] tail error:", e)

        await asyncio.sleep(0.25)


async def console_status_loop():
    while True:
        await asyncio.sleep(10)
        size = OUTBOX.stat().st_size if OUTBOX.exists() else 0
        print(
            f"[relay] status console-only: clients={len(clients)} "
            f"outbox_size={size} tail_pos={tail_pos} sent={sent_count} skipped={skipped_count} "
            f"grid={GRID_WIDTH}x{GRID_HEIGHT} yfix=off uptime_s={int(time.time() - start_time)}"
        )


async def main():
    clear_outbox_on_startup()
    server = await asyncio.start_server(handle_client, HOST, PORT)
    print(f"[relay] listening on ws://{HOST}:{PORT}")
    print(f"[relay] outbox: {OUTBOX.resolve()}")
    print(f"[relay] grid size: {GRID_WIDTH}x{GRID_HEIGHT}")
    print("[relay] no axis flip: outgoing grid_y is unchanged from C export")
    print("[relay] data-only mode: Scratch receives warning_check plus observation_arrival event payloads")
    async with server:
        await asyncio.gather(server.serve_forever(), tail_outbox(), console_status_loop())


if __name__ == "__main__":
    asyncio.run(main())
