#!/usr/bin/env python3

import argparse
import json
import os
import re
import sys
import time

# Define color codes
E = "\x1b["
colors = {
    "K": f"{E}1;30m",
    "R": f"{E}1;31m",
    "G": f"{E}1;32m",
    "B": f"{E}1;34m",
    "Y": f"{E}1;33m",
    "M": f"{E}1;35m",
    "C": f"{E}1;36m",
    "W": f"{E}1;37m",
    "k": f"{E}30m",
    "r": f"{E}31m",
    "g": f"{E}32m",
    "b": f"{E}34m",
    "y": f"{E}33m",
    "m": f"{E}35m",
    "c": f"{E}36m",
    "w": f"{E}37m",
    "u": f"{E}4m",
    "U": f"{E}4m",
    "Z": f"{E}0m",
    "BK": f"{E}40m",
    "BR": f"{E}41m",
    "BG": f"{E}42m",
    "BB": f"{E}44m",
    "BY": f"{E}43m",
    "BM": f"{E}45m",
    "BC": f"{E}46m",
    "BW": f"{E}47m",
}
simple_colors = "GYCMW"


def highlight_text(text, rules, simple_words, regex_patterns, case_sensitive):
    flags = 0 if case_sensitive else re.IGNORECASE

    # Apply simple highlights
    for i, word in enumerate(simple_words):
        color = simple_colors[i % len(simple_colors)]
        text = re.sub(rf"{re.escape(word)}", f"{colors[color]}{word}{colors['Z']}", text, flags=flags)

    # Apply regex highlights
    for pattern, replacement in rules:
        for code in colors:
            replacement = replacement.replace("{" + code + "}", colors[code])
        text = re.sub(pattern, f"{replacement}\\g<0>{colors['Z']}", text, flags=flags)

    return text


def load_rules(config_file):
    rules = []
    try:
        with open(config_file, "r") as f:
            rules = json.load(f)
    except Exception as e:
        print(f"Error reading config file: {e}")
        sys.exit(1)
    return rules


def main():
    parser = argparse.ArgumentParser(
        description="High Beam Highlighter",
        usage=""" cat code.py | highbeam -s def for lambda -r rules {B}{U}

Example ~/.highbeamrc:
[ {"pattern":"view", "color":"{G}"}, {"pattern":"^qo[^ ]*", "color":"{U}{Y}"} ]""",
    )
    parser.add_argument("-c", action="store_true", help="Be case sensitive")
    parser.add_argument("-f", default="~/.highbeamrc", help="Define config file (default ~/.highbeamrc)")
    parser.add_argument("-s", nargs="+", help="Simple highlight words")
    parser.add_argument("-r", nargs="+", help='Define rules with a string, ex. "-r word {G} ^error {R}" ')

    args = parser.parse_args()
    config_file = os.path.expanduser(args.f)
    rules = []
    if os.path.exists(config_file):
        rules = load_rules(config_file)

    simple_words = args.s if args.s else []
    regex_patterns = args.r if args.r else []

    # Convert regex patterns to rules
    for i in range(0, len(regex_patterns), 2):
        if i + 1 < len(regex_patterns):
            rules.append((regex_patterns[i], regex_patterns[i + 1]))
        else:
            raise ValueError("Regex patterns missing color code")

    for text in sys.stdin:
        print(highlight_text(text, rules, simple_words, regex_patterns, args.c), end="")


if __name__ == "__main__":
    main()