#!/usr/bin/env python3

import argparse
import re
import sys


def sub(path, pattern, repl, expected_count=1):
    with open(path) as fh:
        content = fh.read()
    content, count = re.subn(pattern, repl, content, flags=re.MULTILINE)
    if expected_count is not None:
        assert count == expected_count, f"Pattern {pattern} replacement failed!"
    with open(path, "wt") as fh:
        fh.write(content)


def write_version(version: str):
    for p in [
            ".github/workflows/ci-docker.yml",
            ".github/workflows/release-dev.yml",
            ".github/workflows/release.yml"
    ]:
        sub(
            p,
            r'base_version=".*"',
            f'base_version="{version}"'
        )

    sub(
        "docker/Dockerfile",
        r"ARG BUILD_FROM=esphome/esphome-base-amd64:.*",
        f"ARG BUILD_FROM=esphome/esphome-base-amd64:{version}"
    )
    sub(
        "docker/Dockerfile.dev",
        r"FROM esphome/esphome-base-amd64:.*",
        f"FROM esphome/esphome-base-amd64:{version}"
    )
    sub(
        "docker/Dockerfile.lint",
        r"FROM esphome/esphome-lint-base:.*",
        f"FROM esphome/esphome-lint-base:{version}"
    )


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument('new_version', type=str)
    args = parser.parse_args()

    version = args.new_version
    print(f"Bumping to {version}")
    write_version(version)
    return 0


if __name__ == "__main__":
    sys.exit(main() or 0)