import sys
import os
from dotenv import load_dotenv
load_dotenv()
import json
import asyncio
import psutil
from llm_panda.metrics_logger import LANGSMITH_AVAILABLE, log_system_metrics, log_interaction, log_episode_metrics

import math
import time
import threading
import uuid
from typing import Optional
from fastapi import FastAPI, HTTPException, Request, WebSocket, WebSocketDisconnect
from fastapi.staticfiles import StaticFiles
from fastapi.responses import HTMLResponse, Response
import uvicorn
import cv2
import numpy as np
import pybullet as pb
try:
    import edge_tts
except ImportError:  # Browser Vietnamese speech remains the runtime fallback.
    edge_tts = None
# Ensure project source root is in path
PROJECT_SRC = os.path.abspath(os.path.join(os.path.dirname(__file__), ".."))
if PROJECT_SRC not in sys.path:
    sys.path.insert(0, PROJECT_SRC)

from main import PandaApp
from llm_panda.llm_agent import LLMAgent, RemoteLLMPlanner, RuleBasedPlanner
from llm_panda.plan_verifier import PlanVerifier
from llm_panda import config
from llm_panda.execution.cancellation import ExecutionInterrupted
from llm_panda import logger as _logger_setup
_logger_setup.setup()
from frontend.web_runtime import WebRuntime
import pybullet_data
from pybullet_object_models import ycb_objects
from fastapi.middleware.cors import CORSMiddleware
app = FastAPI(title="LLM-Panda Web-Embedded PyBullet Simulation")
app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)
# Serve static files from the frontend directory itself
static_dir = os.path.dirname(__file__)
app.mount("/static", StaticFiles(directory=static_dir), name="static")
# Mount robot meshes and YCB object assets
panda_data_path = os.path.join(pybullet_data.getDataPath(), "franka_panda")
ycb_data_path = ycb_objects.getDataPath()
app.mount("/assets/panda", StaticFiles(directory=panda_data_path), name="panda")
app.mount("/assets/ycb", StaticFiles(directory=ycb_data_path), name="ycb")
if os.path.exists("G:/PyBullet_Datasets/ReplicaCAD_dataset"):
    app.mount("/assets/replicacad", StaticFiles(directory="G:/PyBullet_Datasets/ReplicaCAD_dataset"), name="replicacad")
@app.get("/")
async def get():
    html_file = os.path.join(static_dir, "index.html")
    if os.path.exists(html_file):
        with open(html_file, "r", encoding="utf-8") as f:
            return HTMLResponse(
                content=f.read(),
                status_code=200,
                headers={
                    "Cache-Control": "no-store, no-cache, must-revalidate",
                    "Pragma": "no-cache",
                    "Expires": "0",
                },
            )
    return HTMLResponse(content="<h1>index.html not found!</h1>", status_code=404)


@app.post("/api/speech/synthesize")
async def synthesize_speech(request: Request):
    if edge_tts is None:
        raise HTTPException(
            status_code=503,
            detail="edge-tts is not installed",
        )
    body = await request.json()
    text = str(body.get("text", "")).strip()
    if not text:
        raise HTTPException(status_code=400, detail="Text is empty")
    if len(text) > 2000:
        raise HTTPException(status_code=413, detail="Text is too long")
    voice = os.environ.get(
        "EDGE_TTS_VOICE",
        "vi-VN-NamMinhNeural",
    ).strip() or "vi-VN-NamMinhNeural"
    try:
        communicator = edge_tts.Communicate(text, voice)
        chunks = []
        async for chunk in communicator.stream():
            if chunk.get("type") == "audio":
                chunks.append(chunk["data"])
        audio = b"".join(chunks)
    except Exception as exc:
        raise HTTPException(
            status_code=502,
            detail="Vietnamese TTS service is unavailable",
        ) from exc
    if not audio:
        raise HTTPException(status_code=502, detail="TTS returned no audio")
    return Response(content=audio, media_type="audio/mpeg")


import psutil
from dataclasses import dataclass

MAX_SESSIONS = 100
current_sim = None


class SessionContext:
    def __init__(self, session_id: str):
        self.session_id = session_id
        self.sim = current_sim
        self.executor_lock = asyncio.Lock()
        self.pybullet_lock = threading.RLock()
        self.pipeline_executing = False
        self.web_runtime = WebRuntime()
        self.command_worker_task = None
        self.physics_settle_task = None
        self.drag_state = {"constraint_id": None, "active": False}
        self.last_active = time.time()
        self.active_websockets = 0

active_sessions: dict[str, SessionContext] = {}

def check_server_capacity() -> bool:
    if len(active_sessions) >= MAX_SESSIONS:
        return False
    try:
        mem = psutil.virtual_memory()
        if mem.percent > 90.0:
            return False
    except Exception:
        pass
    return True



def _finite_vec3(value):
    return (
        isinstance(value, (list, tuple))
        and len(value) >= 3
        and all(
            isinstance(component, (int, float))
            and math.isfinite(component)
            for component in value[:3]
        )
    )


def _clamp_manual_target(position, target_type="object"):
    """Keep browser-controlled targets inside the supported workspace."""
    if not _finite_vec3(position):
        raise ValueError("Manual target must contain three finite coordinates")
    x, y, z = (float(value) for value in position[:3])
    x = min(max(x, -0.10), 0.90)
    y = min(max(y, -0.65), 0.65)
    z = min(max(z, 0.03), 0.90)

    cx, cy = config.TABLE_CENTER
    hx, hy = config.TABLE_HALF_EXTENTS[:2]
    if cx - hx <= x <= cx + hx and cy - hy <= y <= cy + hy:
        z = max(z, config.TABLE_TOP_Z + 0.01)
    tx, ty, tz = config.TRAY_POSITION
    thx, thy = config.TRAY_HALF_EXTENTS[:2]
    if tx - thx <= x <= tx + thx and ty - thy <= y <= ty + thy:
        z = max(z, tz + 0.01)
    return [x, y, z]


def _interpolate_manual_path(start, target, max_step):
    """Return bounded linear steps used only by browser drag handling."""
    if len(start) != len(target) or max_step <= 0:
        raise ValueError("Invalid manual drag interpolation")
    max_delta = max(
        (abs(float(end) - float(begin)) for begin, end in zip(start, target)),
        default=0.0,
    )
    step_count = max(1, int(math.ceil(max_delta / max_step)))
    return [
        [
            float(begin) + (float(end) - float(begin)) * step / step_count
            for begin, end in zip(start, target)
        ]
        for step in range(1, step_count + 1)
    ]


def _load_environment_config(path):
    with open(path, "r", encoding="utf-8") as handle:
        data = json.load(handle)
    if not isinstance(data, dict):
        raise ValueError("File JSON không hợp lệ.")
    if "robot" not in data:
        raise ValueError("Thiếu trường 'robot'.")
    has_legacy = "table" in data and "tray" in data
    has_structures = "structures" in data
    if not (has_legacy or has_structures):
        raise ValueError("Thiếu trường 'table'/'tray' hoặc 'structures'.")
    if "objects" in data and not isinstance(data["objects"], list):
        raise ValueError("Trường objects của environment phải là một danh sách.")
    return data


def _list_environment_files(directory):
    files = []
    for filename in sorted(os.listdir(directory)):
        if not filename.endswith(".json") or filename.endswith(".semantics.json"):
            continue
        path = os.path.join(directory, filename)
        try:
            _load_environment_config(path)
        except (OSError, ValueError, json.JSONDecodeError):
            continue
        files.append(filename)
    return files
def get_init_payload(ctx: SessionContext) -> dict:
    sim = ctx.sim
    """Gathers the initial static geometry and YCB object definitions."""
    with ctx.pybullet_lock:
        objects_info = {}
        if sim and sim.scene:
            for oid in sim.scene.object_ids:
                name = sim.scene.get_name(oid)
                cls = sim.scene.get_class(oid)
                folder = ""
                web_obj = ""
                web_mtl = ""
                for obj_cfg in sim.env_builder.env_data.get("objects", []):
                    if obj_cfg["name"] == name:
                        folder = obj_cfg.get("folder", "")
                        web_obj = obj_cfg.get("web_obj", "")
                        web_mtl = obj_cfg.get("web_mtl", "")
                        break
                pos, ori = pb.getBasePositionAndOrientation(oid, physicsClientId=sim.client)
                
                # Query collision shape data for CoM-to-link offset
                collision_data = pb.getCollisionShapeData(oid, -1, physicsClientId=sim.client)
                coll_scale = [1.0, 1.0, 1.0]
                coll_local_pos = [0.0, 0.0, 0.0]
                coll_local_ori = [0.0, 0.0, 0.0, 1.0]
                if collision_data:
                    cd = collision_data[0]
                    coll_scale = list(cd[3])
                    coll_local_pos = list(cd[5])
                    coll_local_ori = list(cd[6])
                # Query visual shape data for exact dimensions/scales and offsets relative to link
                visual_data = pb.getVisualShapeData(oid, physicsClientId=sim.client)
                vis_scale = [1.0, 1.0, 1.0]
                vis_local_pos = [0.0, 0.0, 0.0]
                vis_local_ori = [0.0, 0.0, 0.0, 1.0]
                if visual_data:
                    vd = visual_data[0]
                    vis_scale = list(vd[3])
                    # Transform visual frame to be relative to CoM: CoM_to_visual = CoM_to_link * link_to_visual
                    vis_pos_com, vis_ori_com = pb.multiplyTransforms(
                        coll_local_pos, coll_local_ori,
                        list(vd[5]), list(vd[6])
                    )
                    vis_local_pos = list(vis_pos_com)
                    vis_local_ori = list(vis_ori_com)
                    
                    if not web_obj and vd[4]:
                        mesh_path = vd[4].decode('utf-8')
                        web_obj = os.path.basename(mesh_path)
                        if not folder:
                            folder = os.path.basename(os.path.dirname(mesh_path))
                        candidate_mtls = [
                            mesh_path + ".mtl",
                            os.path.splitext(mesh_path)[0] + ".mtl",
                        ]
                        matching_mtl = next(
                            (
                                candidate
                                for candidate in candidate_mtls
                                if os.path.isfile(candidate)
                            ),
                            None,
                        )
                        if not web_mtl and matching_mtl:
                            web_mtl = os.path.basename(matching_mtl)
                        
                objects_info[name] = {
                    "class": cls,
                    "folder": folder,
                    "web_obj": web_obj,
                    "web_mtl": web_mtl,
                    "initial_pos": list(pos),
                    "initial_ori": list(ori),
                    "vis_scale": vis_scale,
                    "vis_local_pos": vis_local_pos,
                    "vis_local_ori": vis_local_ori,
                    "coll_scale": coll_scale,
                    "coll_local_pos": coll_local_pos,
                    "coll_local_ori": coll_local_ori
                }
        ycb_path = ycb_objects.getDataPath()
        catalog = {}
        try:
            folders = [f for f in os.listdir(ycb_path) if os.path.isdir(os.path.join(ycb_path, f)) and f.startswith("Ycb")]
            for folder in folders:
                name_key = folder.replace("Ycb", "")
                catalog[name_key] = {"folder": folder}
            print(f"DEBUG: ycb_path={ycb_path}, catalog has {len(catalog)} items")
        except Exception as e:
            print(f"DEBUG ERROR in catalog: {e}")

        return {
            "type": "init",
            "structures": sim.env_builder.env_data.get("structures", []),
            "table": {
                "center": [config.TABLE_CENTER[0], config.TABLE_CENTER[1], config.TABLE_BASE_Z],
                "half_extents": list(config.TABLE_HALF_EXTENTS),
                "color": list(getattr(config, 'TABLE_COLOR', [0.8, 0.8, 0.8, 1.0])),
                "top_z": config.TABLE_TOP_Z
            },
            "tray": {
                "position": list(config.TRAY_POSITION),
                "half_extents": list(config.TRAY_HALF_EXTENTS),
                "color": list(getattr(config, 'TRAY_COLOR', [0.2, 0.2, 0.8, 1.0]))
            },
            "robot_base": list(sim.env_builder.env_data["robot"]["base_position"]),
            "slots": {name: list(pos) for name, pos in getattr(config, 'SLOTS', {}).items()},
            "objects": objects_info,
            "catalog": catalog
        }
def get_simulation_state_payload(ctx: SessionContext) -> dict:
    sim = ctx.sim
    """Gathers the current robot joint positions and YCB object poses."""
    with ctx.pybullet_lock:
        if sim is None or sim.client is None or not pb.isConnected(sim.client):
            return {}
        
        # Read all joints (12 joints including arm joints and gripper fingers)
        robot_id = sim.controller.robot_id
        num_joints = pb.getNumJoints(robot_id, physicsClientId=sim.client)
        joints_state = []
        for i in range(num_joints):
            pos, _, _, _ = pb.getJointState(robot_id, i, physicsClientId=sim.client)
            joints_state.append(pos)
            
        # Read objects pos/ori
        objects_state = {}
        for oid in sim.scene.object_ids:
            name = sim.scene.get_name(oid)
            pos, ori = pb.getBasePositionAndOrientation(oid, physicsClientId=sim.client)
            objects_state[name] = {
                "pos": list(pos),
                "ori": list(ori)
            }
        
        return {
            "type": "state",
            "joints": joints_state,
            "objects": objects_state
        }
def set_pybullet_window_visibility(visible: bool):
    """Finds the PyBullet GUI window on Windows and sets its visibility by moving it on/off-screen."""
    import sys
    if sys.platform != "win32":
        return False
    import ctypes
    import time
    
    EnumWindowsProc = ctypes.WINFUNCTYPE(ctypes.c_bool, ctypes.c_void_p, ctypes.c_void_p)
    
    for _ in range(30):
        found_hwnd = []
        
        def foreach_window(hwnd, lParam):
            length = ctypes.windll.user32.GetWindowTextLengthW(hwnd)
            buff = ctypes.create_unicode_buffer(length + 1)
            ctypes.windll.user32.GetWindowTextW(hwnd, buff, length + 1)
            title = buff.value
            if "bullet physics examplebrowser" in title.lower():
                found_hwnd.append(hwnd)
                return False
            return True
            
        ctypes.windll.user32.EnumWindows(EnumWindowsProc(foreach_window), 0)
        if found_hwnd:
            hwnd = found_hwnd[0]
            # Ensure the window is shown/restored from minimized state first (SW_SHOW = 5)
            ctypes.windll.user32.ShowWindow(hwnd, 5)
            
            # SWP_NOSIZE = 0x0001, SWP_NOZORDER = 0x0004
            if visible:
                # Move window to visible region
                ctypes.windll.user32.SetWindowPos(hwnd, 0, 100, 100, 0, 0, 0x0001 | 0x0004)
            else:
                # Move window far off-screen to hide it without destroying OpenGL context
                ctypes.windll.user32.SetWindowPos(hwnd, 0, -10000, -10000, 0, 0, 0x0001 | 0x0004)
            print(f"[Window Visibility] Set PyBullet window (HWND: {hwnd}) visibility to: {visible}", flush=True)
            return True
        time.sleep(0.1)
    return False
def compute_and_format_score(ctx: SessionContext, plan: list, executed_successfully: bool):
    sim = ctx.sim
    """Calculates rearrangement score based on final object locations and collisions."""
    with ctx.pybullet_lock:
        score = 100
        details = []
        
        # Check target placement if we had a plan
        target_location = None
        target_object = None
        for step in plan:
            if step.get("function") in ("place_at", "place_down", "move_to"):
                args = step.get("args", {})
                if "location_id" in args:
                    target_location = args["location_id"]
            if step.get("function") == "pick_up":
                target_object = step.get("args", {}).get("object_id")
                
        if target_object and target_location:
            obj_id = sim.scene.name_to_id(target_object)
            if obj_id is not None:
                pos, _ = pb.getBasePositionAndOrientation(obj_id)
                current_loc = sim.scene.classify_location(pos)
                if current_loc == target_location:
                    details.append(f"✅ **Đã đặt {target_object} đúng vị trí {target_location}**: +50 điểm")
                else:
                    details.append(f"❌ **{target_object} chưa đến đúng vị trí {target_location} (đang ở {current_loc})**: -50 điểm")
                    score -= 50
            else:
                details.append(f"❌ **Không tìm thấy vật thể {target_object}**: -50 điểm")
                score -= 50
        else:
            details.append("ℹ️ Không có yêu cầu di chuyển vật thể cụ thể nào được phát hiện.")
            
        # Knock detection
        knocked_count = 0
        for oid in sim.scene.object_ids:
            if oid == sim.scene.name_to_id(target_object):
                continue
            pos, _ = pb.getBasePositionAndOrientation(oid)
            loc = sim.scene.classify_location(pos)
            if loc == "elsewhere":
                name = sim.scene.get_name(oid)
                details.append(f"⚠️ **Vật thể '{name}' bị lệch khỏi vị trí quy định**: -15 điểm")
                score -= 15
                knocked_count += 1
                
        if knocked_count == 0:
            details.append("✅ **Không có vật thể nào khác bị va chạm hay xê dịch**: +30 điểm")
        else:
            score = max(0, score)
            
        # Execution success
        if executed_successfully:
            details.append("✅ **Thực thi toàn bộ kế hoạch thành công**: +20 điểm")
        else:
            details.append("❌ **Thực thi gặp lỗi gián đoạn**: -30 điểm")
            score -= 30
            
        score = max(0, min(100, score))
        
        md = f"### 🏆 Rearrangement Evaluation Report\n\n**Tổng điểm: {score}/100**\n\n**Chi tiết:**\n" + "\n".join([f"- {d}" for d in details])
        return score, md

def run_pipeline_threadsafe(
    user_prompt: str,
    loop,
    use_mock: bool,
    ctx: SessionContext,
    cancellation_token,
    generation: int,
    task_id: str,
    camera_state: Optional[dict] = None,
):
    """Runs the planning, validation, execution with image streaming, and scoring using closed-loop execution monitor."""
    ctx.pipeline_executing = True
    ctx.sim.pipeline.set_cancellation_token(cancellation_token)

    def publish(payload: dict, journal: bool = True):
        payload = {**payload, "task_id": task_id, "generation": generation}
        if getattr(ctx, "episode_run_tree", None) is not None:
            payload["run_id"] = str(ctx.episode_run_tree.id)
        ctx.web_runtime.publish_threadsafe(payload, loop, journal=journal)
    
    def log(message: str, step: Optional[int] = None):
        try:
            print(f"[WS Log] {message}", flush=True)
        except Exception:
            pass
        publish({"type": "log", "message": message, "step": step})

    # Step 1: Planning
    log("⏳ [1/5] Tôi đang suy nghĩ cách thực hiện...")
    
    if use_mock:
        planner = RuleBasedPlanner()
        log("✅ [1] Tôi đang sử dụng RuleBasedPlanner (Ngoại tuyến)")
    else:
        api_key = os.environ.get("OPENAI_API_KEY") or os.environ.get("LLM_API_KEY")
        if not api_key:
            log("❌ [1] Lỗi: Không tìm thấy OPENAI_API_KEY / LLM_API_KEY. Vui lòng chọn Mock.")
            ctx.pipeline_executing = False
            return
        planner = RemoteLLMPlanner()
        log(f"✅ [1] Tôi đang phân tích yêu cầu...")
        
    with ctx.pybullet_lock:
        ctx.sim.pipeline.agent = LLMAgent(
            planner=planner,
            conversation_manager=ctx.web_runtime.conversation_manager,
        )
        
    try:
        MAX_REPLANS = max(0, int(os.environ.get("MAX_REPLANS", "3")))
    except ValueError:
        MAX_REPLANS = 3
    last_failed_sig = None
    last_failed_state_sig = None
    resolved_goal = None
    
    # Hook pb.stepSimulation to capture states during movements
    original_step = pb.stepSimulation
    last_stream_time = [0.0]
    throttle_interval = 0.033
    
    def wrapped_step():
        cancellation_token.checkpoint()
        original_step()
        if not ctx.sim.gui:
            time.sleep(0.001)
        curr_time = time.time()
        if curr_time - last_stream_time[0] >= throttle_interval:
            last_stream_time[0] = curr_time
            state_payload = get_simulation_state_payload(ctx)
            if state_payload:
                publish(state_payload, journal=False)

    for attempt in range(MAX_REPLANS + 1):
        cancellation_token.checkpoint()
        if attempt > 0:
            ctx.web_runtime.set_status("replanning", generation=generation)
            publish({"type": "execution_status", "state": "replanning"})
            log(f"🔄 [MONITOR] Đang thử lập kế hoạch lại lần {attempt}/{MAX_REPLANS} do: {last_failed_sig}")
            log("⏳ [1/5] Tôi đang suy nghĩ cách thực hiện...")
            
        with ctx.pybullet_lock:
            # Settle physics in case the user placed objects hovering or overlapping in the UI
            if ctx.sim.controller:
                ctx.sim.controller._settle(60)
                
            # Define reachability check for scene state
            def reachability_fn(oid):
                cls = ctx.sim.scene.get_class(oid)
                try:
                    grasp_pos, grasp_ori = ctx.sim.pipeline.grasp_planner.compute_grasp_pose(oid, cls)
                    return ctx.sim.controller.is_reachable(grasp_pos, grasp_ori)
                except Exception:
                    return False
            carried_name = ctx.sim.scene.get_name(ctx.sim.pipeline._carried) if ctx.sim.pipeline._carried is not None else None
            scene_state = ctx.sim.scene.get_scene_state(reachability_fn=reachability_fn, carried_object=carried_name)
            state_text = json.dumps(scene_state, indent=2)

            if attempt == 0:
                ctx.episode_run_tree = None
                if LANGSMITH_AVAILABLE:
                    from langsmith.run_trees import RunTree
                    ctx.episode_run_tree = RunTree(
                        name="robot_episode",
                        run_type="chain",
                        inputs={
                            "user_command": user_prompt,
                            "state_before": scene_state
                        }
                    )
                    ctx.episode_run_tree.post()
            
            # Signature of the physical state to detect if environment actually changed
            state_sig_dict = {
                "objects": [{"name": o["name"], "location": o["location"], "blocking": o.get("blocking_objects", [])} for o in scene_state.get("objects", [])],
                "carried": scene_state.get("carried_object")
            }
            state_sig = json.dumps(state_sig_dict, sort_keys=True)
            
        try:
            with ctx.pybullet_lock:
                llm_start_time = time.time()
                planning_instruction = (
                    resolved_goal.canonical_instruction
                    if attempt > 0 and resolved_goal is not None
                    else user_prompt
                )
                ctx.sim.pipeline.agent.last_outcome = None
                plan_kwargs = {
                    "previous_failed_reasons": last_failed_sig,
                }
                if attempt > 0:
                    plan_kwargs["resolve_goal"] = False
                plan = ctx.sim.pipeline.agent.plan(
                    planning_instruction,
                    state_text,
                    **plan_kwargs,
                )
                outcome = ctx.sim.pipeline.agent.last_outcome
                if attempt == 0 and outcome is not None:
                    resolved_goal = outcome.resolved_goal
                    
                if getattr(ctx, "episode_run_tree", None) is not None:
                    # Update intent and plan to trace
                    # LangSmith automatically handles children if wrapped, but we can also store explicitly
                    pass # Handled by LLMAgent wrapping if we did that, or we'll just end the tree at the end.
                llm_time = time.time() - llm_start_time
                cancellation_token.checkpoint()
                
            # Discard an LLM result that arrived after the user pressed stop.
            cancellation_token.checkpoint()

            # Summarize plan
            targets = []
            destinations = set()
            for step_dict in plan:
                if step_dict.get("function") == "pick_up":
                    obj_id = step_dict.get("args", {}).get("object_id")
                    if obj_id: targets.append(obj_id)
                elif step_dict.get("function") == "move_to":
                    loc_id = step_dict.get("args", {}).get("location_id")
                    if loc_id: destinations.add(loc_id)
            
            def translate_obj(x):
                obj_map = {
                    'banana': 'quả chuối', 
                    'soup_can': 'lon súp', 
                    'chips_can': 'lon chip', 
                    'tennis_ball': 'quả bóng tennis', 
                    'pear': 'quả lê', 
                    'clamp': 'cái kẹp', 
                    'table': 'bàn', 
                    'tray': 'khay'
                }
                return obj_map.get(x, x)
                
            targets_vi = [translate_obj(t) for t in targets]
            dests_vi = [translate_obj(d) for d in destinations]
            
            if plan and plan[0].get("function") in ("clarify", "reject"):
                summary_msg = plan[0].get("args", {}).get("message", "Yêu cầu không hợp lệ.")
            elif targets_vi:
                target_str = ", ".join(targets_vi[:-1]) + " và " + targets_vi[-1] if len(targets_vi) > 1 else targets_vi[0]
                dest_str = " và ".join(dests_vi) if dests_vi else "vị trí mới"
                summary_msg = f"Tôi sẽ gắp {target_str} và di chuyển tới {dest_str}."
            else:
                summary_msg = "Tôi đã xử lý xong yêu cầu của bạn."
                
            log(f"✅ [1] {summary_msg}")
        except ExecutionInterrupted:
            raise
        except Exception as e:
            log(f"❌ [1] Lập kế hoạch thất bại: {e}")
            ctx.pipeline_executing = False
            return
            
        # Send steps to client
        plan_steps = []
        for idx, step_item in enumerate(plan, 1):
            func = step_item.get("function")
            args = step_item.get("args", {})
            obj = args.get("object_id") or args.get("location_id") or ""
            plan_steps.append({
                "step": idx,
                "skill": func,
                "object": obj,
                "args": args
            })
            
        publish({"type": "plan", "plan": plan_steps, "summary": summary_msg, "llm_time": llm_time})
        
        if plan and plan[0].get("function") in ("clarify", "reject"):
            log(f"Robot phản hồi: {summary_msg}")
            header_title = "Phản hồi từ Robot" if plan[0].get("function") == "clarify" else "Yêu cầu bị từ chối"
            score_md = f"### {header_title}\n\n**Robot phản hồi:** {summary_msg}\n"
            publish({
                "type": "result",
                "response_type": plan[0].get("function"),
                "score_display": score_md,
                "score_value": 0,
                "success": False,
            })
            state_payload = get_simulation_state_payload(ctx)
            if state_payload:
                publish(state_payload, journal=False)
            ctx.pipeline_executing = False
            return

        if not plan:
            log("⚠️ [1] Kế hoạch rỗng.")
            log("⏳ [5/5] Đang tính toán điểm sắp xếp...")
            score_val, score_md = compute_and_format_score(ctx, [], True)
            publish({
                "type": "result",
                "score_display": score_md,
                "score_value": score_val,
            })
            state_payload = get_simulation_state_payload(ctx)
            if state_payload:
                publish(state_payload, journal=False)
            ctx.pipeline_executing = False
            return
            
        # Step 2: Validation
        cancellation_token.checkpoint()
        log("⏳ [2/5] Tôi đang kiểm tra mức độ an toàn của kế hoạch...")
        try:
            with ctx.pybullet_lock:
                warnings = PlanVerifier.check(plan, scene_state)
            cancellation_token.checkpoint()
            if not warnings:
                log("✅ [2] Kế hoạch an toàn và hợp lệ.")
            else:
                log(f"⚠️ [2] Phát hiện {len(warnings)} rủi ro tiềm ẩn.")
                for w in warnings:
                    log(f"   - {w}")
        except ExecutionInterrupted:
            raise
        except Exception as e:
            log(f"❌ [2] Lỗi kiểm tra kế hoạch: {e}")
            ctx.pipeline_executing = False
            return
            
        # Step 3: PyBullet Execution (Live!)
        cancellation_token.checkpoint()
        log("⏳ [3/5] Tôi bắt đầu thực hiện các bước...")
        ctx.web_runtime.set_status("executing", generation=generation)
        publish({"type": "execution_status", "state": "executing"})
        
        with ctx.pybullet_lock:
            pb.stepSimulation = wrapped_step
            
        executed_successfully = True
        step_failure_reason = None
        failed_step_index = -1
        failed_step_data = None
        
        try:
            # Trigger an initial simulation state transmission
            state_payload = get_simulation_state_payload(ctx)
            if state_payload:
                publish(state_payload, journal=False)
                
            def format_skill_vn(step_data):
                fn = step_data.get("function", "")
                args = step_data.get("args", {})
                target = args.get("object_id") or args.get("location_id") or ""
                
                obj_map = {
                    'banana': 'quả chuối',
                    'soup_can': 'lon súp',
                    'chips_can': 'lon snack',
                    'tennis_ball': 'quả bóng tennis',
                    'pear': 'quả lê',
                    'clamp': 'chiếc kìm',
                    'table': 'bàn',
                    'tray': 'khay',
                }
                target_vn = obj_map.get(target, target)
                
                if fn == "pick_up": return f"Gắp {target_vn} lên"
                if fn == "move_to": return f"Di chuyển đến {target_vn}"
                if fn == "place_down": return "Đặt đồ vật xuống"
                if fn == "open_gripper": return "Mở tay kẹp"
                if fn == "close_gripper": return "Đóng tay kẹp"
                return f"{fn}"
                
            expected_location = None
            for i, step in enumerate(plan, 1):
                cancellation_token.checkpoint()
                ctx.web_runtime.set_status("executing", generation=generation, current_step=i)
                publish({"type": "current_step", "step": i})
                friendly_text = format_skill_vn(step)
                # Make the first letter lowercase for natural sentence integration
                if friendly_text:
                    friendly_text = friendly_text[0].lower() + friendly_text[1:]
                log(f"🚀 [3/5] Tôi đang {friendly_text}...")
                
                pre_positions = {
                    oid: pb.getBasePositionAndOrientation(oid)[0]
                    for oid in ctx.sim.scene.object_ids
                    if oid != ctx.sim.pipeline._carried
                }
                held_before = ctx.sim.pipeline._carried
                ctx.sim.controller.reset_stuck()
                
                with ctx.pybullet_lock:
                    ctx.sim.pipeline._execute(step, verbose=True)
                    
                # Run step verification using execution monitor
                ok, reason = ctx.sim.pipeline.execution_monitor.verify_step(
                    step, held_before, expected_location, ctx.sim.pipeline._carried, ctx.sim.pipeline._pending_target_pos
                )
                if step.get("function") == "move_to":
                    expected_location = step.get("args", {}).get("location_id")
                    
                if not ok:
                    executed_successfully = False
                    step_failure_reason = reason
                    failed_step_index = i
                    failed_step_data = step
                    break
                    
            if executed_successfully:
                log("✅ [3] Tôi đã hoàn tất các bước.")
            else:
                log(f"❌ [3] Thực thi lỗi tại bước {failed_step_index} '{failed_step_data.get('function')}': {step_failure_reason}")
        except ExecutionInterrupted:
            with ctx.pybullet_lock:
                ctx.sim.pipeline.interrupt_hold()
            raise
        except Exception as e:
            log(f"❌ Lỗi trong lúc thực thi mô phỏng: {e}")
            executed_successfully = False
            step_failure_reason = str(e)
        finally:
            with ctx.pybullet_lock:
                pb.stepSimulation = original_step
                
            # Send a final state to ensure visual sync
            state_payload = get_simulation_state_payload(ctx)
            if state_payload:
                publish(state_payload, journal=False)
                
        # Now check overall goal using execution monitor
        goal_completed = False
        goal_reason = ""
        cancellation_token.checkpoint()
        
        with ctx.pybullet_lock:
            final_scene_state = ctx.sim.scene.get_scene_state(
                reachability_fn=reachability_fn,
                carried_object=(
                    ctx.sim.scene.get_name(ctx.sim.pipeline._carried)
                    if ctx.sim.pipeline._carried is not None
                    else None
                ),
            )
            final_state_text = json.dumps(final_scene_state, indent=2)

        if executed_successfully:
            goal_completed, goal_reason = (
                ctx.sim.pipeline.execution_monitor.verify_goal(
                    user_prompt,
                    final_state_text,
                    resolved_goal=resolved_goal,
                )
            )
            cancellation_token.checkpoint()
            if goal_completed:
                log(f"✅ [KIỂM TRA MỤC TIÊU] ĐẠT: {goal_reason}")
            else:
                log(f"⚠️ [KIỂM TRA MỤC TIÊU] THẤT BẠI: {goal_reason}")
                
        if goal_completed:
            break
            
        # If we failed or goal check failed, set replan inputs
        if not executed_successfully:
            last_failed_sig = f"Step {failed_step_index} '{failed_step_data.get('function')}' failed: {step_failure_reason}"
        else:
            last_failed_sig = f"Goal Check Failed: {goal_reason}"
            
        unreachable_objects = [
            o["name"] for o in final_scene_state.get("objects", [])
            if not o.get("reachable", True)
        ]
        
        last_failed_sig += f"\n\nCURRENT ENVIRONMENT STATE (At the moment of failure):\n{final_state_text}"
        if unreachable_objects:
            last_failed_sig += f"\n\nWARNING: The following objects are currently UNREACHABLE: {unreachable_objects}. You must NOT include them in your next plan."
            
        last_failed_state_sig = state_sig
        
        if attempt < MAX_REPLANS:
            cancellation_token.checkpoint()
            log(f"⚠️ [MONITOR] Phát hiện lỗi. Đang khôi phục cánh tay robot và lập kế hoạch lại (thử lần {attempt + 1})...")
            with ctx.pybullet_lock:
                ctx.sim.pipeline._recover_to_safe_state(verbose=True)
            cancellation_token.checkpoint()
        else:
            log(f"❌ [MONITOR] Đã vượt quá số lần lập kế hoạch lại ({MAX_REPLANS}). Hủy bỏ tác vụ.")
            
    # Step 5: Scoring
    cancellation_token.checkpoint()
    score_val, score_md = compute_and_format_score(ctx, plan, executed_successfully and goal_completed)
    
    error_cause = None
    if not plan:
        error_cause = "llm"
    elif not executed_successfully:
        error_cause = "ik"
    elif not goal_completed:
        error_cause = "slip"
        
    is_success = executed_successfully and goal_completed and score_val >= 50

    # Calculate PostHog metrics
    token_usage = {"prompt_tokens": 0, "completion_tokens": 0}
    if hasattr(ctx.sim.pipeline.agent, "planner") and hasattr(ctx.sim.pipeline.agent.planner, "last_token_usage"):
        token_usage = getattr(ctx.sim.pipeline.agent.planner, "last_token_usage")
    
    system_ram_mb = psutil.virtual_memory().used / (1024 * 1024)
    timeout_network_error = False
    plan_validity_first_try = False
    collision_free = True
    kinematic_error = False
    
    if last_failed_sig:
        if "Rate limit" in last_failed_sig or "Timeout" in last_failed_sig:
            timeout_network_error = True
        if "Kinematic" in last_failed_sig or "IK" in last_failed_sig or "outside workspace" in last_failed_sig.lower():
            kinematic_error = True
        if "collision" in last_failed_sig.lower():
            collision_free = False
            
    if attempt == 0 and not last_failed_sig:
        plan_validity_first_try = True

    try:
        log_episode_metrics(
            session_id=str(uuid.uuid4()),  # Or an actual session if tracked
            llm_latency_sec=0, # Fallback
            token_usage=token_usage,
            system_ram_mb=system_ram_mb,
            timeout_network_error=timeout_network_error,
            plan_validity_first_try=plan_validity_first_try,
            collision_free=collision_free,
            kinematic_error=kinematic_error,
            replan_count=attempt
        )
    except Exception as e:
        log(f"Failed to send PostHog metrics: {e}")

    if getattr(ctx, "episode_run_tree", None) is not None:
        import dataclasses
        outputs = {
            "success": is_success,
            "score": score_val,
            "error_cause": error_cause,
            "plan": plan,
        }
        if resolved_goal:
            outputs["intent"] = resolved_goal.canonical_instruction
            outputs["goal"] = [dataclasses.asdict(p) for p in resolved_goal.predicates]
            
        ctx.episode_run_tree.end(outputs=outputs)
        try:
            ctx.episode_run_tree.patch()
        except Exception as e:
            log(f"Failed to patch LangSmith run: {e}")

    publish({
        "type": "result",
        "score_display": score_md,
        "score_value": score_val,
        "error_cause": error_cause,
        "success": is_success,
    })
    ctx.sim.pipeline.set_cancellation_token(None)
    ctx.pipeline_executing = False


async def command_worker_loop(loop, ctx: SessionContext):
    """Serialize robot commands while allowing the newest prompt to interrupt."""
    
    try:
        while True:
            command = ctx.web_runtime.take_pending()
            if command is None:
                return
            generation = command["generation"]
            task_id = command["task_id"]
            interrupted = False
            await ctx.web_runtime.publish({
                "type": "command_accepted",
                "prompt": command["prompt"],
                "task_id": task_id,
                "generation": generation,
            })
            await ctx.web_runtime.publish({
                "type": "execution_status",
                "state": "planning",
                "task_id": task_id,
                "generation": generation,
            })
            try:
                await asyncio.to_thread(
                    run_pipeline_threadsafe,
                    command["prompt"],
                    loop,
                    command["use_mock"],
                    ctx,
                    command["token"],
                    generation,
                    task_id,
                )
                # Close the race where stop is pressed just as the worker
                # returns from its final blocking operation.
                command["token"].checkpoint()
            except ExecutionInterrupted as exc:
                interrupted = True
                ctx.pipeline_executing = False
                with ctx.pybullet_lock:
                    if ctx.sim is not None and ctx.sim.pipeline is not None:
                        ctx.sim.pipeline.interrupt_hold()
                ctx.web_runtime.set_status("interrupted", generation=generation)
                state_payload = await asyncio.to_thread(
                    get_simulation_state_payload,
                    ctx,
                )
                if state_payload:
                    await ctx.web_runtime.publish(state_payload, journal=False)
                ctx.web_runtime.conversation_manager.record_interruption(
                    command["prompt"]
                )
                await ctx.web_runtime.publish({
                    "type": "interrupted",
                    "reason": str(exc),
                    "task_id": task_id,
                    "generation": generation,
                })
            except Exception as exc:
                ctx.pipeline_executing = False
                await ctx.web_runtime.publish({
                    "type": "result",
                    "success": False,
                    "error_cause": "runtime",
                    "score_display": f"### Lỗi thực thi\n\n{exc}",
                    "task_id": task_id,
                    "generation": generation,
                })
            finally:
                if ctx.sim is not None and ctx.sim.pipeline is not None:
                    ctx.sim.pipeline.set_cancellation_token(None)
                ctx.web_runtime.complete(
                    generation,
                    state="interrupted" if interrupted else "completed",
                )
    finally:
        ctx.command_worker_task = None
        if ctx.web_runtime.has_pending:
            ctx.command_worker_task = asyncio.create_task(command_worker_loop(loop))


async def wait_for_active_task_stop(ctx: SessionContext, timeout=5.0):
    ctx.web_runtime.cancel_all("runtime_mutation")
    deadline = time.monotonic() + timeout
    while ctx.web_runtime.active_token is not None and time.monotonic() < deadline:
        await asyncio.sleep(0.02)
    return ctx.web_runtime.active_token is None



import psutil

async def system_health_loop():
    while True:
        try:
            from llm_panda.metrics_logger import log_system_metrics
            cpu = psutil.cpu_percent(interval=1.0)
            ram = psutil.virtual_memory().percent
            active = len(active_sessions)
            log_system_metrics(active, cpu, ram)
        except Exception:
            pass
        await asyncio.sleep(60)

@app.on_event("startup")
async def startup_event():
    asyncio.create_task(system_health_loop())

@app.websocket("/ws")
async def websocket_endpoint(websocket: WebSocket):
    session_id = websocket.query_params.get("session_id")
    if not session_id:
        await websocket.close(code=1008, reason="Missing session_id")
        return

    # Check capacity if new session
    is_new = session_id not in active_sessions
    if is_new and not check_server_capacity():
        await websocket.accept()
        await websocket.send_json({
            "type": "system_busy", 
            "message": "Hệ thống đã có người sử dụng, vui lòng thử lại sau."
        })
        await websocket.close(code=1013)
        return

    await websocket.accept()
    
    if is_new:
        ctx = SessionContext(session_id)
        active_sessions[session_id] = ctx
    else:
        ctx = active_sessions[session_id]
        
    ctx.active_websockets += 1
    ctx.last_active = time.time()
    await ctx.web_runtime.subscribe(websocket)
    loop = asyncio.get_running_loop()
    
    client_wants_gui = websocket.query_params.get("gui", "false").lower() == "true"
    gui = False
    
    camera_state = {
        "target": [0.5, 0.0, 0.3],
        "dist": 1.3,
        "yaw": 45.0,
        "pitch": -35.0
    }
    
    async def physics_settle_loop():
        try:
            while True:
                if not ctx.pipeline_executing:
                    moving = False
                    with ctx.pybullet_lock:
                        if ctx.sim is not None and ctx.sim.client is not None:
                            for oid in ctx.sim.scene.object_ids:
                                try:
                                    vel, ang_vel = pb.getBaseVelocity(oid, physicsClientId=ctx.sim.client)
                                    if (vel[0]**2 + vel[1]**2 + vel[2]**2 > 1e-4) or (ang_vel[0]**2 + ang_vel[1]**2 + ang_vel[2]**2 > 1e-3):
                                        moving = True
                                        break
                                except Exception:
                                    pass
                    if moving:
                        async with ctx.executor_lock:
                            def step():
                                with ctx.pybullet_lock:
                                    for _ in range(8):
                                        pb.stepSimulation(physicsClientId=ctx.sim.client)
                            await asyncio.to_thread(step)
                        
                        state_payload = await asyncio.to_thread(get_simulation_state_payload, ctx)
                        if state_payload:
                            await ctx.web_runtime.publish(state_payload, journal=False)
                await asyncio.sleep(1/30.0)
        except asyncio.CancelledError:
            pass
        except Exception as e:
            print(f"physics_settle_loop error: {e}")
            
    async with ctx.executor_lock:
        # PandaApp is now global
        await asyncio.to_thread(set_pybullet_window_visibility, client_wants_gui)
        
    try:
        env_files = _list_environment_files(
            os.path.join(PROJECT_SRC, "envs")
        )
        await websocket.send_json({"type": "env_list", "envs": env_files})
    
        init_payload = await asyncio.to_thread(get_init_payload, ctx)
        ctx.web_runtime.latest_init = init_payload
        await websocket.send_json(init_payload)
        
        state_payload = await asyncio.to_thread(get_simulation_state_payload, ctx)
        if state_payload:
            ctx.web_runtime.latest_state = state_payload
            await websocket.send_json(state_payload)
        await websocket.send_json(ctx.web_runtime.snapshot())

        if ctx.physics_settle_task is None or ctx.physics_settle_task.done():
            ctx.physics_settle_task = asyncio.create_task(physics_settle_loop())
        
        while True:
            data = await websocket.receive_text()
            ctx.last_active = time.time()
            payload = json.loads(data)
        
            p_type = payload.get("type")
            if p_type == "run":
                prompt = payload.get("prompt", "")
                use_mock = payload.get("mock", False)
                if prompt.strip():
                    command = ctx.web_runtime.submit(prompt.strip(), use_mock)
                    await ctx.web_runtime.publish({
                        "type": "interrupt_requested",
                        "generation": command["generation"],
                        "task_id": command["task_id"],
                    })
                    if ctx.command_worker_task is None or ctx.command_worker_task.done():
                        ctx.command_worker_task = asyncio.create_task(
                            command_worker_loop(loop, ctx)
                        )

            elif p_type == "cancel":
                if ctx.web_runtime.cancel_all("user_cancelled"):
                    await ctx.web_runtime.publish({
                        "type": "interrupt_requested",
                        "generation": ctx.web_runtime.active_generation,
                    })
            
            
            elif p_type == "mouse_pick_start":
                if ctx.web_runtime.active_token is not None:
                    if not await wait_for_active_task_stop(ctx):
                        await websocket.send_json({
                            "type": "drag_rejected",
                            "reason": "Robot chưa dừng an toàn.",
                        })
                        continue
                rayFrom = payload.get("rayFrom")
                rayTo = payload.get("rayTo")
                requested_type = payload.get("target_type")
                requested_name = payload.get("object_id")
                hit_position = payload.get("hit_position")
                if requested_type != "object":
                    continue
                if (
                    _finite_vec3(rayFrom)
                    and _finite_vec3(rayTo)
                    and _finite_vec3(hit_position)
                ):
                    async with ctx.executor_lock:
                        def do_pick():
                            with ctx.pybullet_lock:
                                if ctx.drag_state.get("constraint_id") is not None:
                                    pb.removeConstraint(
                                        ctx.drag_state["constraint_id"],
                                        physicsClientId=ctx.sim.client,
                                    )
                                drag_id = uuid.uuid4().hex

                                object_id = ctx.sim.scene.name_to_id(
                                    requested_name
                                )
                                if (
                                    object_id is None
                                    or ctx.sim.pipeline._carried == object_id
                                ):
                                    return None
                                mass = pb.getDynamicsInfo(
                                    object_id,
                                    -1,
                                    physicsClientId=ctx.sim.client,
                                )[0]
                                if mass <= 0:
                                    return None
                                aabb_min, aabb_max = pb.getAABB(
                                    object_id,
                                    -1,
                                    physicsClientId=ctx.sim.client,
                                )
                                min_extent = min(
                                    max(0.0, high - low)
                                    for low, high in zip(aabb_min, aabb_max)
                                )
                                ccd_radius = max(
                                    0.001,
                                    min(0.02, min_extent * 0.2),
                                )
                                drag_force = max(
                                    30.0,
                                    min(250.0, float(mass) * 150.0),
                                )
                                pb.changeDynamics(
                                    object_id,
                                    -1,
                                    ccdSweptSphereRadius=ccd_radius,
                                    contactProcessingThreshold=0.0,
                                    physicsClientId=ctx.sim.client,
                                )
                                position, orientation = (
                                    pb.getBasePositionAndOrientation(
                                        object_id,
                                        physicsClientId=ctx.sim.client,
                                    )
                                )
                                inverse_position, inverse_orientation = (
                                    pb.invertTransform(
                                        position,
                                        orientation,
                                    )
                                )
                                local_pivot, _ = pb.multiplyTransforms(
                                    inverse_position,
                                    inverse_orientation,
                                    hit_position,
                                    [0, 0, 0, 1],
                                )
                                constraint_id = pb.createConstraint(
                                    parentBodyUniqueId=object_id,
                                    parentLinkIndex=-1,
                                    childBodyUniqueId=-1,
                                    childLinkIndex=-1,
                                    jointType=pb.JOINT_POINT2POINT,
                                    jointAxis=[0, 0, 0],
                                    parentFramePosition=local_pivot,
                                    childFramePosition=hit_position,
                                    physicsClientId=ctx.sim.client,
                                )
                                pb.changeConstraint(
                                    constraint_id,
                                    hit_position,
                                    maxForce=drag_force,
                                    physicsClientId=ctx.sim.client,
                                )
                                ctx.drag_state.update({
                                    "drag_id": drag_id,
                                    "target_type": "object",
                                    "object_name": requested_name,
                                    "constraint_id": constraint_id,
                                    "active": True,
                                    "target_pos": list(hit_position[:3]),
                                    "max_force": drag_force,
                                })
                                return dict(ctx.drag_state)
                        started = await asyncio.to_thread(do_pick)

                    if not started:
                        await websocket.send_json({
                            "type": "drag_rejected",
                            "reason": (
                                "Không thể kéo vật đang được robot giữ, "
                                "vật tĩnh hoặc vật không còn tồn tại."
                            ),
                        })
                        continue
                    await websocket.send_json({
                        "type": "drag_started",
                        "drag_id": started["drag_id"],
                        "target_type": started["target_type"],
                        "object_id": started.get("object_name"),
                    })

                    if (
                        started["target_type"] == "object"
                        and (
                            ctx.drag_state.get("task") is None
                            or ctx.drag_state["task"].done()
                        )
                    ):
                        async def drag_loop():
                            while ctx.drag_state.get("active") and ctx.drag_state.get("constraint_id") is not None:
                                async with ctx.executor_lock:
                                    def step():
                                        with ctx.pybullet_lock:
                                            for _ in range(4):
                                                pb.stepSimulation(physicsClientId=ctx.sim.client)
                                    await asyncio.to_thread(step)
                                state_payload = await asyncio.to_thread(get_simulation_state_payload, ctx)
                                if state_payload:
                                    try:
                                        await ctx.web_runtime.publish(state_payload, journal=False)
                                    except Exception:
                                        break
                                await asyncio.sleep(1/60.0)
                        ctx.drag_state["task"] = asyncio.create_task(drag_loop())

            elif p_type == "mouse_pick_update":
                targetPos = payload.get("targetPos")
                drag_id = payload.get("drag_id")
                if (
                    drag_id
                    and drag_id == ctx.drag_state.get("drag_id")
                    and _finite_vec3(targetPos)
                ):
                    target_pos = _clamp_manual_target(
                        targetPos,
                        ctx.drag_state.get("target_type", "object"),
                    )
                    async with ctx.executor_lock:
                        def do_update():
                            with ctx.pybullet_lock:
                                if ctx.drag_state.get("constraint_id") is not None:
                                    start_pos = ctx.drag_state.get(
                                        "target_pos",
                                        target_pos,
                                    )
                                    for waypoint in _interpolate_manual_path(
                                        start_pos,
                                        target_pos,
                                        0.01,
                                    ):
                                        pb.changeConstraint(
                                            ctx.drag_state["constraint_id"],
                                            waypoint,
                                            maxForce=ctx.drag_state.get(
                                                "max_force",
                                                100.0,
                                            ),
                                            physicsClientId=ctx.sim.client,
                                        )
                                        for _ in range(2):
                                            pb.stepSimulation(
                                                physicsClientId=ctx.sim.client,
                                            )
                                    ctx.drag_state["target_pos"] = list(target_pos)
                        await asyncio.to_thread(do_update)
                    state_payload = await asyncio.to_thread(
                        get_simulation_state_payload,
                        ctx,
                    )
                    if state_payload:
                        await ctx.web_runtime.publish(state_payload, journal=False)

            elif p_type == "mouse_pick_end":
                drag_id = payload.get("drag_id")
                if drag_id != ctx.drag_state.get("drag_id"):
                    continue
                ctx.drag_state["active"] = False
                async with ctx.executor_lock:
                    def do_end():
                        with ctx.pybullet_lock:
                            if ctx.drag_state.get("constraint_id") is not None:
                                pb.removeConstraint(ctx.drag_state["constraint_id"], physicsClientId=ctx.sim.client)
                                ctx.drag_state["constraint_id"] = None
                    await asyncio.to_thread(do_end)
                ctx.drag_state["drag_id"] = None
                ctx.drag_state["target_type"] = None
                for _ in range(20):
                    with ctx.pybullet_lock:
                        pb.stepSimulation(physicsClientId=ctx.sim.client)
                state_payload = await asyncio.to_thread(
                    get_simulation_state_payload,
                    ctx,
                )
                if state_payload:
                    await ctx.web_runtime.publish(state_payload, journal=False)

            elif p_type == "spawn_catalog_object":
                if ctx.web_runtime.active_token is not None:
                    if not await wait_for_active_task_stop(ctx):
                        continue
                folder = payload.get("folder")
                position = payload.get("position")
                if not _finite_vec3(position):
                    continue
                spawn_position = _clamp_manual_target(position, "object")
                async with ctx.executor_lock:
                    def spawn_object():
                        with ctx.pybullet_lock:
                            return ctx.sim.env_builder.spawn_catalog_object(
                                folder,
                                spawn_position,
                            )
                    try:
                        _, object_name = await asyncio.to_thread(spawn_object)
                    except (ValueError, OSError) as exc:
                        await websocket.send_json({
                            "type": "spawn_rejected",
                            "reason": str(exc),
                        })
                        continue
                for _ in range(30):
                    with ctx.pybullet_lock:
                        pb.stepSimulation(physicsClientId=ctx.sim.client)
                init_payload = await asyncio.to_thread(
                    get_init_payload,
                    ctx,
                )
                ctx.web_runtime.latest_init = init_payload
                state_payload = await asyncio.to_thread(
                    get_simulation_state_payload,
                    ctx,
                )
                await ctx.web_runtime.publish({
                    "type": "object_added",
                    "name": object_name,
                    "config": init_payload["objects"][object_name],
                    "state": (
                        state_payload.get("objects", {}).get(object_name)
                        if state_payload
                        else None
                    ),
                }, journal=False)
                if state_payload:
                    await ctx.web_runtime.publish(state_payload, journal=False)
                await ctx.web_runtime.publish({
                    "type": "log",
                    "message": f"Đã thêm '{object_name}' vào môi trường.",
                })

            elif p_type == "delete_object":
                if ctx.web_runtime.active_token is not None:
                    if not await wait_for_active_task_stop(ctx):
                        continue
                object_name = payload.get("name")
                async with ctx.executor_lock:
                    def delete_object():
                        with ctx.pybullet_lock:
                            object_id = ctx.sim.scene.name_to_id(
                                object_name
                            )
                            if object_id is None:
                                return False
                            if ctx.sim.pipeline._carried == object_id:
                                return False
                            pb.removeBody(
                                object_id,
                                physicsClientId=ctx.sim.client,
                            )
                            ctx.sim.scene.unregister_object(object_id)
                            return True
                    deleted = await asyncio.to_thread(delete_object)
                if not deleted:
                    await websocket.send_json({
                        "type": "delete_rejected",
                        "reason": "Không thể xóa vật đang được giữ hoặc không tồn tại.",
                    })
                    continue
                await ctx.web_runtime.publish({
                    "type": "object_removed",
                    "name": object_name,
                }, journal=False)
                state_payload = await asyncio.to_thread(
                    get_simulation_state_payload,
                    ctx,
                )
                if state_payload:
                    await ctx.web_runtime.publish(state_payload, journal=False)
                await ctx.web_runtime.publish({
                    "type": "log",
                    "message": f"Đã xóa '{object_name}' khỏi môi trường.",
                })

            elif p_type == "reset":
                if not await wait_for_active_task_stop(ctx):
                    await ctx.web_runtime.publish({
                        "type": "log",
                        "message": "Không thể reset khi robot chưa dừng an toàn.",
                    })
                    continue
                async with ctx.executor_lock:
                    def reset_sim():
                        with ctx.pybullet_lock:
                            ctx.sim.reset_objects(seed=0)
                            ctx.sim.controller.reset_to_home()
                    await asyncio.to_thread(reset_sim)
                ctx.web_runtime.conversation_manager.clear()
                ctx.web_runtime.clear_task_history()

                init_payload = await asyncio.to_thread(
                    get_init_payload,
                    ctx,
                )
                ctx.web_runtime.latest_init = init_payload
                await ctx.web_runtime.publish(init_payload, journal=False)

                state_payload = await asyncio.to_thread(get_simulation_state_payload, ctx)
                if state_payload:
                    await ctx.web_runtime.publish(state_payload, journal=False)
                await ctx.web_runtime.publish({"type": "log", "message": "🔄 Đã reset môi trường mô phỏng."})
            
            elif p_type == "change_env":
                if not await wait_for_active_task_stop(ctx):
                    await ctx.web_runtime.publish({
                        "type": "log",
                        "message": "Không thể đổi môi trường khi robot chưa dừng an toàn.",
                    })
                    continue
                env_file = payload.get("env")
                if (
                    not isinstance(env_file, str)
                    or os.path.basename(env_file) != env_file
                ):
                    await websocket.send_json({
                        "type": "environment_rejected",
                        "reason": "Tên environment không hợp lệ.",
                    })
                    continue
                env_path = os.path.join(PROJECT_SRC, "envs", env_file)
                try:
                    _load_environment_config(env_path)
                except (
                    OSError,
                    ValueError,
                    json.JSONDecodeError,
                ) as exc:
                    await websocket.send_json({
                        "type": "environment_rejected",
                        "reason": str(exc),
                    })
                    continue
                async with ctx.executor_lock:
                    def change_sim_env():
                        with ctx.pybullet_lock:
                            ctx.sim.change_environment(env_path)
                    await asyncio.to_thread(change_sim_env)
                ctx.web_runtime.conversation_manager.clear()
                ctx.web_runtime.clear_task_history()
            
                # Send the new static data
                init_payload = await asyncio.to_thread(get_init_payload, ctx)
                ctx.web_runtime.latest_init = init_payload
                await ctx.web_runtime.publish(init_payload, journal=False)
            
                # Send the new dynamic state
                state_payload = await asyncio.to_thread(get_simulation_state_payload, ctx)
                if state_payload:
                    await ctx.web_runtime.publish(state_payload, journal=False)
                await ctx.web_runtime.publish({"type": "log", "message": f"🌍 Đã chuyển sang môi trường {env_file}."})
            
            elif p_type == "update_object_pose":
                if ctx.web_runtime.active_token is not None:
                    continue
                name = payload.get("name")
                pos = payload.get("pos")
                ori = payload.get("ori")
                async with ctx.executor_lock:
                    def move_obj():
                        with ctx.pybullet_lock:
                            oid = ctx.sim.scene.name_to_id(name)
                            if oid is not None:
                                pb.resetBasePositionAndOrientation(oid, pos, ori, physicsClientId=ctx.sim.client)
                    await asyncio.to_thread(move_obj)
            
                state_payload = await asyncio.to_thread(get_simulation_state_payload, ctx)
                if state_payload:
                    await ctx.web_runtime.publish(state_payload, journal=False)
                
            elif p_type == "update_joints":
                if ctx.web_runtime.active_token is not None:
                    continue
                joints_data = payload.get("joints")
                async with ctx.executor_lock:
                    def move_joints():
                        with ctx.pybullet_lock:
                            for idx, q in enumerate(joints_data[:7]):
                                pb.resetJointState(ctx.sim.controller.robot_id, idx, q, physicsClientId=ctx.sim.client)
                            if len(joints_data) >= 11:
                                pb.resetJointState(ctx.sim.controller.robot_id, 9, joints_data[9], physicsClientId=ctx.sim.client)
                                pb.resetJointState(ctx.sim.controller.robot_id, 10, joints_data[10], physicsClientId=ctx.sim.client)
                    await asyncio.to_thread(move_joints)
            
                state_payload = await asyncio.to_thread(get_simulation_state_payload, ctx)
                if state_payload:
                    await ctx.web_runtime.publish(state_payload, journal=False)
            elif p_type in ("camera_rotate", "camera_pan", "camera_zoom"):
                pass
                
    except WebSocketDisconnect:
        print(f"[{session_id}] Client disconnected")
    except Exception as e:
        print(f"[{session_id}] WS Exception: {e}")
    finally:
        if 'ctx' in locals():
            ctx.drag_state["active"] = False
            if ctx.drag_state.get("constraint_id") is not None:
                try:
                    with ctx.pybullet_lock:
                        pb.removeConstraint(
                            ctx.drag_state["constraint_id"],
                            physicsClientId=ctx.sim.client,
                        )
                except Exception:
                    pass
                ctx.drag_state["constraint_id"] = None
            await ctx.web_runtime.unsubscribe(websocket)
            
            # REQUIREMENT: DELETE UPON DISCONNECT (XÓA LUÔN)
            print(f"[{session_id}] Cleaning up session...", flush=True)
            ctx.active_websockets -= 1
            if ctx.active_websockets <= 0 and session_id in active_sessions:
                del active_sessions[session_id]
            if ctx.physics_settle_task:
                ctx.physics_settle_task.cancel()
            if ctx.command_worker_task:
                ctx.command_worker_task.cancel()
            # Removed pb.disconnect to preserve global_sim for next user

if __name__ == "__main__":
    port = int(os.environ.get("PORT", 8000))
    uvicorn.run("interactive_server:app", host="0.0.0.0", port=port, reload=False)