#!/usr/bin/env python3

import os
import re
import subprocess
import sys

import dnf

DNF = '/bin/dnf5'


def get_installed() -> set[str]:
	base = dnf.Base()
	base.fill_sack()
	q = base.sack.query()
	return {f'{p.name}.{p.arch}' for p in q.installed()}


def colorize_green_fg(s: str) -> str:
	return f'\x1b[32m{s}\033[0m'


if sys.argv[1] != 'search':
	os.execvp(DNF, sys.argv[1:])
	sys.exit(0)

# call CLI to get all the formatting and highlighting for free
res = subprocess.run(
	[DNF, '--color=always', *sys.argv[1:]],
	capture_output=True,
	encoding='utf-8',
	check=False,
)
lines = res.stdout.split('\n')
installed = get_installed()
printlines = []
# matching bold, magenta and normal codes from
# https://github.com/rpm-software-management/dnf/blob/master/dnf/cli/term.py
dnf_color = re.compile(r'\x1b\[35m|\x1b\[1m|\x1b\(B\x1b\[m')
i = colorize_green_fg('*')
for line in lines:
	if not line or not line.startswith(' '):
		printlines.append(line)
		continue
	(package, _) = line.split(':', maxsplit=1)
	package = dnf_color.sub('', package.lstrip())
	indicator = i if package in installed else ' '
	printlines.append(f'{indicator} {line[1:]}')  # [:1] to strip ' '
print('\n'.join(printlines))