#!/usr/bin/env python3
# Copyright 2025 The Pigweed Authors
#
# Licensed under the Apache License, Version 2.0 (the "License"); you may not
# use this file except in compliance with the License. You may obtain a copy of
# the License at
#
#     https://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS, WITHOUT
# WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the
# License for the specific language governing permissions and limitations under
# the License.
from __future__ import annotations

import json
import re
import tarfile
from pathlib import Path

import requests

ORIGIN_URL = 'https://pigweed.googlesource.com/pigweed/pigweed'
LOG_URL = f'{ORIGIN_URL}/+log/main/{{}}?format=JSON&n=1'
TAR_URL = f'{ORIGIN_URL}/+archive/{{}}/{{}}.tar.gz'

PROTO_PATHS = [
    'pw_build/proto',
]


def main():
    base_dir = Path(__file__).parent

    for proto_path in PROTO_PATHS:
        proto_dir = base_dir / proto_path
        proto_dir.mkdir(parents=True, exist_ok=True)

        to_remove: set[Path] = set()
        for root, _, files in proto_dir.walk():
            for file in files:
                if file.endswith('.proto'):
                    to_remove.add(root / file)

        resp = requests.get(LOG_URL.format(proto_path))
        commit = str(json.loads(resp.text[4:])['log'][0]['commit'])
        print('Updating', proto_path, 'to', commit)

        resp = requests.get(TAR_URL.format(commit, proto_path), stream=True).raw
        with tarfile.open(mode='r|*', fileobj=resp) as tar:
            for item in tar:
                if item.name.endswith('.proto'):
                    print('Extracting', item.name)
                    tar.extract(item, path=proto_dir, filter='data')
                    item_path = proto_dir / item.name
                    to_remove.discard(item_path)

                    contents = item_path.read_text()
                    contents = re.sub(
                        r'^import "pw_',
                        'import "pigweed/pw_',
                        contents,
                        flags=re.MULTILINE,
                    )
                    item_path.write_text(contents)

        if to_remove:
            print('Removing stale protos')
            for rem in to_remove:
                rem.unlink()

    print('Done')


if __name__ == '__main__':
    main()