import cv2
import sys
import numpy as np
import os

if len(sys.argv) < 2:
    print("Please provide the input files.")
    sys.exit(1)

basecolor_file = None
normal_file = None
rma_file = None

# Loop through the command-line arguments
for arg in sys.argv[1:]:
    if '_alb.' in arg.lower():
        basecolor_file = arg
    elif '_nrm.' in arg.lower():
        normal_file = arg
    elif '_rma.' in arg.lower():
        rma_file = arg

# Check if all files are provided
if not all([basecolor_file, normal_file, rma_file]):
    print("Please provide all input files: base_color, normal, roughness, ambient_occlusion, and metallic.")
    sys.exit(1)

basecolor = cv2.imread(basecolor_file, cv2.IMREAD_COLOR)
normal = cv2.imread(normal_file, cv2.IMREAD_COLOR)
orm = cv2.imread(rma_file, cv2.IMREAD_COLOR)

basecolorb, basecolorg, basecolorr = basecolor.T
normalb, normalg, normalr = normal.T
aob, metallicb, roughnessb, = orm.T

BCR = np.array([basecolorb, basecolorg, basecolorr, roughnessb])
NMO = np.array([metallicb, normalg, normalr, aob])

#result = re.search(r'TX_(.*)_', basecolor_file)
output_filename = os.path.basename(basecolor_file)
output_filename = output_filename.removeprefix("TX_")
output_filename = output_filename.removesuffix("_ALB.PNG")

if not os.path.exists(os.path.join(os.path.dirname(basecolor_file),"merged")):
   os.mkdir(os.path.join(os.path.dirname(basecolor_file),"merged"))

output_dir = os.path.join(os.path.dirname(basecolor_file),"merged")

#output_dir = re.search(r'^(.+)\\.+$', basecolor_file).group(1)

cv2.imwrite(output_dir + f"\\{output_filename}_BCR.png", BCR.T)
cv2.imwrite(output_dir + f"\\{output_filename}_NMO.png", NMO.T)