#!/usr/bin/env python3
"""Reward server for GRPO garment training.

Runs on one or more GPUs and serves garment simulation reward over HTTP.
The GRPO trainer sends *_specification.json content + sample_id.
This server does:
  1. Warp physics simulation (parallel across configured GPUs)
  2. Find GT mesh from GarmentCodeData
  3. Compute chamfer distance
  4. Return reward scalar (4-tier: parse=H20 / spec=H20 / sim / chamfer)

Usage:
    python sim_scripts/reward_server.py \
        --port 59876 \
        --garmentcode-root /filesdir/GarmentCode \
        --gcd-root /filesdir/GarmentCodeData_v2 \
        --gpus 3 \
        --work-dir /tmp/grpo_reward_work
"""

from __future__ import annotations

import argparse
import json
import logging
import os
import shutil
import subprocess
import sys
import threading
import time
import uuid
from dataclasses import dataclass, field
from functools import lru_cache
from pathlib import Path
from typing import Any, Callable, Optional

import numpy as np
import yaml
try:
    import torch
except Exception:
    torch = None

logger = logging.getLogger(__name__)


# ---------------------------------------------------------------------------
# Reward tiers (must match H20 side grpo_config.py)
# ---------------------------------------------------------------------------

@dataclass
class RewardTiers:
    r_spec: float = 0.2       # spec ok but sim fail
    r_sim_base: float = 0.3   # sim ok, no chamfer available
    chamfer_max: float = 0.05  # normalization cap (meters)


# ---------------------------------------------------------------------------
# Mesh I/O + Chamfer distance
# ---------------------------------------------------------------------------

def _load_obj_vertices(obj_path: Path) -> np.ndarray:
    """Read vertex positions from OBJ file -> (N, 3) float64 array."""
    verts = []
    with open(obj_path, "r") as f:
        for line in f:
            if line.startswith("v "):
                parts = line.strip().split()
                verts.append([float(parts[1]), float(parts[2]), float(parts[3])])
    return np.array(verts, dtype=np.float64)


def _load_ply_vertices(path: Path) -> np.ndarray:
    """Read vertex positions from PLY file (binary or ASCII) -> (N, 3) float64 array."""
    from plyfile import PlyData
    plydata = PlyData.read(str(path))
    x = np.array(plydata["vertex"]["x"], dtype=np.float64)
    y = np.array(plydata["vertex"]["y"], dtype=np.float64)
    z = np.array(plydata["vertex"]["z"], dtype=np.float64)
    return np.stack([x, y, z], axis=1)


def _load_mesh_vertices(path: Path) -> np.ndarray:
    if path.suffix.lower() == ".ply":
        return _load_ply_vertices(path)
    return _load_obj_vertices(path)


@lru_cache(maxsize=4096)
def _load_mesh_vertices_cached(path_str: str) -> np.ndarray:
    """Cache mesh vertices to avoid repeated disk I/O for GT samples."""
    return _load_mesh_vertices(Path(path_str))


def _directed_chamfer(source: np.ndarray, target: np.ndarray, batch_size: int = 4096) -> float:
    n = source.shape[0]
    min_dists = np.empty(n, dtype=np.float64)
    for start in range(0, n, batch_size):
        end = min(start + batch_size, n)
        diff = source[start:end, None, :] - target[None, :, :]
        min_dists[start:end] = (diff * diff).sum(axis=-1).min(axis=1)
    return float(np.sqrt(min_dists).mean())


def _directed_chamfer_torch(
    source: np.ndarray,
    target: np.ndarray,
    device: "torch.device",
    batch_size: int = 4096,
) -> float:
    """Directed chamfer on GPU using torch.cdist."""
    src = torch.from_numpy(source).to(device=device, dtype=torch.float32)
    tgt = torch.from_numpy(target).to(device=device, dtype=torch.float32)
    total = 0.0
    count = 0
    with torch.no_grad():
        for start in range(0, src.shape[0], batch_size):
            chunk = src[start:start + batch_size]
            dists = torch.cdist(chunk, tgt, p=2)
            mins = dists.min(dim=1).values
            total += float(mins.sum().item())
            count += int(mins.numel())
    return total / max(count, 1)


def chamfer_distance(pred: np.ndarray, gt: np.ndarray, gpu_id: Optional[int] = None) -> float:
    """Symmetric chamfer distance between two point clouds.

    Prefer GPU computation when torch+CUDA are available.
    """
    if torch is not None and torch.cuda.is_available():
        try:
            if gpu_id is None:
                device = torch.device("cuda:0")
            else:
                device = torch.device(f"cuda:{gpu_id}")
            return _directed_chamfer_torch(pred, gt, device) + _directed_chamfer_torch(gt, pred, device)
        except Exception as e:
            logger.warning("GPU chamfer failed, fallback to CPU: %s", e)
    return _directed_chamfer(pred, gt) + _directed_chamfer(gt, pred)


def find_gt_sim_mesh(sample_id: str, gcd_root: Path) -> Optional[Path]:
    """Find GT simulated mesh: <gcd_root>/garments_*/default_body/<sample_id>/<sample_id>_sim.{ply,obj}"""
    for garments_dir in sorted(gcd_root.iterdir()):
        if not garments_dir.is_dir():
            continue
        sample_dir = garments_dir / "default_body" / sample_id
        if not sample_dir.is_dir():
            continue
        for suffix in (".ply", ".obj"):
            mesh = sample_dir / f"{sample_id}_sim{suffix}"
            if mesh.exists():
                return mesh
    return None


def build_gt_mesh_index(gcd_root: Path) -> dict[str, Path]:
    index: dict[str, Path] = {}
    for garments_dir in sorted(gcd_root.iterdir()):
        if not garments_dir.is_dir():
            continue
        default_body = garments_dir / "default_body"
        if not default_body.is_dir():
            continue
        for sample_dir in default_body.iterdir():
            if not sample_dir.is_dir():
                continue
            sample_id = sample_dir.name
            for suffix in (".ply", ".obj"):
                mesh = sample_dir / f"{sample_id}_sim{suffix}"
                if mesh.exists():
                    index[sample_id] = mesh
                    break
    return index


# ---------------------------------------------------------------------------
# Simulation
# ---------------------------------------------------------------------------

_FATAL_FAIL_KEYS = (
    "crashes", "frame_timeout", "simulation_timeout",
    "cloth_body_intersection", "cloth_self_intersection", "static_equilibrium",
)


def check_sim_success(sim_props_path: Path) -> bool:
    if not sim_props_path.exists():
        return False
    try:
        props = yaml.safe_load(sim_props_path.read_text())
        fails = props["sim"]["stats"]["fails"]
    except Exception:
        return False
    return not any(fails.get(k, []) for k in _FATAL_FAIL_KEYS)


def run_simulation(
    spec_path: Path,
    garmentcode_root: Path,
    sim_config: str,
    gpu_id: int,
    out_dir: Path,
) -> dict[str, Any]:
    """Run Warp simulation via test_garment_sim.py subprocess."""
    name = "sim_result"
    cmd = [
        sys.executable, "test_garment_sim.py",
        "-p", str(spec_path),
        "-s", sim_config,
        "--out_path", str(out_dir),
        "--out_name", name,
        "--disable_timestamp",
    ]
    env = os.environ.copy()
    env["CUDA_VISIBLE_DEVICES"] = str(gpu_id)

    proc = subprocess.run(
        cmd, cwd=str(garmentcode_root), env=env,
        capture_output=True, text=True, timeout=600,
    )

    rd = out_dir / name
    front = rd / f"{name}_render_front.png"
    back = rd / f"{name}_render_back.png"
    sim_props = rd / "sim_props.yaml"
    sim_mesh = rd / f"{name}_sim.obj"

    render_ok = proc.returncode == 0 and front.exists() and back.exists()
    physics_ok = check_sim_success(sim_props) if sim_props.exists() else False

    return {
        "success": render_ok and physics_ok,
        "sim_mesh": str(sim_mesh) if sim_mesh.exists() else None,
    }


# ---------------------------------------------------------------------------
# Full reward for one spec: simulate + chamfer + score
# ---------------------------------------------------------------------------

def compute_reward_for_spec(
    spec_json_content: str,
    sample_id: str,
    gc_root: Path,
    sim_config: str,
    gpu_id: int,
    work_path: Path,
    gcd_root: Path,
    tiers: RewardTiers,
    keep_work_dir: bool = False,
    gt_mesh_lookup: Optional[Callable[[str], Optional[Path]]] = None,
    worker_pool: Optional[Any] = None,
    sim_timeout_s: int = 30,
) -> dict[str, Any]:
    """Simulate spec, compute chamfer vs GT, return reward."""
    result = {
        "sample_id": sample_id,
        "sim_ok": False,
        "chamfer": None,
        "reward": tiers.r_spec,
        "tier": "sim_fail",
        "work_dir": None,
        "timing_s": {},
    }
    t0 = time.perf_counter()

    job_dir = work_path / f"simjob_{uuid.uuid4().hex[:8]}"
    job_dir.mkdir(parents=True, exist_ok=True)
    spec_path = job_dir / "garment_specification.json"
    spec_path.write_text(spec_json_content, encoding="utf-8")

    try:
        t_sim0 = time.perf_counter()
        if worker_pool is not None:
            sim = worker_pool.submit_and_wait(
                str(spec_path), str(job_dir), timeout=max(5, sim_timeout_s + 5),
            )
        else:
            sim = run_simulation(spec_path, gc_root, sim_config, gpu_id, job_dir)
        result["timing_s"]["sim"] = round(time.perf_counter() - t_sim0, 4)
    except Exception as e:
        logger.warning("Simulation error for %s: %s", sample_id, e)
        if keep_work_dir:
            result["work_dir"] = str(job_dir)
        else:
            shutil.rmtree(str(job_dir), ignore_errors=True)
        result["timing_s"]["total"] = round(time.perf_counter() - t0, 4)
        return result

    if sim.get("error"):
        logger.warning("Worker sim error for %s: %s", sample_id, sim["error"])

    if not sim["success"]:
        if keep_work_dir:
            result["work_dir"] = str(job_dir)
        else:
            shutil.rmtree(str(job_dir), ignore_errors=True)
        result["timing_s"]["total"] = round(time.perf_counter() - t0, 4)
        return result

    result["sim_ok"] = True

    # Chamfer distance vs GT
    chamfer = None
    if sample_id and sim["sim_mesh"]:
        t_lookup0 = time.perf_counter()
        if gt_mesh_lookup is not None:
            gt_mesh = gt_mesh_lookup(sample_id)
        else:
            gt_mesh = find_gt_sim_mesh(sample_id, gcd_root)
        result["timing_s"]["gt_lookup"] = round(time.perf_counter() - t_lookup0, 4)
        if gt_mesh is not None:
            try:
                t_cd0 = time.perf_counter()
                pred_verts = _load_obj_vertices(Path(sim["sim_mesh"]))
                gt_verts = _load_mesh_vertices_cached(str(gt_mesh))
                chamfer_gpu = sim.get("gpu_id", gpu_id) if worker_pool is not None else gpu_id
                chamfer = chamfer_distance(pred_verts, gt_verts, gpu_id=chamfer_gpu)
                result["timing_s"]["chamfer"] = round(time.perf_counter() - t_cd0, 4)
            except Exception as e:
                logger.warning("Chamfer error for %s: %s", sample_id, e)

    if keep_work_dir:
        result["work_dir"] = str(job_dir)
    else:
        shutil.rmtree(str(job_dir), ignore_errors=True)

    if chamfer is not None:
        result["chamfer"] = round(chamfer, 6)
        norm = min(chamfer / tiers.chamfer_max, 1.0)
        result["reward"] = round(tiers.r_sim_base + (1.0 - tiers.r_sim_base) * (1.0 - norm), 4)
    else:
        result["reward"] = tiers.r_sim_base

    result["tier"] = "sim_ok"
    result["timing_s"]["total"] = round(time.perf_counter() - t0, 4)
    return result


# ---------------------------------------------------------------------------
# Job tracking (for async batch endpoint)
# ---------------------------------------------------------------------------

@dataclass
class BatchJob:
    job_id: str
    items: list[dict]
    results: list[Optional[dict]] = field(default_factory=list)
    done: threading.Event = field(default_factory=threading.Event)
    submitted_at: float = 0.0


# ---------------------------------------------------------------------------
# HTTP Server
# ---------------------------------------------------------------------------

def run_server(
    port: int,
    garmentcode_root: str,
    gcd_root: str,
    sim_config: str,
    gpus: list[int],
    work_dir: str,
    chamfer_max: float,
    keep_work_dir: bool,
    use_persistent_workers: bool = True,
    sim_timeout_s: int = 30,
):
    import bottle
    from bottle import request

    gc_root = Path(garmentcode_root).resolve()
    gcd_path = Path(gcd_root)
    work_path = Path(work_dir)
    work_path.mkdir(parents=True, exist_ok=True)
    tiers = RewardTiers(chamfer_max=chamfer_max)

    worker_pool = None
    if use_persistent_workers:
        from sim_worker import SimWorkerPool
        print(f"[RewardServer] Starting persistent worker pool ({len(gpus)} GPUs)...")
        worker_pool = SimWorkerPool(gpus, garmentcode_root, sim_config, sim_timeout_s=sim_timeout_s)
        import atexit
        atexit.register(worker_pool.shutdown)
        print(f"[RewardServer] Worker pool ready.")

    gpu_cycle = [0]
    gpu_locks = {g: threading.Semaphore(1) for g in gpus}
    active_jobs: dict[str, BatchJob] = {}
    jobs_lock = threading.Lock()
    logger.info("[RewardServer] Building GT mesh index from %s ...", gcd_path)
    gt_mesh_index = build_gt_mesh_index(gcd_path)
    logger.info("[RewardServer] Indexed %d GT meshes", len(gt_mesh_index))

    def _get_gt_mesh_cached(sample_id: str) -> Optional[Path]:
        if not sample_id:
            return None
        return gt_mesh_index.get(sample_id)

    def _next_gpu() -> int:
        g = gpus[gpu_cycle[0] % len(gpus)]
        gpu_cycle[0] += 1
        return g

    def _process_one_locked(spec_content: str, sample_id: str, gpu: int) -> dict:
        need_lock = worker_pool is None
        if need_lock:
            gpu_locks[gpu].acquire()
        try:
            return compute_reward_for_spec(
                spec_content, sample_id,
                gc_root, sim_config, gpu, work_path, gcd_path, tiers,
                keep_work_dir=keep_work_dir, gt_mesh_lookup=_get_gt_mesh_cached,
                worker_pool=worker_pool,
                sim_timeout_s=sim_timeout_s,
            )
        except Exception as e:
            return {"sample_id": sample_id, "sim_ok": False, "reward": 0.0,
                    "tier": "error", "error": str(e), "chamfer": None}
        finally:
            if need_lock:
                gpu_locks[gpu].release()

    app = bottle.Bottle()

    @app.route("/health", method="GET")
    def health():
        return json.dumps({"status": "ok", "gpus": gpus, "active_jobs": len(active_jobs)})

    @app.route("/simulate", method="POST")
    def simulate_one():
        """Synchronous: simulate one spec, compute chamfer & reward.

        POST body: {"spec_json": "...", "sample_id": "rand_XXXX"}
        Returns: {"sim_ok": bool, "chamfer": float|null, "reward": float, "tier": str}
        """
        item = json.loads(request.body.read())
        spec = item.get("spec_json", "")
        if not spec:
            return json.dumps({"error": "missing spec_json", "sim_ok": False, "reward": 0.0})
        gpu = _next_gpu()
        return json.dumps(_process_one_locked(spec, item.get("sample_id", ""), gpu))

    @app.route("/simulate_batch", method="POST")
    def simulate_batch():
        """Async: submit batch. POST body: [{"spec_json":"...", "sample_id":"..."}, ...]

        Returns: {"job_id": str}.  Poll GET /result/<job_id>.
        """
        items = json.loads(request.body.read())
        job_id = uuid.uuid4().hex
        job = BatchJob(job_id=job_id, items=items, submitted_at=time.time())
        with jobs_lock:
            active_jobs[job_id] = job

        def _worker():
            from concurrent.futures import ThreadPoolExecutor, as_completed
            results = [None] * len(items)
            with ThreadPoolExecutor(max_workers=len(gpus)) as pool:
                futs = {}
                for i, item in enumerate(items):
                    gpu = _next_gpu()
                    futs[pool.submit(
                        _process_one_locked,
                        item.get("spec_json", ""), item.get("sample_id", ""), gpu,
                    )] = i
                for fut in as_completed(futs):
                    results[futs[fut]] = fut.result()
            job.results = results
            job.done.set()

        threading.Thread(target=_worker, daemon=True).start()
        return json.dumps({"job_id": job_id})

    @app.route("/result/<job_id>", method="GET")
    def get_result(job_id):
        with jobs_lock:
            job = active_jobs.get(job_id)
        if job is None:
            bottle.response.status = 404
            return json.dumps({"error": "unknown job_id"})
        if not job.done.is_set():
            return json.dumps({"status": "pending"})
        elapsed = time.time() - job.submitted_at
        with jobs_lock:
            active_jobs.pop(job_id, None)
        return json.dumps({
            "status": "done",
            "results": job.results,
            "elapsed_s": round(elapsed, 2),
        })

    print(f"[RewardServer] Starting on port {port} with {len(gpus)} GPUs: {gpus}")
    print(f"[RewardServer] GarmentCode:     {gc_root}")
    print(f"[RewardServer] GarmentCodeData: {gcd_path}")
    print(f"[RewardServer] Chamfer max:     {tiers.chamfer_max}")
    print(f"[RewardServer] Keep work dir:   {keep_work_dir}")
    print(f"[RewardServer] Persistent sim:  {use_persistent_workers}")
    print(f"[RewardServer] Sim timeout(s):  {sim_timeout_s}")
    bottle.run(app, host="0.0.0.0", port=port, server="tornado")


# ---------------------------------------------------------------------------
# CLI
# ---------------------------------------------------------------------------

def parse_gpu_list(raw: str) -> list[int]:
    """Parse physical CUDA GPU ids for the simulation server."""
    gpus: list[int] = []
    for item in raw.split(","):
        item = item.strip()
        if not item:
            continue
        try:
            gpus.append(int(item))
        except ValueError as exc:
            raise argparse.ArgumentTypeError(
                f"--gpus must be a comma-separated list of integers, got {raw!r}"
            ) from exc
    if not gpus:
        raise argparse.ArgumentTypeError("--gpus must contain at least one GPU id")
    if len(set(gpus)) != len(gpus):
        raise argparse.ArgumentTypeError(f"--gpus contains duplicates: {raw!r}")
    return gpus

def main():
    parser = argparse.ArgumentParser(description="GRPO Reward Server (sim + chamfer + reward)")
    parser.add_argument("--port", type=int, default=59876)
    parser.add_argument("--garmentcode-root", type=str, default="/filesdir/GarmentCode")
    parser.add_argument("--gcd-root", type=str, default="/filesdir/GarmentCodeData_v2")
    parser.add_argument("--sim-config", type=str, default="assets/Sim_props/default_sim_props.yaml")
    parser.add_argument(
        "--gpus",
        type=parse_gpu_list,
        default=parse_gpu_list(os.environ.get("SIM_GPUS", "3")),
        help=(
            "Physical GPU ids used by the simulation server. Defaults to SIM_GPUS "
            "or GPU 3, so a 4-card machine can train on 0,1,2 and simulate on 3."
        ),
    )
    parser.add_argument("--work-dir", type=str, default="/tmp/grpo_reward_work")
    parser.add_argument("--chamfer-max", type=float, default=0.05)
    parser.add_argument("--sim-timeout", type=int, default=30,
                        help="Per-sample timeout in seconds; long tasks are marked as failed")
    parser.add_argument("--keep-work-dir", action="store_true", help="Keep per-request simulation folders")
    parser.add_argument("--no-persistent-workers", action="store_true",
                        help="Disable persistent sim workers (fall back to subprocess per request)")
    args = parser.parse_args()

    logging.basicConfig(level=logging.INFO)
    run_server(
        port=args.port,
        garmentcode_root=args.garmentcode_root,
        gcd_root=args.gcd_root,
        sim_config=args.sim_config,
        gpus=args.gpus,
        work_dir=args.work_dir,
        chamfer_max=args.chamfer_max,
        keep_work_dir=args.keep_work_dir,
        use_persistent_workers=not args.no_persistent_workers,
        sim_timeout_s=args.sim_timeout,
    )


if __name__ == "__main__":
    main()