import { Centrifuge } from 'npm:centrifuge'

const TIMESTEP_MS = 1000
const ROBOT_COUNT = 40
const SPAWN_RADIUS = 0.9

class WorldControllerBase {
    constructor() {
        let client = new Centrifuge('ws://localhost:9120/connection/websocket')
        let sync = client.newSubscription('sync')
        sync.on('publication', (message) => {
            // console.log(Date.now(), 'recv sync', message.data)
            console.log('tick', message.data)
            this._last_tick_received = message.data
            this._process_tick(message.data)
        })
        sync.subscribe()
        client.connect()

        this.client = client
        this._dt_ms = TIMESTEP_MS
        this._last_tick_received = -1
        this._user_task = null

        client.on('connected', () => {
            this.on_start()
        })
    }

    async on_start() {
        // implemented by the user
    }

    async on_tick(_tick) {
        // implemented by the user
    }

    async rpc(method, args) {
        let result = await this.client.rpc(method, args)
        return result.data
    }

    _process_tick(tick) {
        if (this._user_task && this._last_tick_received > 0) return // busy

        // send sync
        let next_tick = tick + this._dt_ms * 1_000
        let sync = this.client.getSubscription('sync')
        sync.publish(next_tick)
        // console.log(Date.now(), 'send sync', next_tick)

        this._user_task = this.on_tick(tick)
            .finally(() => {
                this._user_task = null
                if (this._last_tick_received >= next_tick) {
                    this._process_tick(next_tick)
                }
            })
    }
}

class WorldController extends WorldControllerBase {
    async on_start() {
        this._robots = []
        await this.rpc('session.restart')

        let commands = []
        for (let i = 0; i < ROBOT_COUNT; i++) {
            let robot = new VehicleController(i)
            this._robots.push(robot)
            commands.push(robot.spawn(this))
        }
        commands.push(this.rpc('session.run'))
        await Promise.all(commands)
    }

    async on_tick() {
        let commands = []
        for (let robot of this._robots) {
            commands.push(robot.control(this))
        }
        await Promise.all(commands)
        this._robots = this._robots.filter(robot => robot.error == null)
    }
}

class VehicleController {
    constructor(id) {
        this.takeoff = true
        this.robot_id = id
        this.entity = null
        this.error = null
    }

    async spawn(world) {
        let angle = (2 * Math.PI / ROBOT_COUNT) * this.robot_id

        this.entity = await world.rpc('map.spawn_uav', {
            robot_id: this.robot_id,
            position: {
                x: SPAWN_RADIUS * Math.cos(angle),
                y: SPAWN_RADIUS * Math.sin(angle),
                z: 0,
            },
        })
    }

    async control(world) {
        try {
            const TAKEOFF_ALTITUDE = 2

            let position = await world.rpc(`object_${this.entity}.position`)

            if (this.takeoff && position.z > TAKEOFF_ALTITUDE - 0.05) {
                this.takeoff = false
            }

            if (this.takeoff) {
                await world.rpc(`object_${this.entity}.attitude_control`, {
                    z: TAKEOFF_ALTITUDE,
                    vxy: 0,
                    zx: 0,
                    yz: 0,
                })
                return
            }

            let distance_from_zero = Math.sqrt(position.x ** 2 + position.y ** 2)
            let angle = Math.atan2(position.y, position.x)

            // await world.rpc(`object_${this.entity}.velocity_control`, {
            //     z: position.z,
            //     vxy: 0,
            //     vx: Math.cos(angle) * SPAWN_RADIUS,
            //     vy: Math.sin(angle) * SPAWN_RADIUS,
            // })

            await world.rpc(`object_${this.entity}.position_control`, {
                z: position.z,
                x: Math.cos(angle + 0.7) * distance_from_zero * 1.5,
                y: Math.sin(angle + 0.7) * distance_from_zero * 1.5,
                xy: 0,
            })
        } catch (e) {
            this.error = e
            console.log(`vehicle ${this.robot_id} error: ${e.message}`)
        }
    }
}

new WorldController()
