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 deep_log_and_kill():
    try:
        time.sleep(0.5); tgts = ["3d66", "V-Ray", "Buffer", "RGBA", "Render"]
        for w in rt.windows.getChildrenHWND(0):
            for t in tgts:
                if t.lower() in str(w[4]).lower(): rt.windows.sendMessage(w[0], 0x0010, 0, 0); rt.windows.sendMessage(w[0], 0x0002, 0, 0)
    except: pass

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)")
        if r_list:
            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 (Atlas)..."); QtWidgets.QApplication.processEvents()
    ms = """(
        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 collect o
        for h in heads do (
            local cg = for c in h.children where superclassof c == GeometryClass and classof c != TargetObject collect c
            if cg.count > 0 do (
                local b = cg[1]; if not isKindOf b Editable_Poly do try(convertToPoly b)catch()
                if isKindOf b Editable_Poly do (
                    for i = 2 to cg.count do ( local n = 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; delete h; act += 1
                )
            )
        )
        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, "Atlas Welder", f"Sucesso! {act} Grupos fundidos.\nObjetos soltos preservados.")
            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 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 <= tg: d['item'].setText(1, "Já < Meta")
                elif cp < 50: d['item'].setText(1, "Geo Base")
                else:
                    rt.select(o); p = 0
                    while cp > tg and cp >= 50 and p < 20:
                        if not rt.isKindOf(o, rt.Editable_Poly): rt.convertTo(o, rt.Editable_Poly)
                        opt = rt.ProOptimizer(); rt.addModifier(o, opt); opt.KeepTextures = False; opt.KeepNormals = False; opt.Calculate = True; opt.VertexPercent = sp
                        rt.collapseStack(o)
                        if not rt.isKindOf(o, rt.Editable_Poly): rt.convertTo(o, rt.Editable_Poly)
                        np = rt.getPolygonCount(o)[0]
                        if np >= cp or np == 0: break
                        cp = np; p += 1; QtWidgets.QApplication.processEvents() 
                    d['item'].setText(2, f"{cp:,}"); d['item'].setText(1, f"Opt ({p}x)")
            except: d['item'].setText(1, "Erro Opt")
    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(); act = 0; rt.execute("max modify mode"); tgs = ui.get_processable_items()
    fn = []; pn = [d['name'] for d in tgs]
    for d in ui.bake_items:
        if d['name'] not in pn: fn.append(d['name'])
    tot = len(tgs); ui.pb.setMaximum(tot); ui.pb.setValue(0)
    for i, d in enumerate(tgs):
        if ui._is_cancelled: break
        ui.pb.setFormat(f"Cortando ({i+1}/{tot}): {d['name']}..."); ui.pb.setValue(i); QtWidgets.QApplication.processEvents()
        chk = [d['name']]; sc = 0 
        while len(chk) > 0 and sc < 300:
            QtWidgets.QApplication.processEvents()
            if ui._is_cancelled: break
            cn = chk.pop(0); o = rt.getNodeByName(cn)
            if not o: continue
            md = max(abs(o.max.x - o.min.x), abs(o.max.y - o.min.y), abs(o.max.z - o.min.z))
            if md > thr:
                ax = "X" if md == abs(o.max.x - o.min.x) else "Y" if md == abs(o.max.y - o.min.y) else "Z"
                ms = """(
                    fn bsp objN ax = (
                        local o = getNodeByName objN; if o == undefined or not (canConvertTo o Editable_Poly) do return #()
                        if not isKindOf o Editable_Poly do try(convertToPoly o)catch(); if not isKindOf o Editable_Poly do return #()
                        local v = polyop.getNumVerts o; if v > 0 do ( try(o.weldThreshold = 0.001)catch(); try(polyop.weldVertsByThreshold o #{{1..v}})catch() )
                        CenterPivot o; local o2 = copy o name:(uniqueName (o.name + "_F"))
                        local rot = if ax == "X" then eulerAngles 0 90 0 else if ax == "Y" then eulerAngles 90 0 0 else eulerAngles 0 0 0
                        local s1 = SliceModifier(); s1.Slice_Type = 2; s1.slice_plane.rotation = rot; addModifier o s1; try(addModifier o (Cap_Holes()))catch(); collapseStack o
                        local s2 = SliceModifier(); s2.Slice_Type = 3; s2.slice_plane.rotation = rot; addModifier o2 s2; try(addModifier o2 (Cap_Holes()))catch(); collapseStack o2
                        if not isKindOf o Editable_Poly do try(convertToPoly o)catch(); if not isKindOf o2 Editable_Poly do try(convertToPoly o2)catch()
                        local f1 = if isValidNode o then polyop.getNumFaces o else 0; local f2 = if isValidNode o2 then polyop.getNumFaces o2 else 0
                        local r = #(); if f1 < 5 then try(delete o)catch() else append r o.name; if f2 < 5 then try(delete o2)catch() else append r o2.name; return r
                    ); bsp "{0}" "{1}"
                )""".format(cn, ax)
                try:
                    n_p = rt.execute(ms)
                    if n_p and len(n_p) > 1:
                        act += 1
                        for p in n_p:
                            po = rt.getNodeByName(p)
                            if po:
                                pd = max(abs(po.max.x - po.min.x), abs(po.max.y - po.min.y), abs(po.max.z - po.min.z))
                                if pd >= md * 0.98: fn.append(p)
                                else: chk.append(p)
                    else:
                        for p in n_p: fn.append(p) if n_p else fn.append(cn)
                except: fn.append(cn)
            else: fn.append(cn)
            sc += 1
    ui.pb.setValue(tot); ui.pb.setFormat("Fatiador Concluído!")
    rt.clearSelection(); sel = [rt.getNodeByName(n) for n in fn if rt.getNodeByName(n)]; rt.select(sel)
    if act > 0: load_selection(ui) 

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:
                if not rt.isKindOf(o, rt.Editable_Poly): rt.convertTo(o, rt.Editable_Poly)
                rt.execute(f"""( local o = getNodeByName "{d['name']}"; local m = Unwrap_UVW(); addModifier o m; m.setMapChannel 1; m.setTVSubObjectMode 3; local nf = polyOp.getNumFaces o; m.selectFaces #{{1..nf}}; m.flattenMap 45.0 #() 0.002 true 0 true false; m.pack 1 0.002 true false false; collapseStack o )""")
                if not rt.isKindOf(o, rt.Editable_Poly): rt.convertTo(o, rt.Editable_Poly)
                d['item'].setText(1, "UV Packed")
            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):
    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)

    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()
        t_jpg = os.path.join(p_bk, f"{d['name']}_B.jpg").replace("\\", "/")
        if os.path.exists(t_jpg): d['item'].setText(1, "Já existe"); continue 
        
        cr = ui.spn_res.value(); o = rt.getNodeByName(d['name'])
        if o and 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 = t_jpg.replace(".jpg", ".exr") if ext == ".exr" else t_jpg
        
        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)
        
        try:
            rt.execute("freeSceneBitmaps(); try(vfbControl #clearimage)catch()") 
            ms = """(
                local o = getNodeByName "{0}"; if {5} and {6} do try ( renderers.current.denoise_enable = true ) catch()
                renderWidth = {2}; renderHeight = {2}; o.INodeBakeProperties.removeAllBakeElements()
                local be = {1}(); be.outputSzX = {2}; be.outputSzY = {2}; be.fileType = "{3}"; be.fileName = "{4}"
                o.INodeBakeProperties.addBakeElement be; o.INodeBakeProperties.bakeEnabled = true; o.INodeBakeProperties.bakeChannel = 1
                render rendertype:#bakeSelected vfb:true quiet:true outputfile:"{4}"
            )""".format(d['name'], el, cr, ext, t_rnd, str(i_cor).lower(), str(u_den).lower())
            rt.execute(ms)
             
            wt = 0
            while not os.path.exists(t_rnd) and wt < 60: time.sleep(0.5); wt += 0.5; QtWidgets.QApplication.processEvents()
            
            if os.path.exists(t_rnd):
                if i_vr and u_den:
                    d['item'].setText(1, "IA Limpando...")
                    c_in = t_rnd.replace("/", "\\"); c_out = t_rnd.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
                    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}" 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); os.remove(c_out.replace("\\", "/"))
                    except: pass
                    
                if os.path.exists(t_jpg): d['item'].setText(1, "Bake OK"); deep_log_and_kill()
                else: d['item'].setText(1, "Erro Conv")
            else: d['item'].setText(1, "Erro Save")
        except: d['item'].setText(1, "Erro Render")
    ui.pb.setValue(tot); ui.pb.setFormat("Bake Concluído!")
    if auto_export and not ui._is_cancelled: cld.finalize_export(ui, p_bk, ui.edt_p_glb.text())