"""Apply the local budgets to the verified b12032 UI copy."""
from pathlib import Path
import re

ROOT = Path(__file__).resolve().parents[1]
ORIGINAL = ROOT / 'runtime/build/tools/ui/dist'
SERVED = ROOT / 'runtime/webui'
OLD_BUNDLE = 'bundle.BWYoSFMS.js'
NEW_BUNDLE = 'bundle.BWYoSFMS.local-reliability.js'
OLD_BUDGETS = '{[es.HIGH]:8192,[es.LOW]:512,[es.MAX]:-1,[es.MEDIUM]:2048}'
NEW_BUDGETS = '{[es.HIGH]:128,[es.LOW]:64,[es.MAX]:128,[es.MEDIUM]:96}'


def atomic_write(path, text):
    temporary = path.with_name(path.name + '.tmp')
    temporary.write_text(text)
    temporary.replace(path)


def main():
    bundle = (ORIGINAL / '_app/immutable' / OLD_BUNDLE).read_text()
    assert bundle.count(OLD_BUDGETS) == 1, 'Unexpected upstream UI; review the patch before updating'
    patched = bundle.replace(OLD_BUDGETS, NEW_BUDGETS)
    atomic_write(SERVED / '_app/immutable' / NEW_BUNDLE, patched)
    bootstrap = (ROOT / 'web-search/bootstrap.js').read_text()
    atomic_write(SERVED / 'local-qwen-web-search.js', bootstrap)
    index = (SERVED / 'index.html').read_text()
    index = re.sub(r'local-qwen-web-search\.js\?v=\d+',
                   'local-qwen-web-search.js?v=10', index)
    index = index.replace(OLD_BUNDLE, NEW_BUNDLE)
    assert index.count(NEW_BUNDLE) == 2
    atomic_write(SERVED / 'index.html', index)
    print('Prepared local UI: Low 64, Medium 96, High/Max 128 thinking tokens')


if __name__ == '__main__':
    main()