## PPO_Bot — enemy tracker + state vector wired into the game loop.
## Training: trajectory collected per tick, PPO update on round end.

import std/os
import arraymancer
import tankroyale_botapi
import network
import actions
import training
import ./enemy_tracker
import ./state_vector

const botJsonPath = currentSourcePath().parentDir / "PPO_Bot.json"

type PPOBot = ref object of Bot
  tracker:    EnemyTracker
  buffer:     TrajectoryBuffer
  prevEnergy: float32       # own energy last tick
  prevEnemyE: float32       # enemy energy last tick (from tracker)
  lastState:  Tensor[float32]
  lastAction: Tensor[float32]
  lastLogP:   float32
  lastValue:  float32
  hasLastTrans: bool

var ac = initActorCritic()

method onScannedBot*(bot: PPOBot, e: ScannedBotEvent) =
  bot.tracker.update(e.x, e.y, e.direction, e.speed, e.energy)

method onRoundStarted*(bot: PPOBot, e: RoundStartedEvent) =
  bot.tracker       = initEnemyTracker()
  bot.buffer        = initTrajectoryBuffer()
  bot.prevEnergy    = 0.0'f32
  bot.prevEnemyE    = 0.0'f32
  bot.hasLastTrans  = false

method onRoundEnded*(bot: PPOBot, e: RoundEndedEventForBot) =
  # Add round-end score bonus to last transition (if any)
  let roundReward = computeRoundReward(e.results.totalScore.float32)
  if bot.hasLastTrans and bot.buffer.len > 0:
    bot.buffer.transitions[^1].reward += roundReward

  # PPO update synchronous; background thread is issue #17
  # ponytail: blocking update per round; move to thread pool when #17 lands
  ppoUpdate(ac, bot.buffer, lastValue = 0.0'f32)
  bot.buffer.clear()
  bot.hasLastTrans = false

method run(bot: PPOBot) =
  # Seed energy on first tick
  bot.prevEnergy = getEnergy().float32
  bot.prevEnemyE = if bot.tracker.hasContact: bot.tracker.current.energy.float32 else: 0.0'f32

  while isRunning():
    bot.tracker.deadReckon()

    setRadarTurnRate(bot.tracker.getRadarTurnRate(
      getX(), getY(), getDirection(), getRadarDirection()))

    let botData = BotStateData(
      x:           getX(),
      y:           getY(),
      direction:   getDirection(),
      speed:       getSpeed(),
      energy:      getEnergy(),
      gunDirection: getGunDirection(),
      gunHeat:     getGunHeat(),
      arenaWidth:  float64(getArenaWidth()),
      arenaHeight: float64(getArenaHeight()),
    )

    let state = buildStateVector(botData, bot.tracker)
    let (rawActs, logP) = ac.actorForward(state)
    let value = ac.criticForward(state)
    let acts = mapActions(rawActs, getSpeed().float32, getGunHeat().float32)

    # Compute tick reward from energy deltas
    let curEnergy   = getEnergy().float32
    let curEnemyE   = if bot.tracker.hasContact: bot.tracker.current.energy.float32 else: bot.prevEnemyE
    let myDelta     = curEnergy   - bot.prevEnergy
    let enemyDelta  = curEnemyE   - bot.prevEnemyE
    let tickReward  = computeTickReward(myDelta, enemyDelta)

    # Finalise previous transition with the reward from this tick's state change
    if bot.hasLastTrans:
      let tr = Transition(
        state:   bot.lastState,
        action:  bot.lastAction,
        logProb: bot.lastLogP,
        reward:  tickReward,
        value:   bot.lastValue,
      )
      bot.buffer.add(tr)

    # Store current for next tick
    bot.lastState   = state
    bot.lastAction  = rawActs
    bot.lastLogP    = logP
    bot.lastValue   = value
    bot.prevEnergy  = curEnergy
    bot.prevEnemyE  = curEnemyE
    bot.hasLastTrans = true

    setTargetSpeed(acts.targetSpeed.float)
    setTurnRate(acts.turnRate.float)
    setGunTurnRate(acts.gunTurnRate.float)
    if acts.shouldFire:
      discard setFire(acts.firePower.float)
    go()

when isMainModule:
  var bot = PPOBot(
    tracker: initEnemyTracker(),
    buffer:  initTrajectoryBuffer(),
  )
  start(bot, botJsonPath)