"""Synthetic telemetry WS server for dashboard development.

Emits packets in the exact shape the stratoflights backend broadcasts
(see stratoflights_api/consumers.py), simulating a full flight profile:
ascent at ~5 m/s to 30 km, burst, parachute descent. Point the dashboard at
ws://localhost:8765/ and press "Подключиться".

Usage:
  pip install websockets
  python3 test_server.py [--interval 1.0] [--speedup 20]

--speedup compresses flight time: 20 means one wall-clock second advances the
flight by 20 s (a 3 h flight replays in ~9 min).

Push mode (--push): instead of serving, connect to the Django backend as a
ground station and push the same synthetic flight there. Both leaflet_svelte
(tracking) and this dashboard can then read the shared stream from
ws://<host>/api/ws/satellite/<uuid>/telemetry/:

  python3 test_server.py --push ws://localhost:8000 \
      --satellite <uuid> --token <drf-token> --speedup 20
"""

import argparse
import asyncio
import bisect
import json
import math
import random
import time
from datetime import datetime

import websockets

BURST_ALT = 30000.0
ASCENT_MS = 5.0
GROUND_ALT = 200.0
BASE_LAT, BASE_LON = 61.66, 129.38
START_WALL = time.time()

# Random lat/lon drift around the trajectory (metres): each packet takes a
# random step; past DRIFT_MAX_M the next step is forced back toward zero.
DRIFT_STEP_M = 150.0
DRIFT_MAX_M = 5000.0

TRAJ = None   # [(t_s, lat, lon, alt), ...] from --trajectory, else synthetic
_TRAJ_TIMES = []
_drift = {"n": 0.0, "e": 0.0}


def load_trajectory(path):
    """stratoflights/Tawhiri prediction JSON -> [(sec_from_launch, lat, lon, alt)]."""
    d = json.load(open(path))
    stages = (d.get("result") or d)["prediction"]
    pts = [p for s in stages for p in s["trajectory"]]
    t0 = datetime.fromisoformat(pts[0]["datetime"].replace("Z", "+00:00"))
    traj = []
    for p in pts:
        t = (datetime.fromisoformat(p["datetime"].replace("Z", "+00:00")) - t0).total_seconds()
        traj.append((t, p["latitude"], p["longitude"], p["altitude"]))
    return traj


def traj_at(flight_s):
    """Interpolated (lat, lon, alt, landed) along TRAJ."""
    if flight_s <= TRAJ[0][0]:
        _, lat, lon, alt = TRAJ[0]
        return lat, lon, alt, False
    if flight_s >= TRAJ[-1][0]:
        _, lat, lon, alt = TRAJ[-1]
        return lat, lon, alt, True
    i = bisect.bisect_right(_TRAJ_TIMES, flight_s)
    t0, la0, lo0, al0 = TRAJ[i - 1]
    t1, la1, lo1, al1 = TRAJ[i]
    f = (flight_s - t0) / (t1 - t0)
    return la0 + f * (la1 - la0), lo0 + f * (lo1 - lo0), al0 + f * (al1 - al0), False


def synth_at(flight_s):
    """Fallback synthetic profile: 5 m/s ascent to 30 km, parachute descent."""
    ascent_end = (BURST_ALT - GROUND_ALT) / ASCENT_MS
    if flight_s < ascent_end:
        alt = GROUND_ALT + ASCENT_MS * flight_s
        landed = False
    else:
        t = flight_s - ascent_end
        alt = BURST_ALT
        dt, step = 0.0, 5.0
        while dt < t and alt > GROUND_ALT:
            rate = 5.0 + 45.0 * (alt / BURST_ALT) ** 2
            alt -= rate * step
            dt += step
        alt = max(alt, GROUND_ALT)
        landed = alt <= GROUND_ALT
    lat = BASE_LAT + flight_s * 1.2e-5
    lon = BASE_LON + flight_s * 6e-5
    return lat, lon, alt, landed


def drift_step():
    for k in ("n", "e"):
        d = random.uniform(-DRIFT_STEP_M, DRIFT_STEP_M)
        if abs(_drift[k]) > DRIFT_MAX_M:
            d = -math.copysign(abs(d), _drift[k])  # forced step back
        _drift[k] += d


def sample(flight_s: float) -> dict:
    lat, lon, alt, landed = traj_at(flight_s) if TRAJ else synth_at(flight_s)
    drift_step()
    lat += _drift["n"] / 111320.0
    lon += _drift["e"] / (111320.0 * math.cos(math.radians(lat)))
    vs_phase = "landed" if landed else "flight"

    alt_noise = alt + random.uniform(-8, 8)
    pressure = 101325.0 * math.exp(-alt / 8400.0)
    temp_ext = 15.0 - 6.5 * min(alt, 11000) / 1000.0 + (0.002 * max(alt - 20000, 0))
    temp_int = max(temp_ext + 25.0, -15.0)
    hum = max(5.0, 60.0 - alt / 400.0)
    vbat = 3900 - flight_s * 0.06 + random.uniform(-5, 5)

    flags = ["GPS_FIX", "GPS_TIME_VALID", "SD_OK", "MS5611_OK", "SHT45_OK", "PWR_GNSS", "PWR_SD"]
    if temp_int < 0:
        flags.append("HEATER1_ON")
    if vbat < 3300:  # firmware BATT_LOW_MV (board.h)
        flags.append("LOW_BATT")

    utc = time.gmtime()
    return {
        "id": "test",
        # flight-time timestamp so derived vertical speed is realistic under --speedup
        "timestamp": int(START_WALL + flight_s),
        "lat": round(lat, 7),
        "lon": round(lon, 7),
        "alt": round(alt_noise, 1),
        "payload": {
            "utc": time.strftime("%H:%M:%S", utc) + ".00",
            "callsign": "YKSA-1",
            "uptime_s": int(flight_s) + 600,
            "sats": random.randint(7, 12) if vs_phase != "landed" else random.randint(4, 8),
            "vbat_mv": int(vbat),
            "pressure_pa": int(pressure),
            "baro_alt": round(alt_noise + random.uniform(-30, 30), 1) if alt < 25000 else None,
            "baro_temp_c": round(temp_int, 2),
            "humidity_pct": round(hum, 1),
            "sht_temp_c": round(temp_int - 3.0, 2),
            "tmp_ext_c": round(temp_ext, 2),
            "flags": flags,
        },
        "raw_data": {},
    }


async def handler(ws, interval: float, speedup: float, start_s: float = 0.0):
    print(f"client connected: {ws.remote_address}")
    start = time.time()
    n = 0
    try:
        while True:
            flight_s = start_s + (time.time() - start) * speedup
            packet = sample(flight_s)
            await ws.send(json.dumps(packet))
            n += 1
            if n % 25 == 0:
                print(f"  sent {n} packets, flight time {flight_s/60:.1f} min, alt {packet['alt']:.0f} m")
            await asyncio.sleep(interval)
    except websockets.ConnectionClosed:
        print(f"client disconnected after {n} packets")


async def push(base_url, satellite, token, interval, speedup, start_s):
    """Ground-station mode: push the synthetic flight into the Django backend."""
    base = base_url.rstrip("/").replace("http://", "ws://").replace("https://", "wss://")
    url = f"{base}/api/ws/station/{satellite}/telemetry/?token={token}"
    n = 0
    while True:
        try:
            async with websockets.connect(url) as ws:
                print(f"pushing to {url.split('?')[0]}")

                async def log_errors():
                    async for msg in ws:
                        print(f"  server: {msg}")

                reader = asyncio.ensure_future(log_errors())
                while True:
                    flight_s = start_s + (time.time() - START_WALL) * speedup
                    await ws.send(json.dumps(sample(flight_s)))
                    n += 1
                    if n % 25 == 0:
                        print(f"  sent {n} packets, flight time {flight_s/60:.1f} min")
                    await asyncio.sleep(interval)
        except (OSError, websockets.WebSocketException) as e:
            print(f"push link down ({e}); retrying in 3 s")
            await asyncio.sleep(3)


async def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--port", type=int, default=8765)
    ap.add_argument("--interval", type=float, default=1.0, help="seconds between packets")
    ap.add_argument("--speedup", type=float, default=20.0, help="flight seconds per wall second")
    ap.add_argument("--start", type=float, default=0.0,
                    help="flight seconds to skip at connect (5900 ≈ just before burst)")
    ap.add_argument("--push", metavar="URL", help="push to Django backend instead of serving")
    ap.add_argument("--satellite", help="satellite UUID (push mode)")
    ap.add_argument("--token", help="DRF auth token (push mode)")
    ap.add_argument("--trajectory", metavar="JSON",
                    help="fly along a stratoflights/Tawhiri prediction JSON "
                         "(с рандомным дрейфом ±5 км) вместо встроенного профиля")
    args = ap.parse_args()

    if args.trajectory:
        global TRAJ, _TRAJ_TIMES
        TRAJ = load_trajectory(args.trajectory)
        _TRAJ_TIMES = [p[0] for p in TRAJ]
        print(f"trajectory: {len(TRAJ)} pts, {_TRAJ_TIMES[-1]/60:.0f} min, "
              f"launch {TRAJ[0][1]:.5f},{TRAJ[0][2]:.5f}")

    if args.push:
        if not args.satellite or not args.token:
            ap.error("--push requires --satellite and --token")
        await push(args.push, args.satellite, args.token, args.interval, args.speedup, args.start)
        return

    async with websockets.serve(
        lambda ws: handler(ws, args.interval, args.speedup, args.start), "localhost", args.port
    ):
        print(f"synthetic telemetry on ws://localhost:{args.port}/  (speedup x{args.speedup})")
        await asyncio.Future()


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