"""The minimal ProVega backend: ProgressiVis loads the data, the spec says how it is sent.

The ProgressiVis backend this replaces (progressivis-test/backend.py) took its
chunk size and pace from environment variables. The spec's chunking.reading
never reached it, so changing either meant editing and restarting the backend,
and a backend producing data faster than the chart could draw it buried Vega
under rows it could not keep up with.

Here the chart sends its chunking.reading when it connects (ProVega emits
"provega:start" to every Socket.IO data source) and the backend applies the
two properties that concern it:

    chunk_size   rows per message
    frequency    milliseconds between messages

Which file and which columns stay the backend's business. ProgressiVis loads
the file into a PTable progressively, at its own pace; what reaches the
browser is cut to the spec's measure. The protocol, server to browser:

    provega:meta   {"received": {...}, "backend": "..."}   on provega:start
    provega:meta   {"rows": n, "chunks": n}                 once the file is loaded
    data           [record, ...]                            one chunk
    (disconnect)                                            nothing more to send

The platform runs this as its default backend, on the Socket.IO namespace
/progressivis; a spec reaches it with "url": "sio:/progressivis". The same
file runs on its own, on any machine:

    pip install progressivis python-socketio aiohttp
    python provega_backend.py my-data.csv --port 8765

and a spec reaches that one with "url": "ws://localhost:8765".
"""
import asyncio
import inspect
import json
import math
import os

# On Linux ProgressiVis keeps tables in memory-mapped files, in one temporary
# folder for the whole process. With several charts streaming at once (the
# platform runs one scheduler per chart, each in its own thread) they tripped
# over each other's files: under a load test, 3 streams in 10 failed with
# FileNotFoundError on ..._table/Date_Time_strings. In memory, as it already is
# on Windows and macOS, every stream completed. Set before ProgressiVis is
# imported, which is when it reads the variable.
os.environ.setdefault("PROGRESSIVIS_STORAGE", "numpy")

# ProgressiVis 0.2.x annotates with np.typing without importing it, and NumPy
# does not load that submodule by itself: importing progressivis then fails.
import numpy.typing  # noqa: F401
import progressivis
from progressivis import Scheduler, SimpleCSVLoader, Sink

BACKEND = f"ProgressiVis {getattr(progressivis, '__version__', '')}".strip()

# Same defaults as ProVega applies in the browser, and bounds that keep one
# spec from asking for a message per millisecond or the whole file at once.
CHUNK_SIZE = (10, 1, 10000)       # default, min, max
FREQUENCY_MS = (1000, 50, 60000)


def spec_number(reading, key, bounds):
    default, low, high = bounds
    value = reading.get(key, default)
    if isinstance(value, bool) or not isinstance(value, (int, float)) or not math.isfinite(value):
        value = default
    return int(min(max(value, low), high))


def received_properties(reading):
    """The properties this backend takes from the spec, as it will apply them."""
    reading = reading if isinstance(reading, dict) else {}
    return {
        "chunk_size": spec_number(reading, "chunk_size", CHUNK_SIZE),
        "frequency": spec_number(reading, "frequency", FREQUENCY_MS),
    }


async def deliver(send, event, payload):
    result = send(event, payload)
    if inspect.isawaitable(result):
        await result


async def stream(path, reading, send, listening, columns=None):
    """Load path with ProgressiVis and send it in the spec's chunks.

    send(event, payload) delivers one message to the chart, and may be a
    coroutine function; listening() says whether anyone still wants it.
    columns, when given, names the fields as the file does: ProgressiVis
    renames those that are not identifiers ("Date/Time" becomes "Date_Time").
    """
    properties = received_properties(reading)
    size, delay = properties["chunk_size"], properties["frequency"] / 1000
    await deliver(send, "provega:meta", {"received": properties, "backend": BACKEND})

    scheduler = Scheduler()
    with scheduler:
        loader = SimpleCSVLoader(path, scheduler=scheduler)
        # A module whose output nobody reads is never run.
        sink = Sink(scheduler=scheduler)
        sink.input.inp = loader.output.result

    async def send_chunks():
        sent, announced = 0, False
        try:
            while listening():
                table = loader.result
                loaded = len(table) if table is not None else 0
                done = loader.is_terminated()
                if done and not announced:
                    announced = True
                    await deliver(send, "provega:meta", {"rows": loaded, "chunks": math.ceil(loaded / size)})
                if loaded - sent >= size or (done and loaded > sent):
                    end = min(sent + size, loaded)
                    # Through pandas' JSON: it writes NaN as null and NumPy
                    # numbers as plain ones, which the browser can parse.
                    frame = table.loc[sent:end - 1].to_df()
                    if columns and len(columns) == len(frame.columns):
                        frame.columns = columns
                    rows = json.loads(frame.to_json(orient="records"))
                    await deliver(send, "data", rows)
                    sent = end
                    await asyncio.sleep(delay)
                elif done:
                    return
                else:
                    await asyncio.sleep(0.02)  # ProgressiVis has not loaded a chunk's worth yet
        finally:
            await scheduler.stop()

    # A file ProgressiVis cannot read raises out of start(), and so out of here.
    await asyncio.gather(scheduler.start(), send_chunks())


def main():
    """Run the backend on its own, outside the platform."""
    import argparse
    import os

    import socketio
    from aiohttp import web

    parser = argparse.ArgumentParser(description="Stream a CSV file to ProVega charts, loaded by ProgressiVis.")
    parser.add_argument("csv", help="the CSV file to stream")
    parser.add_argument("--host", default="127.0.0.1", help="address to listen on (default 127.0.0.1)")
    parser.add_argument("--port", type=int, default=8765, help="port to listen on (default 8765)")
    args = parser.parse_args()
    if not os.path.isfile(args.csv):
        parser.error(f"no such file: {args.csv}")

    # Any page may connect: this is a tool for one's own machine and data.
    sio = socketio.AsyncServer(async_mode="aiohttp", cors_allowed_origins="*")
    app = web.Application()
    sio.attach(app)
    streams = set()

    @sio.on("provega:start")
    async def start(sid, payload):
        if sid in streams:
            return
        streams.add(sid)
        reading = payload.get("reading") if isinstance(payload, dict) else None
        print(f"[{sid}] spec says {received_properties(reading)}")

        def send(event, data):
            return sio.emit(event, data, to=sid)

        try:
            await stream(args.csv, reading, send, lambda: sid in streams)
        except Exception as e:
            print(f"[{sid}] {e}")
            await sio.emit("provega:error", {"message": f"The backend could not read {args.csv}: {e}"}, to=sid)
        finally:
            if sid in streams:
                streams.discard(sid)
                await sio.disconnect(sid)  # how the chart learns there is nothing more

    @sio.event
    async def disconnect(sid, *_):
        streams.discard(sid)

    print(f"{BACKEND} backend for {args.csv} on ws://{args.host}:{args.port}")
    web.run_app(app, host=args.host, port=args.port, print=None)


if __name__ == "__main__":
    main()