import argparse
import re
import os
import sys


COLOR_RULES = {
    "black": "0.05 0.05 0.05", "dark": "0.15 0.15 0.15", "grey": "0.5 0.5 0.5",
    "white": "0.9 0.9 0.9", "chrome": "0.8 0.8 0.8", "metal": "0.7 0.7 0.7",
    "silver": "0.75 0.75 0.75", "gold": "0.8 0.7 0.2", "copper": "0.7 0.4 0.2",
    "glass": "0.1 0.1 0.12", "window": "0.1 0.1 0.12", "mirror": "0.9 0.9 0.9",
    "tyre": "0.03 0.03 0.03", "tire": "0.03 0.03 0.03", "rubber": "0.03 0.03 0.03",
    "wood": "0.4 0.25 0.1", "carbon": "0.15 0.15 0.15", "plastic": "0.8 0.8 0.8",
    "red": "0.8 0.1 0.1", "green": "0.1 0.8 0.1", "blue": "0.1 0.1 0.8",
    "yellow": "0.8 0.8 0.1"
}


def hex_to_kd(hex_str):
    r, g, b = int(hex_str[0:2], 16), int(hex_str[2:4], 16), int(hex_str[4:6], 16)
    return f"{r/255.0:.6f} {g/255.0:.6f} {b/255.0:.6f}"


def get_color(name):
    match = re.search(r'\(([A-Fa-f0-9]{6})\)', name)
    if match:
        return hex_to_kd(match.group(1))

    name_lower = name.lower()
    for kw, kd in COLOR_RULES.items():
        if kw in name_lower: return kd
    return None


def main():
    parser = argparse.ArgumentParser(description="Auto-fix missing Kd colors in MTL files based on material names.")
    parser.add_argument("input", help="Path to the input .mtl file")
    parser.add_argument("-o", "--output", help="Path to output file (default: <input>_fixed.mtl)")
    parser.add_argument("-w", "--overwrite", action="store_true", help="Overwrite the input file directly")
    args = parser.parse_args()

    if not os.path.exists(args.input):
        print(f"Error: {args.input} not found.")
        sys.exit(1)

    with open(args.input, 'r') as f:
        lines = f.readlines()

    out_lines = []
    current_mat = ""
    changes = 0

    for line in lines:
        if line.startswith("newmtl "):
            current_mat = line.split()[1].strip()
            out_lines.append(line)
        elif line.startswith("Kd ") and current_mat:
            new_kd = get_color(current_mat)
            if new_kd:
                out_lines.append(f"Kd {new_kd}\n")
                changes += 1
            else:
                out_lines.append(line)
        else:
            out_lines.append(line)

    out_path = args.input if args.overwrite else (args.output or args.input.replace(".mtl", "_fixed.mtl"))

    with open(out_path, 'w') as f:
        f.writelines(out_lines)

    print(f"Done! {changes} materials updated. Saved to {out_path}")


if __name__ == "__main__":
    main()