#!/usr/bin/env python3
"""Add guarded trainer transactions to ALE. Rebuild worldserver after applying."""
from pathlib import Path
import sys
root=Path(sys.argv[1]);here=Path(__file__).resolve().parent
header=root/'modules/mod-ale/src/LuaEngine/methods/Playerbots/PlayerBotAIMethods.h'
s=header.read_text()
if 'int GetClassTrainerStatus(' not in s:
 s=s.replace('#include "Strategy.h"','#include "Strategy.h"\n#include "Trainer.h"\n#include "ObjectMgr.h"\n#include "DBCStores.h"')
 s=s.replace('    inline BotState CheckBotState', (here/'native/trainer-methods.inc').read_text()+'\n    inline BotState CheckBotState')
 header.write_text(s)
p=root/'modules/mod-ale/src/LuaEngine/LuaFunctions.cpp';s=p.read_text()
if '"GetClassTrainerStatus"' not in s:
 s=s.replace('    { "GetBot", &LuaPlayerBotAI::GetBot },','    { "GetBot", &LuaPlayerBotAI::GetBot },\n    { "GetClassTrainerStatus", &LuaPlayerBotAI::GetClassTrainerStatus },\n    { "TrainAtClassTrainer", &LuaPlayerBotAI::TrainAtClassTrainer },')
 p.write_text(s)
print('Trainer methods installed; rebuild required')
p=root/'modules/mod-playerbots/src/Ai/World/Rpg/Action/NewRpgBaseAction.cpp';s=p.read_text()
marker='// AzerothAI: unsupervised sightseeing flights are unsuitable for leveling residents.'
if marker not in s:
 import re
 pattern=r'(bool NewRpgBaseAction::CheckRpgStatusAvailable\(NewRpgStatus (\w+)\)\s*\{)'
 match=re.search(pattern,s);assert match,'Flight status signature changed'
 s=s[:match.end()]+'\n    '+marker+'\n    if ('+match.group(2)+' == RPG_TRAVEL_FLIGHT && bot->GetLevel() < 60) return false;\n'+s[match.end():]
 p.write_text(s)
p=root/'modules/mod-playerbots/src/Ai/World/Rpg/Strategy/RpgStrategy.cpp'
if not p.exists():p=next((root/'modules/mod-playerbots').rglob('RpgStrategy.cpp'))
s=p.read_text()
if '"rpg train"' not in s:
 s=s.replace('    // Sub actions','    // Training was missing from the NPC activity triggers.\n    triggers.push_back(new TriggerNode("rpg train", { NextAction("rpg train", 1.205f) }));\n\n    // Sub actions')
 p.write_text(s)
print('Training trigger and low-level sightseeing-flight guard installed')