#!/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 re
import shutil
import subprocess
import sys
import util

_FIND_TRAILER = util.list_template('''self.trailers()
.filter(|trailer| trailer.key().lower() == "{}")
.map(|trailer| trailer.value())''')

_BUGGED_MISSING_ERR = '''\
Searching for bugs requires the CLI tool "bugged" to be installed.
This tool is only available to Googlers (see go/bugged for install instructions)
'''


def _search_bugs(query: str | None):
  args = [
    'bugged',
    'search',
    '--error-if-not-found=true',
    '--columns=issue,reporter,assignee,status,summary',
  ]
  if query:
    args.append(query)
  ps = util.run_command(
    args,
    stdout=subprocess.PIPE,
    check=False,
    text=True,
  )
  if ps.returncode == 5:  # No results found
    print('No results found. Try another query?')
    return _search_bugs(input('>>> '))
  elif ps.returncode != 0:
    print('Issue tracker search failed', file=sys.stderr)
    exit(1)

  # The first column is the bug number. We need this internally, but don't want
  # to show it to the user.
  lines = [
    re.split(r'\s+', line, maxsplit=1)
    for line in ps.stdout.rstrip().split('\n')
  ]
  print(f'#  {lines[0][1]}')
  issues = {}
  # Reverse it because we want the important information to be at the bottom.
  # We do this since the prompt is at the bottom.
  for i, (issue, line) in reversed(list(enumerate(lines))[1:]):
    issues[str(i)] = issue
    fmt = f'{i:>2} {line}'
    # Print every second line in cyan so you can easily line up the numbers
    # with the summary
    print(f'\033[36m{fmt}\033[0m' if i % 2 else fmt)

  # Maybe someone with experience with TUIs could improve on the UX, but this
  # works pretty well and is much easier to code.
  print('Type the number in the left column to select an issue')
  print('Type "q [query]" to search for something different')
  while True:
    answer = input('>>> ')
    if answer.startswith('q '):
      return _search_bugs(answer.removeprefix('q '))
    elif answer == '':
      # Run the default query
      return _search_bugs(None)
    elif answer in issues:
      return issues[answer]
    print('Invalid response')


def _add_handler(args, revs: str):
  direct = {
    x['change_id']
    for x in util.jj_log(
      revisions=f'{revs}',
      templates={
        'change_id': 'change_id',
      },
      ignore_working_copy=True,
    )
  }

  ancestors = util.jj_log(
    revisions=f'mutable()::({revs})',
    templates={
      'change_id': 'change_id',
      'parents': util.MUTABLE_PARENTS,
      'trailers': 'trailers',
      'desc': 'description',
      'bugs': _FIND_TRAILER.format('bug'),
      'fixed': _FIND_TRAILER.format('fixed'),
    },
    ignore_working_copy=True,
  )

  revs = {rev['change_id']: rev for rev in ancestors}

  for rev in ancestors:
    rev['fixed'] = set(rev['fixed'])
    rev['bugs'] = set(rev['bugs'])

  bugs = set(args.bug)

  if args.query or (not bugs and not args.inherit):
    if args.query:
      print('Searching for bugs on issue tracker...')
    else:
      print('No bugs provided. Searching for bugs on issue tracker...')
    bugged = shutil.which('bugged')
    if bugged is None:
      print(_BUGGED_MISSING_ERR, file=sys.stderr)
      exit(1)
    bugs.add(_search_bugs(args.query))

  for rev in reversed(ancestors):
    if rev['change_id'] not in direct:
      # Suppose you have A <- B <- C:
      # When running `jj bug add --inherit -r C`, C should inherit from B
      # When running `jj bug add --inherit -r B|C`, C should inherit from B,
      # which should have already inherited from A.
      continue
    bugs_to_add = set(bugs)
    if args.inherit:
      # This doesn't need to be done recursively because of the toposort
      # guaruntee of `jj log`
      if rev['parents']:
        for change_id in rev['parents']:
          bugs_to_add.update(revs[change_id]['bugs'])

    bugs_to_add -= rev['bugs']
    bugs_to_add -= rev['fixed']
    (rev['fixed'] if args.fixed else rev['bugs']).update(bugs_to_add)

    # Skip the ones we already have
    tag = 'Fixed' if args.fixed else 'Bug'
    # We may not need to add bugs, especially in the case of inherit.
    if bugs_to_add:
      util.add_trailers(
        rev=rev,
        trailers={tag: [','.join(bugs_to_add)]},
        commit=True,
      )


def main(args):
  logging.basicConfig(level=logging.getLevelNamesMapping()[args.verbosity])

  revs = util.join_revsets(args.revision)
  args.func(args, revs)


if __name__ == '__main__':
  parser = argparse.ArgumentParser()
  parser.add_argument(
    '--verbosity',
    help='Verbosity of logging',
    default='INFO',
    choices=logging.getLevelNamesMapping().keys(),
    type=lambda x: x.upper(),
  )

  subparsers = parser.add_subparsers(required=True)

  parser_add = subparsers.add_parser('add')
  parser_add.set_defaults(func=_add_handler)
  parser_add.add_argument(
    '-r',
    '--revision',
    help='The revisions to add the bug to',
    action='append',
    required=True,
  )
  parser_add.add_argument(
    '-f',
    '--fixed',
    help='Mark the revision as fixing the bug',
    action='store_true',
  )
  parser_add.add_argument(
    '-i',
    '--inherit',
    help='Inherit the bug from parent commits. Bugs are not transitively '
    + 'inherited unless the parent is also part of -r',
    action='store_true',
  )
  parser_add.add_argument(
    '-b',
    '--bug',
    help='The bug to add to the revisions',
    action='append',
    default=[],
  )
  parser_add.add_argument(
    '-q',
    '--query',
    help='Query to use when searching for a bug in issue tracker',
    default=None,
  )

  main(parser.parse_args())