#!/usr/bin/env python3
# Copyright 2025 The Chromium Authors
# Use of this source code is governed by a BSD-style license that can be
# found in the LICENSE file.

import argparse
import logging
import pathlib
import re
import subprocess

_SRC_ROOT = pathlib.Path(__file__).parents[3]

_GOOGLE_JAVA_FORMAT = (_SRC_ROOT / 'third_party' / 'google-java-format' /
                       'google-java-format')

_IMPORTS = """import org.chromium.build.annotations.NullMarked;
import org.chromium.build.annotations.NullMarked;
import org.chromium.build.annotations.Nullable;
import org.chromium.build.annotations.MonotonicNonNull;
import org.chromium.build.annotations.Contract;
import org.chromium.build.annotations.EnsuresNonNull;
import org.chromium.build.annotations.EnsuresNonNullIf;
import org.chromium.build.annotations.Initializer;
import org.chromium.build.annotations.RequiresNonNull;
import static org.chromium.build.NullUtil.assumeNonNull;
import static org.chromium.build.NullUtil.assertNonNull;
"""

_MODIFIER_KEYWORDS = (r'(?:(?:' + '|'.join([
    'abstract',
    'default',
    'final',
    'native',
    'private',
    'protected',
    'public',
    'static',
    'synchronized',
    r'/\*\s*package\s*\*/',
]) + r')\s+)*')

_NULLABLE_RE = re.compile(r'(\n *)@Nullable'
                          r'('
                          r'(?:\s*@\w+(?:\(.*?\))?)*'
                          r'\s+(?:' + _MODIFIER_KEYWORDS + r')?' +
                          r'(?:<.*?>)?'
                          r')')
_CLASSES_REGEX = re.compile(
    r'(^(?:public|protected|private|/\*\s*package\s*\*/)?\s*'
    r'(?:(?:static|abstract|final|sealed)\s+)*'
    r'(?:class|@?interface|enum)\s+\w+)',
    flags=re.MULTILINE)


def _mark_file(path):
    data = path.read_text()
    if '@NullMarked' in data:
        logging.warning('Skipping %s. Already has @NullMarked', path)
        return False
    if '@NullUnmarked' in data:
        logging.warning('Skipping %s. Already has @NullUnmarked', path)
        return False

    data = data.replace('import androidx.annotation.Nullable;\n', '')
    # Move @Nullable before methods to right before return type.
    data = _NULLABLE_RE.sub(r'\1\2 @Nullable ', data)
    # Fix up type-use position.
    data = re.sub(r'@Nullable\s+((?:\w+\.)+)(\w+)', r'\1@Nullable \2', data)
    data = re.sub(r'@Nullable\s+([\w<>]+)\[\]', r'\1 @Nullable[]', data)

    # Remove @NonNull
    data = data.replace('@NonNull', '')
    # Add imports
    data = re.sub(r'(^package .*\n)',
                  r'\1' + _IMPORTS,
                  data,
                  flags=re.MULTILINE,
                  count=1)

    # Add @NullMarked
    data = _CLASSES_REGEX.sub(r'@NullMarked\n\1', data, count=1)

    # Make all Void's @Nullable
    if re.search(r'\bVoid\b', data):
        data = re.sub(r'\bVoid\b', '@Nullable Void', data)
        data = data.replace('@Nullable @Nullable Void', '@Nullable Void')

    # Make all Supplier<Tab> -> Supplier<@Nullable Tab>
    if 'Supplier<Tab>' in data:
        data = data.replace('Supplier<Tab>', 'Supplier<@Nullable Tab>')

    logging.info('Processed: %s', path)
    path.write_text(data)
    return True


def main():
    logging.basicConfig(format='%(message)s', level=logging.INFO)

    parser = argparse.ArgumentParser()
    parser.add_argument('files', nargs='+')
    args = parser.parse_args()

    changed_paths = []
    for f in args.files:
        if _mark_file(pathlib.Path(f)):
            changed_paths.append(f)

    if changed_paths:
        cmd = [
            str(_GOOGLE_JAVA_FORMAT), '--aosp', '--skip-javadoc-formatting',
            '--skip-removing-unused-imports', '--replace'
        ] + changed_paths
        logging.info('Running: %s', ' '.join(cmd))
        subprocess.check_call(cmd)
    print(f'Added @NullMarked to {len(changed_paths)} path(s).')


if __name__ == '__main__':
    main()