import gradio as gr
import json, time, uuid, os, csv, random
import numpy as np

TRIALS_JSON = "trials.json"
LOG_PATH = "responses.csv"
# Check for persistent storage (Hugging Face Spaces)
if os.path.exists("/data"):
    LOG_PATH = "/data/responses.csv"

def load_trials(path=TRIALS_JSON):
    with open(path, "r") as f:
        return json.load(f)

TRIALS = load_trials()

def ensure_log_header():
    if not os.path.exists(LOG_PATH):
        with open(LOG_PATH, "w", newline="") as f:
            w = csv.writer(f)
            w.writerow([
                "ts", "participant_id", "trial_id", "task",
                "left_method", "right_method",
                "choice", "voxel_choice", "rt_ms",
                "prompt_text", "prompt_image"
            ])

ensure_log_header()

import os, hashlib
import trimesh

UNCOLORED_DIR = "cache_uncolored"
os.makedirs(UNCOLORED_DIR, exist_ok=True)

def to_uncolored_obj(mesh_path: str) -> str:
    # cache by path string (and you can include file mtime if you want)
    key = hashlib.md5(mesh_path.encode("utf-8")).hexdigest()
    out_path = os.path.join(UNCOLORED_DIR, f"{key}.glb")
    if os.path.exists(out_path):
        return out_path

    m = trimesh.load(mesh_path, force="mesh", process=False)

    # --- THIS is the key line: only verts + faces, no visual/material info ---
    # Preserving normals is crucial for correct shading
    m = trimesh.Trimesh(
        vertices=m.vertices, 
        faces=m.faces, 
        # vertex_normals=m.vertex_normals, 
        process=False
    )

    # export as OBJ (no textures/colors)
    m.export(out_path)
    return out_path

def to_gray_obj(mesh_path: str, gray=180) -> str:
    """
    Convert any mesh (obj/glb/gltf/...) to a geometry-only OBJ with constant gray vertex colors.
    gray: 0..255
    """
    # Updated cache key to 'v4_double_sided' to force regeneration
    key = hashlib.md5((mesh_path + f":g{gray}:v4_double_sided").encode("utf-8")).hexdigest()
    out_path = os.path.join(UNCOLORED_DIR, f"{key}.glb")
    if os.path.exists(out_path):
        return out_path

    # process=True fixes winding, merges vertices, etc.
    loaded = trimesh.load(mesh_path, process=True)

    # GLB/GLTF often loads as a Scene -> merge geometries
    if isinstance(loaded, trimesh.Scene):
        if len(loaded.geometry) == 0:
            raise ValueError(f"Empty scene: {mesh_path}")
        m = trimesh.util.concatenate(tuple(loaded.geometry.values()))
    else:
        m = loaded

    # Fix normals to ensure smooth shading (computes them if missing)
    # m.fix_normals()

    # Constant vertex colors (RGBA)
    vc = np.tile(np.array([gray, gray, gray, 255], dtype=np.uint8), (len(m.vertices), 1))
    m.visual = trimesh.visual.ColorVisuals(mesh=m, vertex_colors=vc)
    
    # Enable double-sided rendering to fix visibility issues
    # if hasattr(m.visual, 'material'):
        # m.visual.material.doubleSided = True
    # For some Trimesh versions/types, we might need to ensure it's PBR or handled correctly
    # But usually setting the attribute on the material object helps GLB export.

    # Export as GLB (supports vertex colors + normals well)
    m.export(out_path)
    return out_path

def to_ccm_obj(mesh_path: str, gray=180) -> str:
    """
    Convert any mesh (obj/glb/gltf/...) to a geometry-only OBJ with constant gray vertex colors.
    gray: 0..255
    """
    # Updated cache key to 'v4_double_sided' to force regeneration
    key = hashlib.md5((mesh_path + f":g{gray}:v4_ccm").encode("utf-8")).hexdigest()
    out_path = os.path.join(UNCOLORED_DIR, f"{key}.glb")
    if os.path.exists(out_path):
        return out_path

    # process=True fixes winding, merges vertices, etc.
    loaded = trimesh.load(mesh_path, process=True)

    # GLB/GLTF often loads as a Scene -> merge geometries
    if isinstance(loaded, trimesh.Scene):
        if len(loaded.geometry) == 0:
            raise ValueError(f"Empty scene: {mesh_path}")
        m = trimesh.util.concatenate(tuple(loaded.geometry.values()))
    else:
        m = loaded

    # Fix normals to ensure smooth shading (computes them if missing)
    # m.fix_normals()

    # Constant vertex colors (RGBA)
    # vc = np.tile(np.array([gray, gray, gray, 255], dtype=np.uint8), (len(m.vertices), 1))
    num_vertices = len(m.vertices)
    # color based on coord
    vc = m.vertices.copy()
    vc = (vc - vc.min(axis=0)) / (vc.max(axis=0) - vc.min(axis=0))
    vc = (vc * 255).astype(np.uint8)
    vc_alpha = np.ones((num_vertices, 1), dtype=np.uint8) * 255
    vc = np.concatenate([vc, vc_alpha], axis=1)
    # normalize vertices to 0,1

    m.visual = trimesh.visual.ColorVisuals(mesh=m, vertex_colors=vc)
    
    # Enable double-sided rendering to fix visibility issues
    # if hasattr(m.visual, 'material'):
        # m.visual.material.doubleSided = True
    # For some Trimesh versions/types, we might need to ensure it's PBR or handled correctly
    # But usually setting the attribute on the material object helps GLB export.

    # Export as GLB (supports vertex colors + normals well)
    m.export(out_path)
    return out_path

def to_normal_obj(mesh_path: str, gray=180) -> str:
    """
    Convert any mesh (obj/glb/gltf/...) to a geometry-only OBJ with constant gray vertex colors.
    gray: 0..255
    """
    # Updated cache key to 'v4_double_sided' to force regeneration
    key = hashlib.md5((mesh_path + f":g{gray}:v4_normal").encode("utf-8")).hexdigest()
    out_path = os.path.join(UNCOLORED_DIR, f"{key}.glb")
    if os.path.exists(out_path):
        return out_path

    # process=True fixes winding, merges vertices, etc.
    loaded = trimesh.load(mesh_path, process=True)

    # GLB/GLTF often loads as a Scene -> merge geometries
    if isinstance(loaded, trimesh.Scene):
        if len(loaded.geometry) == 0:
            raise ValueError(f"Empty scene: {mesh_path}")
        m = trimesh.util.concatenate(tuple(loaded.geometry.values()))
    else:
        m = loaded

    # Fix normals to ensure smooth shading (computes them if missing)
    m.fix_normals()

    # Normalize normals to 0..1 range for visualization
    # Normal range is [-1, 1], so (n + 1) / 2 maps to [0, 1]
    vc = (m.vertex_normals + 1) / 2
    vc = (vc * 255).astype(np.uint8)
    
    # Add alpha channel
    num_vertices = len(m.vertices)
    vc_alpha = np.ones((num_vertices, 1), dtype=np.uint8) * 255
    vc = np.concatenate([vc, vc_alpha], axis=1)

    m.visual = trimesh.visual.ColorVisuals(mesh=m, vertex_colors=vc)
    
    # Enable double-sided rendering to fix visibility issues
    # if hasattr(m.visual, 'material'):
        # m.visual.material.doubleSided = True
    # For some Trimesh versions/types, we might need to ensure it's PBR or handled correctly
    # But usually setting the attribute on the material object helps GLB export.

    # Export as GLB (supports vertex colors + normals well)
    m.export(out_path)
    return out_path

def _trial_to_view(trial, flip_lr: bool):
    if flip_lr:
        left_method, right_method = "continuous", "ours"
        left_voxel, right_voxel = trial["cont_voxel"], trial["ours_voxel"]
        left_mesh,  right_mesh  = trial["cont_mesh"],  trial["ours_mesh"]
        left_video = trial.get("cont_video", None)
        right_video = trial.get("ours_video", None)
    else:
        left_method, right_method = "ours", "continuous"
        left_voxel, right_voxel = trial["ours_voxel"], trial["cont_voxel"]
        left_mesh,  right_mesh  = trial["ours_mesh"],  trial["cont_mesh"]
        left_video = trial.get("ours_video", None)
        right_video = trial.get("cont_video", None)

    # Make sure voxel OBJ is geometry-only (and optionally also strip mesh if you want)
    left_voxel  = to_ccm_obj(left_voxel)
    right_voxel = to_ccm_obj(right_voxel)

    # (Optional) also strip mesh to geometry-only if needed:
    left_mesh  = to_normal_obj(left_mesh)   # only if you want to force mesh untextured via OBJ export
    right_mesh = to_normal_obj(right_mesh)

    task = trial["task"]
    prompt_text = trial.get("prompt_text", "")
    prompt_image = trial.get("prompt_image", None)

    show_text = (task == "text")
    show_img = (task == "image")

    return dict(
        trial_id=trial["trial_id"],
        task=task,
        show_text=show_text,
        show_img=show_img,
        prompt_text=prompt_text,
        prompt_image=prompt_image,
        left_method=left_method,
        right_method=right_method,
        left_voxel=left_voxel,
        right_voxel=right_voxel,
        left_mesh=left_mesh,
        right_mesh=right_mesh,
        left_video=left_video,
        right_video=right_video,
    )

def start_session(seed=None):
    pid = str(uuid.uuid4())[:8]
    rng = random.Random(seed if seed is not None else time.time())

    order = list(range(len(TRIALS)))
    rng.shuffle(order)

    # DEBUG: move specific trial to front
    # target_id = "img_treehouse_rmapple"
    # for i, t in enumerate(TRIALS):
    #     if t["trial_id"] == target_id:
    #         if i in order:
    #             order.remove(i)
    #             order.insert(0, i)
    #         break

    state = {
        "pid": pid,
        "order": order,
        "i": 0,
        "t0": time.time(),
        "flip_lr": rng.choice([False, True]),  # for first trial
        "rng_state": rng.getstate(),
    }

    view = _trial_to_view(TRIALS[order[0]], state["flip_lr"])
    return (
        state, pid,
        gr.update(value=f"Trial 1 / {len(TRIALS)}"),
        gr.update(visible=view["show_text"], value=view["prompt_text"]),
        gr.update(visible=view["show_img"], value=view["prompt_image"]),
        gr.update(value="Left Model"),
        gr.update(value="Right Model"),
        view["left_voxel"], view["right_voxel"],
        view["left_mesh"], view["right_mesh"],
        view["left_video"], view["right_video"],
        gr.update(value=""),  # completion code
        gr.update(value=None), # reset choice
        gr.update(value=None)  # reset voxel_choice
    )

def submit(choice, voxel_choice, state):
    i = state["i"]
    idx = state["order"][i]
    trial = TRIALS[idx]
    rt_ms = int((time.time() - state["t0"]) * 1000)

    # log
    with open(LOG_PATH, "a", newline="") as f:
        w = csv.writer(f)
        w.writerow([
            time.time(), state["pid"], trial["trial_id"], trial["task"],
            ("continuous" if state["flip_lr"] else "ours"),
            ("ours" if state["flip_lr"] else "continuous"),
            choice, voxel_choice, rt_ms,
            trial.get("prompt_text", ""),
            trial.get("prompt_image", ""),
        ])

    # advance
    state["i"] += 1
    if state["i"] >= len(TRIALS):
        code = f"COMPLETE-{state['pid']}"
        return (
            state,
            gr.update(value="Done ✅ (copy the completion code below)"),
            gr.update(visible=False, value=""),
            gr.update(visible=False, value=None),
            gr.update(value="LEFT"),
            gr.update(value="RIGHT"),
            None, None,
            None, None,
            None, None,
            gr.update(value=code),
            gr.update(value=None),
            gr.update(value=None)
        )

    # next trial
    rng = random.Random()
    rng.setstate(state["rng_state"])
    state["flip_lr"] = rng.choice([False, True])
    state["rng_state"] = rng.getstate()
    state["t0"] = time.time()

    ni = state["i"]
    nidx = state["order"][ni]
    view = _trial_to_view(TRIALS[nidx], state["flip_lr"])

    return (
        state,
        gr.update(value=f"Trial {ni+1} / {len(TRIALS)}"),
        gr.update(visible=view["show_text"], value=view["prompt_text"]),
        gr.update(visible=view["show_img"], value=view["prompt_image"]),
        gr.update(value="Left Model"),
        gr.update(value="Right Model"),
        view["left_voxel"], view["right_voxel"],
        view["left_mesh"], view["right_mesh"],
        view["left_video"], view["right_video"],
        gr.update(value=""),
        gr.update(value=None),
        gr.update(value=None)
    )

def get_responses_file(password):
    # Set your password in HF Space Settings as ADMIN_PASSWORD, or use default "123456"
    correct_pass = os.environ.get("ADMIN_PASSWORD", "123456")
    if password != correct_pass:
        raise gr.Error("Incorrect Password")
    
    if not os.path.exists(LOG_PATH):
        raise gr.Error("No responses collected yet.")
        
    return LOG_PATH

def clear_responses(password):
    correct_pass = os.environ.get("ADMIN_PASSWORD", "123456")
    if password != correct_pass:
        raise gr.Error("Incorrect Password")
        
    if os.path.exists(LOG_PATH):
        os.remove(LOG_PATH)
        ensure_log_header() # Re-create header immediately
        return "Responses cleared successfully."
    else:
        return "No responses file found to clear."

with gr.Blocks(css="""
#prompt-img {max-width: 320px; margin: 0 auto;}
#prompt-img img {width: 100% !important; height: auto !important;}
.model3d-window {height: 600px !important;}
""") as demo:
    st = gr.State()

    gr.Markdown("## 3D Asset User Study (Image + Text)")

    with gr.Accordion("📝 Instructions (Click to Read)", open=True):
        gr.Markdown("""
        ### How to complete this study:
        1. **Observe the Prompt**: You will see a text description or a reference image.
        2. **Compare Models**: Rotate and zoom into the "Left" and "Right" 3D models to inspect them.
        3. **Vote**: Select which model, voxel (Geometry and/or Texture) better matches the prompt or has higher quality. 
        4. **Submit**: Click "Submit" to move to the next trial.
        
        *You could consider:
        1. Alignment with prompt.
        2. Stuctural quality, such as holes, cracks, floaters, etc.
        3. Since the process of voxel->mesh can be lossy (missing parts), please also inspect voxels' quality.
        4. If you are not sure about the winner, pick not sure.
        5. Unexpected holes on voxels are typically not desired.*
        """)

    pid_box = gr.Textbox(label="Participant ID (auto)", interactive=False)
    progress = gr.Markdown()

    # Prompt area
    prompt_text = gr.Markdown(visible=False)
    prompt_img = gr.Image(label="Prompt Image", visible=False, elem_id="prompt-img")

    gr.Markdown("### Compare the two methods")

    with gr.Row():
        with gr.Column():
            left_title = gr.Textbox(label="Left Method", interactive=False)
            left_voxel = gr.Model3D(label="Left Voxel (untextured)", clear_color=[1,1,1,1], elem_classes="model3d-window")
            left_mesh  = gr.Model3D(label="Left Mesh (untextured)",  clear_color=[1,1,1,1], elem_classes="model3d-window")
            left_vid   = gr.Video(label="Left Video (optional)", autoplay=True)

        with gr.Column():
            right_title = gr.Textbox(label="Right Method", interactive=False)
            right_voxel = gr.Model3D(label="Right Voxel (untextured)", clear_color=[1,1,1,1], elem_classes="model3d-window")
            right_mesh  = gr.Model3D(label="Right Mesh (untextured)",  clear_color=[1,1,1,1], elem_classes="model3d-window")
            right_vid   = gr.Video(label="Right Video (optional)", autoplay=True)

    voxel_choice = gr.Radio(["Left", "Right", "Not Sure"], label="Which generated voxel is better?", value=None)
    choice = gr.Radio(["Left", "Right", "Not Sure"], label="Which generated mesh is better?", value=None)
    
    btn = gr.Button("Submit")

    code_box = gr.Textbox(label="Completion Code", interactive=False)

    demo.load(
        start_session,
        inputs=None,
        outputs=[st, pid_box, progress, prompt_text, prompt_img,
                left_title, right_title,
                left_voxel, right_voxel, left_mesh, right_mesh, left_vid, right_vid,
                code_box, choice, voxel_choice]
    )

    btn.click(
        submit,
        inputs=[choice, voxel_choice, st],
        outputs=[st, progress, prompt_text, prompt_img,
                left_title, right_title,
                left_voxel, right_voxel, left_mesh, right_mesh, left_vid, right_vid,
                code_box, choice, voxel_choice]
    )

    gr.Markdown("---")
    with gr.Accordion("Admin Zone (Download Responses)", open=False):
        with gr.Row():
            admin_pass = gr.Textbox(label="Admin Password", type="password", placeholder="Enter password to download")
            download_btn = gr.Button("Download CSV")
            clear_btn = gr.Button("Clear Responses", variant="stop")
        
        with gr.Row():
            admin_file = gr.File(label="Responses File", interactive=False)
            admin_status = gr.Textbox(label="Status", interactive=False)

        download_btn.click(
            get_responses_file,
            inputs=[admin_pass],
            outputs=[admin_file]
        )
        
        clear_btn.click(
            clear_responses,
            inputs=[admin_pass],
            outputs=[admin_status]
        )

demo.launch(allowed_paths=["/data"])