import json
import logging
import shutil
import time
import uuid
from pathlib import Path

from runtime import ROOT, MODEL_ROOT, configure_storage, startup_checks
configure_storage()
logging.basicConfig(filename=ROOT/"logs/workflow.log",level=logging.INFO,
                    format="%(asctime)s %(levelname)s %(message)s")
runtime_report = startup_checks()

import gradio as gr
import torch
import trimesh
from preflight import inspect_glb
import upstream.app as mia

mia.init_models()
# Initialize existing upstream component references used by its workflow.
# The original UI is not launched.
upstream_blocks = mia.init_blocks()
print("MODELS_READY dedicated CUDA inference",flush=True)

def preserve_and_inspect(path):
    if not path:
        raise gr.Error("Upload a GLB first.")
    folder = ROOT/"inputs"/uuid.uuid4().hex
    folder.mkdir()
    original = folder/"original.glb"
    shutil.copyfile(path,original)
    try:
        report = inspect_glb(original)
    except (ValueError,KeyError,IndexError,TypeError) as exc:
        report = {"inference_allowed":False,"warning":"Invalid or unsupported GLB: "+str(exc)}
    report.update(original=str(original),cuda_available=torch.cuda.is_available(),gpu=runtime_report["gpu"])
    (folder/"preflight.json").write_text(json.dumps(report,indent=2))
    return str(original),report,str(original)

def execute(original,report):
    if not original or not report or not report.get("inference_allowed"):
        raise gr.Error((report or {}).get("warning") or "Inspect an eligible GLB before inference.")
    if not torch.cuda.is_available():
        raise gr.Error("CUDA is unavailable; CPU inference is disabled.")
    # Recheck the preserved file before expensive processing.
    checked=inspect_glb(original)
    if not checked["inference_allowed"]:
        raise gr.Error(checked["warning"])
    job=ROOT/"outputs"/uuid.uuid4().hex
    job.mkdir()
    working=job/"character.glb"
    shutil.copyfile(original,working)
    db=mia.DB()
    started=time.time()
    stages=("Loading mesh","Preparing coordinates","Predicting skeleton and weights",
            "Preparing rest pose","Exporting rigged asset","Finishing")
    try:
        for stage,_ in zip(stages,mia._pipeline(input_path=str(working),db=db,
                                               reset_to_rest=True,rest_pose_type=None)):
            yield stage,None,None,None
        for p in (db.anim_path,db.anim_vis_path):
            if not p or not Path(p).is_file() or Path(p).stat().st_size==0:
                raise RuntimeError("Upstream export did not produce a nonempty asset")
        # Validate export parsing without changing it.
        result=trimesh.load(db.anim_vis_path,force="scene",process=False)
        if not result.geometry:
            raise RuntimeError("Exported GLB contains no geometry")
        with open(db.anim_vis_path,"rb") as stream:
            import struct
            stream.seek(12)
            length,kind=struct.unpack("<II",stream.read(8))
            doc=json.loads(stream.read(length))
        if not doc.get("skins") or not any("skin" in n for n in doc.get("nodes",[])):
            raise RuntimeError("Exported GLB has no skin binding")
        evidence=dict(status="PASS",original=original,glb=db.anim_vis_path,fbx=db.anim_path,
                      elapsed_seconds=time.time()-started,export_meshes=len(result.geometry),
                      skin_count=len(doc["skins"]),runtime=runtime_report["startup_id"])
        (job/"result.json").write_text(json.dumps(evidence,indent=2))
        logging.info("WORKFLOW_PASS %s",json.dumps(evidence))
        print("WORKFLOW_PASS "+json.dumps(evidence),flush=True)
        yield "Complete: rigged GLB and FBX exported.",db.anim_vis_path,db.anim_vis_path,db.anim_path
    except Exception as exc:
        logging.exception("Workflow failed in %s",job)
        (job/"error.txt").write_text(str(exc))
        raise gr.Error("Workflow failed: "+str(exc)+". Original preserved; details in /data/logs/workflow.log.")
    finally:
        mia.clear(db)

def example():
    return preserve_and_inspect(MODEL_ROOT/"data/examples/ninja.glb")

def dense_character_test():
    # Densify an existing character only for preflight; never modify its source.
    source=MODEL_ROOT/"data/examples/ninja.glb"
    mesh=trimesh.load(source,force="mesh",process=False)
    while len(mesh.faces)<=500_000:
        mesh=mesh.subdivide()
    folder=ROOT/"inputs"/("dense-test-"+uuid.uuid4().hex)
    folder.mkdir()
    path=folder/"dense-character.glb"
    mesh.export(path)
    report=inspect_glb(path)
    assert not report["inference_allowed"],"Dense character must be blocked before inference"
    report["test"]="PASS: dense character preprocessing warning"
    report["source"]=str(source)
    report["original"]=str(path)
    (folder/"preflight.json").write_text(json.dumps(report,indent=2))
    print("DENSE_CHARACTER_TEST "+json.dumps(report),flush=True)
    return report

with gr.Blocks(title="Make-It-Animatable") as demo:
    gr.Markdown("# Motion Rig\nUpload a character GLB, inspect it, then prepare and export its rig.")
    original=gr.State()
    report_state=gr.State()
    upload=gr.File(label="Upload GLB",file_types=[".glb"],type="filepath")
    inspect_button=gr.Button("Inspect / preflight")
    report_view=gr.JSON(label="Mesh preflight")
    original_download=gr.File(label="Preserved original",interactive=False)
    run=gr.Button("Prepare and export rig",variant="primary")
    status=gr.Textbox(label="Status",interactive=False)
    preview=gr.Model3D(label="Export preview",interactive=False)
    with gr.Row():
        glb=gr.File(label="Download rigged GLB",interactive=False)
        fbx=gr.File(label="Download rigged FBX",interactive=False)
    with gr.Accordion("Runtime and acceptance checks",open=False):
        gr.JSON(value=runtime_report,label="Startup GPU and dependency checks")
        normal_test=gr.Button("Load bundled normal character for testing")
        dense_test=gr.Button("Test dense character preflight")
        dense_result=gr.JSON(label="Dense test result")
    inspect_button.click(preserve_and_inspect,inputs=upload,
                         outputs=[original,report_view,original_download]).success(
        lambda x:x,inputs=report_view,outputs=report_state)
    normal_test.click(example,outputs=[original,report_view,original_download]).success(
        lambda x:x,inputs=report_view,outputs=report_state)
    run.click(execute,inputs=[original,report_state],outputs=[status,preview,glb,fbx],
              concurrency_limit=1,concurrency_id="workflow")
    dense_test.click(dense_character_test,outputs=dense_result,
                     concurrency_limit=1,concurrency_id="workflow")
demo.queue(max_size=4,default_concurrency_limit=1).launch(
    server_name="0.0.0.0",server_port=7860,show_error=True,ssr_mode=False,
    allowed_paths=[str(ROOT/"inputs"),str(ROOT/"outputs")])