import os, time, subprocess
from pymxs import runtime as rt
try: from PySide6 import QtWidgets, QtCore, QtGui
except ImportError: from PySide2 import QtWidgets, QtCore, QtGui
import vr4life_cloud as cld

def get_vdenoise_path():
    for k, v in os.environ.items():
        if "VRAY" in k.upper() and "MAX" in k.upper() and "MAIN" in k.upper():
            p = os.path.join(v, "vdenoise.exe")
            if os.path.exists(p): return f'"{p}"'
    return "vdenoise.exe" 

def load_bake_elements(ui):
    ui.cmb_bake_elem.clear()
    found = ["CompleteMap", "VRayCompleteMap", "Corona_Beauty", "CShading_Beauty"]
    try: 
        r_list = rt.execute("for c in bake_elements.classes collect (c as string)")
        for c in r_list:
            if str(c) not in found: found.append(str(c))
    except: pass
    ui.cmb_bake_elem.addItems(sorted(list(set(found))))
    try:
        rnd = str(rt.execute("renderers.current as string"))
        t = "VRayCompleteMap" if "V_Ray" in rnd else "Corona_Beauty" if "Corona" in rnd else "CompleteMap"
        idx = ui.cmb_bake_elem.findText(t)
        if idx >= 0: ui.cmb_bake_elem.setCurrentIndex(idx)
    except: pass

def load_selection(ui):
    ui.tree.clear(); ui.bake_items = []
    ms = """( local g = #(); fn ex objs = ( for o in objs do ( if isGroupHead o then ( setGroupOpen o true; ex o.children ) else ( if superclassof o == GeometryClass and (classOf o != TargetObject) and (canConvertTo o Editable_Poly) do appendIfUnique g o ) ) ); ex selection; g )"""
    sel = rt.execute(ms)
    if not sel: return
    tl = [{'name': str(o.name), 'poly': rt.getPolygonCount(o)[0] if rt.canConvertTo(o, rt.Editable_Poly) else 0} for o in sel]
    tl.sort(key=lambda x: x['poly'], reverse=True)
    for d in tl:
        i = QtWidgets.QTreeWidgetItem([d['name'], "Pronto", f"{d['poly']:,}", "", ""])
        i.setFlags(i.flags() | QtCore.Qt.ItemIsUserCheckable)
        i.setCheckState(0, QtCore.Qt.Checked) 
        ui.tree.addTopLevelItem(i)
        ui.bake_items.append({'name': d['name'], 'item': i})
    ui.pb.setFormat(f"Carregados: {len(tl)}")
    ui.pb.setValue(0)
    ui.upd_res_col()

def attach_grouped_objects(ui, from_auto=False):
    ui.pb.setFormat("Fundindo grupos inteiros...")
    QtWidgets.QApplication.processEvents()
    ms = """(
        fn hasProtectedName nd = (
            local nm = toLower nd.name
            if (matchPattern nm pattern:"*_mult*") or (matchPattern nm pattern:"*_catg*") or (matchPattern nm pattern:"*_mat*") do return true
            for c in nd.children do ( if hasProtectedName c do return true )
            return false
        )
        local act = 0; local new_objs = #(); local orig_sel = selection as array; 
        local loose = for o in orig_sel where not isGroupHead o and not isGroupMember o collect o; 
        local heads = for o in orig_sel where isGroupHead o and not (hasProtectedName o) collect o
        
        for h in heads do ( 
            local valid_cg = for c in h.children where superclassof c == GeometryClass and classof c != TargetObject collect c
            if valid_cg.count > 0 do ( 
                local b = valid_cg[1]
                if not isKindOf b Editable_Poly do try(convertToPoly b)catch()
                if isKindOf b Editable_Poly do ( 
                    for i = 2 to valid_cg.count do ( 
                        local n = valid_cg[i]
                        if not isKindOf n Editable_Poly do try(convertToPoly n)catch()
                        if isKindOf n Editable_Poly do try(polyOp.attach b n)catch() 
                    )
                    b.name = h.name
                    setGroupMember b false
                    append new_objs b
                    act += 1 
                ) 
            )
            try(delete h)catch() 
        )
        local final_sel = loose; for no in new_objs do append final_sel no; if final_sel.count > 0 do select final_sel; act
    )"""
    try:
        act = rt.execute(ms)
        load_selection(ui) 
        if not from_auto: 
            if act > 0: QtWidgets.QMessageBox.information(ui, "Sucesso", f"{act} Grupos fundidos.") 
            else: QtWidgets.QMessageBox.information(ui, "Aviso", "Nenhum Grupo encontrado.")
    except Exception as e: print("Erro Solda:", e)
    ui.pb.setValue(0); ui.pb.setFormat("Pronto")

def super_attach_objects(ui, from_auto=False):
    ui.pb.setFormat("Super Solda VR...")
    QtWidgets.QApplication.processEvents()
    ms = """( 
        fn hasProtectedName nd = (
            local nm = toLower nd.name
            if (matchPattern nm pattern:"*_mult*") or (matchPattern nm pattern:"*_catg*") or (matchPattern nm pattern:"*_mat*") do return true
            for c in nd.children do ( if hasProtectedName c do return true )
            return false
        )
        local sel = for o in selection where superclassof o == GeometryClass and classof o != TargetObject and not (hasProtectedName o) collect o
        if sel.count > 1 do ( 
            local b = sel[1]
            if not isKindOf b Editable_Poly do try(convertToPoly b)catch()
            if isKindOf b Editable_Poly do ( 
                for i = 2 to sel.count do ( 
                    local n = sel[i]
                    if not isKindOf n Editable_Poly do try(convertToPoly n)catch()
                    if isKindOf n Editable_Poly do try(polyOp.attach b n)catch() 
                )
                b.name = uniqueName "VR_Bloco_Fundido"
                select b 
            ) 
        )
        sel.count 
    )"""
    try:
        act = rt.execute(ms)
        load_selection(ui)
        if not from_auto: 
            if act > 1: QtWidgets.QMessageBox.information(ui, "Sucesso", f"{act} objetos fundidos.") 
            else: QtWidgets.QMessageBox.information(ui, "Aviso", "Selecione pelo menos 2 objetos.")
    except Exception as e: print("Erro Super Solda:", e)
    ui.pb.setValue(0); ui.pb.setFormat("Pronto")

def optimize_geometry(ui):
    sp = ui.spn_pct.value()
    tg = ui.spn_min_poly.value()
    tgs = ui.get_processable_items()
    tot = len(tgs)
    ui.pb.setMaximum(tot)
    ui.pb.setValue(0)
    rt.execute("max modify mode")
    
    for i, d in enumerate(tgs):
        if ui._is_cancelled: break
        ui.pb.setFormat(f"Opt ({i+1}/{tot}): {d['name']}...")
        ui.pb.setValue(i)
        QtWidgets.QApplication.processEvents()
        
        o = rt.getNodeByName(d['name'])
        if o:
            try:
                cp = rt.getPolygonCount(o)[0] if rt.canConvertTo(o, rt.Editable_Poly) else 0
                if cp < 50: 
                    d['item'].setText(1, "Geo Base")
                else:
                    rt.select(o)
                    if not rt.isKindOf(o, rt.Editable_Poly): 
                        rt.convertTo(o, rt.Editable_Poly)
                        
                    opt = rt.ProOptimizer()
                    rt.addModifier(o, opt)
                    opt.KeepTextures = True
                    opt.KeepNormals = True
                    opt.KeepBorders = True
                    opt.LockMat = True
                    opt.Calculate = True
                    
                    current_pct = 100.0
                    p = 0
                    while True:
                        if ui._is_cancelled: break
                        current_pct = current_pct * (sp / 100.0)
                        opt.VertexPercent = current_pct
                        rt.redrawViews()
                        QtWidgets.QApplication.processEvents()
                        
                        np = rt.getPolygonCount(o)[0]
                        p += 1
                        d['item'].setText(2, f"{np:,}")
                        d['item'].setText(1, f"Opt ({p}x)")
                        QtWidgets.QApplication.processEvents()
                        
                        if np <= tg or np < 50 or p >= 20 or (np >= cp and p > 1): break 
                        cp = np
                        
                    rt.collapseStack(o)
                    if not rt.isKindOf(o, rt.Editable_Poly): 
                        rt.convertTo(o, rt.Editable_Poly)
                        
                    if not ui._is_cancelled: 
                        d['item'].setText(2, f"{rt.getPolygonCount(o)[0]:,}")
                        d['item'].setText(1, f"Opt OK ({p}x)")
            except Exception as e: 
                d['item'].setText(1, "Erro Opt")
                print(f"Erro Opt {d['name']}: {e}")
                
    ui.pb.setValue(tot)
    ui.pb.setFormat("Opt Concluída!")

def slice_large_objects(ui, from_auto=False):
    if not ui.bake_items: return
    thr = ui.spn_max_sz.value()
    tgs = ui.get_processable_items()
    tot = len(tgs)
    ui.pb.setMaximum(tot)
    ui.pb.setValue(0)
    rt.execute("max modify mode")
    
    print("\n" + "="*50)
    print("🔍 INICIANDO LOG: MULTI-UV (P4)")
    print(f"-> Tamanho de corte (Threshold): {thr}")
    print("="*50)
    
    for i, d in enumerate(tgs):
        if ui._is_cancelled: break
        ui.pb.setFormat(f"Multi-UV ({i+1}/{tot}): {d['name']}...")
        ui.pb.setValue(i)
        QtWidgets.QApplication.processEvents()
        
        print(f"\n--- Processando objeto: {d['name']} ---")
        
        # CÓDIGO CORRIGIDO: Tudo embrulhado num "fn" (função) nativo para o return funcionar.
        ms = f"""(
            fn applyMultiMatID objName thrVal = (
                local o = getNodeByName objName
                if o == undefined do return -1
                if not isKindOf o Editable_Poly do convertToPoly o
                
                local dx = abs(o.max.x - o.min.x)
                local dy = abs(o.max.y - o.min.y)
                local dz = abs(o.max.z - o.min.z)
                local md = amax #(dx, dy, dz)
                
                if md <= thrVal do ( 
                    for f = 1 to (polyop.getNumFaces o) do polyop.setFaceMatID o f 1
                    return 1 
                )
                
                local numChunks = (ceil (md / thrVal)) as integer
                if numChunks > 10 do numChunks = 10
                
                local axis = if md == dx then 1 else if md == dy then 2 else 3
                local minVal = if axis == 1 then o.min.x else if axis == 2 then o.min.y else o.min.z
                local step = md / numChunks
                
                for f = 1 to (polyop.getNumFaces o) do (
                    local center = polyop.getFaceCenter o f
                    local val = if axis == 1 then center.x else if axis == 2 then center.y else center.z
                    local id = (floor ((val - minVal) / step)) as integer + 1
                    if id > numChunks do id = numChunks
                    if id < 1 do id = 1
                    polyop.setFaceMatID o f id
                )
                return numChunks
            )
            try ( applyMultiMatID "{d['name']}" {thr} ) catch ( -2 )
        )"""
        
        try:
            print("-> Executando MaxScript de Fatiamento Lógico...")
            r = rt.execute(ms)
            print(f"-> Resposta do MaxScript: {r}")
            
            if r == -1:
                print("🚨 ERRO: Objeto não encontrado na cena pelo MaxScript.")
                d['item'].setText(1, "Erro: Objeto")
            elif r == -2:
                print("🚨 ERRO: MaxScript falhou internamente (Topologia corrompida?).")
                d['item'].setText(1, "Erro: Script")
            elif r and int(r) > 1:
                print(f"✅ SUCESSO: Objeto dividido em {int(r)} IDs!")
                d['item'].setText(1, f"IDs: {int(r)}")
            else:
                print("-> Objeto menor que a Meta. Mantido como 1 ID.")
                d['item'].setText(1, "1 ID (Normal)")
                
        except Exception as e: 
            d['item'].setText(1, "Erro Python")
            print(f"🚨 ERRO CRÍTICO PYTHON no Multi-UV: {e}")
            
    ui.pb.setValue(tot)
    ui.pb.setFormat("Multi-IDs Gerados!")
    print("\n" + "="*50)
    print("✅ FIM DO LOG MULTI-UV")
    print("="*50 + "\n")

def prepare_mesh(ui):
    tgs = ui.get_processable_items()
    tot = len(tgs)
    ui.pb.setMaximum(tot)
    ui.pb.setValue(0)
    rt.execute("max modify mode")
    
    for i, d in enumerate(tgs):
        if ui._is_cancelled: break
        ui.pb.setFormat(f"UV ({i+1}/{tot}): {d['name']}...")
        ui.pb.setValue(i)
        QtWidgets.QApplication.processEvents()
        o = rt.getNodeByName(d['name'])
        if o:
            rt.select(o)
            try:
                ms = f"""( 
                    local o = getNodeByName "{d['name']}"
                    if not isKindOf o Editable_Poly do convertToPoly o
                    local m = Unwrap_UVW()
                    addModifier o m
                    m.setMapChannel 3
                    m.setTVSubObjectMode 3
                    local maxID = 1
                    for f = 1 to (polyop.getNumFaces o) do ( 
                        local id = polyop.getFaceMatID o f
                        if id > maxID do maxID = id 
                    )
                    for x = 1 to maxID do (
                        m.selectByMatID x
                        local sel = m.getSelectedFaces()
                        if not sel.isEmpty do ( 
                            m.flattenMap 45.0 #() 0.002 true 0 true false
                            m.pack 1 0.002 true false false 
                        )
                    )
                    collapseStack o 
                )"""
                rt.execute(ms)
                if not rt.isKindOf(o, rt.Editable_Poly): 
                    rt.convertTo(o, rt.Editable_Poly)
                d['item'].setText(1, "UV Packed (C3)")
            except: 
                d['item'].setText(1, "Erro UV")
                
    ui.pb.setValue(tot)
    ui.pb.setFormat("UV Pack Concluído!")

def process_bake_logic(ui, auto_export=False):
    ui._is_cancelled = False
    el = ui.cmb_bake_elem.currentText()
    p_bk = ui.edt_p_bake.text()
    rnd = str(rt.execute("renderers.current as string"))
    i_vr = "V_Ray" in rnd
    i_cor = "Corona" in rnd
    u_den = ui.chk_denoise.isChecked()
    a_res = ui.chk_a_res.isChecked()
    
    if not os.path.exists(p_bk): os.makedirs(p_bk)
    tgs = ui.get_processable_items()
    tot = len(tgs)
    ui.pb.setMaximum(tot)
    ui.pb.setValue(0)
    
    print("\n" + "="*50)
    print("🎬 INICIANDO LOG: MOTOR DE BAKE (P6)")
    print(f"-> Pasta Alvo: {p_bk}")
    print(f"-> Renderizador: {rnd}")
    print("="*50)

    for i, d in enumerate(tgs):
        if ui._is_cancelled: break
        ui.pb.setFormat(f"Bake ({i+1}/{tot}): {d['name']}...")
        ui.pb.setValue(i)
        QtWidgets.QApplication.processEvents()
        
        print(f"\n--- Preparando Bake: {d['name']} ---")
        
        o = rt.getNodeByName(d['name'])
        if not o: 
            print("🚨 ERRO: Objeto não encontrado.")
            continue
        
        try:
            # 1. Conta quantos IDs existem no objeto
            ms_mid = f"""(local o = getNodeByName "{d['name']}"; local mid = 1; if o != undefined and isKindOf o Editable_Poly do ( for f = 1 to (polyop.getNumFaces o) do ( local id = polyop.getFaceMatID o f; if id > mid do mid = id ) ); mid)"""
            max_id = rt.execute(ms_mid)
            if not max_id: max_id = 1
            max_id = int(max_id)
            print(f"-> Total de IDs a renderizar: {max_id}")
            
            rt.execute("max select none")
            rt.select(o)
            rt.execute("max modify mode")
            rt.execute("max zoomext sel all")
            rt.redrawViews()
            time.sleep(0.5)
            
            for mid in range(1, max_id + 1):
                if ui._is_cancelled: break
                print(f"\n-> Assando ID: {mid} de {max_id}")
                
                if max_id == 1:
                    t_jpg_id = os.path.join(p_bk, f"{d['name']}_B.jpg").replace("\\", "/")
                else:
                    t_jpg_id = os.path.join(p_bk, f"{d['name']}_B_ID{mid}.jpg").replace("\\", "/")
                    
                if os.path.exists(t_jpg_id): 
                    d['item'].setText(1, f"Já existe ID{mid}")
                    print("-> Ficheiro já existe. Pulando.")
                    continue
                
                cr = ui.spn_res.value()
                if a_res:
                    md = max(abs(o.max.x - o.min.x), abs(o.max.y - o.min.y), abs(o.max.z - o.min.z))
                    cr = 256 if md <= ui.s256.value() else 512 if md <= ui.s512.value() else 1024 if md <= ui.s1024.value() else cr
                
                ext = ".exr" if (i_vr and u_den) else ".jpg"
                t_rnd_id = t_jpg_id.replace(".jpg", ".exr") if ext == ".exr" else t_jpg_id
                
                if max_id > 1:
                    print("-> Isolando UV do ID atual...")
                    # CÓDIGO CORRIGIDO 1: Embrulhado em 'fn'
                    ms_uv = f"""(
                        fn isolateUV currentMid = (
                            global temp_uv_mod = Unwrap_UVW()
                            addModifier $ temp_uv_mod
                            temp_uv_mod.setMapChannel 3
                            temp_uv_mod.setTVSubObjectMode 3
                            local allF = #{{1..(polyop.getNumFaces $)}}
                            temp_uv_mod.selectByMatID currentMid
                            local tgtF = temp_uv_mod.getSelectedFaces()
                            local hideF = allF - tgtF
                            
                            if hideF.numberset > 0 do (
                                temp_uv_mod.selectFaces hideF
                                temp_uv_mod.moveSelected [-10, -10, 0]
                            )
                            return 1
                        )
                        try ( isolateUV {mid} ) catch ( 0 )
                    )"""
                    res_uv = rt.execute(ms_uv)
                    if res_uv == 0: print("🚨 ERRO: O MaxScript não conseguiu manipular as UVs.")
                
                print(f"-> Disparando o Render: {t_rnd_id}")
                # CÓDIGO CORRIGIDO 2: Embrulhado em 'fn'
                ms_bake = f"""(
                    fn doBake fileOut szX szY fileExt cor den elem = (
                        try(freeSceneBitmaps(); vfbControl #clearimage)catch()
                        if cor and den do try ( renderers.current.denoise_enable = true ) catch()
                        $.INodeBakeProperties.removeAllBakeElements()
                        
                        local be = (execute (elem + "()"))
                        be.outputSzX = szX
                        be.outputSzY = szY
                        be.fileType = fileExt
                        be.fileName = fileOut
                        
                        $.INodeBakeProperties.addBakeElement be
                        $.INodeBakeProperties.bakeEnabled = true
                        $.INodeBakeProperties.bakeChannel = 3
                        render rendertype:#bakeSelected vfb:true quiet:true production:true outputfile:fileOut
                        return 1
                    )
                    try ( doBake "{t_rnd_id}" {cr} {cr} "{ext}" {str(i_cor).lower()} {str(u_den).lower()} "{el}" ) catch ( 0 )
                )"""
                res_bk = rt.execute(ms_bake)
                if res_bk == 0: print("🚨 ERRO: O Renderizador rejeitou o comando de Bake!")
                
                wt = 0
                while not os.path.exists(t_rnd_id) and wt < 60: 
                    time.sleep(0.5)
                    wt += 0.5
                    QtWidgets.QApplication.processEvents()
                
                if os.path.exists(t_rnd_id):
                    if i_vr and u_den:
                        print("-> Passando Denoise IA...")
                        c_in = t_rnd_id.replace("/", "\\")
                        c_out = t_rnd_id.replace(".exr", "_denoised.exr").replace("/", "\\")
                        try:
                            p = subprocess.Popen(f'{get_vdenoise_path()} -inputFile="{c_in}" -outputFile="{c_out}" -display=0', shell=True, creationflags=0x08000000); dt = 0
                            while p.poll() is None and dt < 60: QtWidgets.QApplication.processEvents(); time.sleep(0.5); dt += 0.5
                            if p.poll() is None: p.terminate()
                        except: pass
                        f_exr = c_out.replace("\\", "/") if os.path.exists(c_out.replace("\\", "/")) else t_rnd_id
                        rt.execute(f"""( try(vfbControl #clearimage)catch(); try ( local bI = openBitmap @"{f_exr}"; if bI != undefined do ( local bO = bitmap bI.width bI.height filename:@"{t_jpg_id}" hdr:true; copy bI bO; save bO; close bI; close bO; free bI; free bO ) ) catch() )""")
                        rt.execute("freeSceneBitmaps(); gc light:true"); time.sleep(0.5)
                        try: os.remove(t_rnd_id); os.remove(c_out.replace("\\", "/"))
                        except: pass
                        
                if max_id > 1: 
                    print("-> Removendo modificador temporário de UV...")
                    rt.execute("try(deleteModifier $ globalvars.get #temp_uv_mod)catch()")
                
                if os.path.exists(t_jpg_id):
                    d['item'].setText(1, "Bake OK")
                    print("-> SUCESSO: Textura gravada no disco.")
                else:
                    d['item'].setText(1, "Erro Conv")
                    print("🚨 ERRO: A textura não foi salva. Timeout?")
                
            if not ui._is_cancelled: 
                print("-> Restaurando UV final do objeto...")
                rt.execute(f"""( local o = getNodeByName "{d['name']}"; try( ChannelInfo.CopyChannel o 3 3; ChannelInfo.PasteChannel o 3 1; collapseStack o )catch() )""")
            
        except Exception as e: 
            d['item'].setText(1, "Erro Render")
            print(f"🚨 EXCEÇÃO PYTHON: {e}")
            
    ui.pb.setValue(tot)
    ui.pb.setFormat("Bake Concluído!")
    print("\n" + "="*50)
    print("✅ FIM DO LOG DE BAKE")
    print("="*50 + "\n")
    
    if auto_export and not ui._is_cancelled: 
        cld.finalize_export(ui, p_bk, ui.edt_p_glb.text())

def upd_res_col(ui):
    a_r = ui.chk_a_res.isChecked()
    m_r = ui.spn_res.value()
    for d in ui.bake_items:
        o = rt.getNodeByName(d['name'])
        if o:
            try:
                md = max(abs(o.max.x - o.min.x), abs(o.max.y - o.min.y), abs(o.max.z - o.min.z))
                r = 256 if md <= ui.s256.value() else 512 if md <= ui.s512.value() else 1024 if md <= ui.s1024.value() else m_r if a_r else m_r
                d['item'].setText(3, f"{r}px")
                d['item'].setText(4, f"{md:.1f}")
                if r == 256: d['item'].setForeground(3, QtGui.QColor(0, 255, 255))
                elif r == 512: d['item'].setForeground(3, QtGui.QColor(0, 255, 0))
                elif r == 1024: d['item'].setForeground(3, QtGui.QColor(255, 165, 0))
                else: d['item'].setForeground(3, QtGui.QColor(255, 50, 50))
            except: pass

def get_processable_items(ui):
    valid_items = []
    for d in ui.bake_items:
        obj = rt.getNodeByName(d['name'])
        if obj: valid_items.append(d)
    return valid_items