#!/usr/bin/env python

import dnf
import subprocess
import os
import re
import sys

from typing import Set

DNF = 'dnf'

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_g_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',
	)
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_g_fg('*')
for line in lines:
	if not line or line.startswith('='):
		printlines += (line + '\n')
		continue
	(package, _) = line.split(' : ', maxsplit=1)
	package = dnf_color.sub('', package.rstrip())
	indicator = i if package in installed else ' '
	printlines += f'{indicator} {line}\n'
print(printlines, end='')