import argparse
import sys
import os
import time
from pathlib import Path
import cv2
import numpy as np
import torch

# Add server to path so we can import sam3 locally
server_path = os.path.abspath(os.path.join(os.path.dirname(__file__), "..", "..", "sam31server"))
if server_path not in sys.path:
    sys.path.insert(0, server_path)

SAM3_CHECKPOINT = (
    "C:/Users/nbright/.cache/huggingface/hub/models--facebook--sam3.1/"
    "snapshots/daa63191845a41281374e725f4c9e51c7a824460/sam3.1_multiplex.pt"
)

from sam3.model_builder import build_sam3_image_model
from sam3.model.sam3_image_processor import Sam3Processor

def build_point_grid(n_per_side: int) -> np.ndarray:
    """Generates a 2D grid of points evenly spaced in [0, 1] x [0, 1]."""
    offset = 1.0 / (2 * n_per_side)
    points_one_side = np.linspace(offset, 1 - offset, n_per_side)
    pts_x, pts_y = np.meshgrid(points_one_side, points_one_side)
    grid = np.stack([pts_x.flatten(), pts_y.flatten()], axis=1)
    return grid

def mask_iou(mask1, mask2):
    inter = np.logical_and(mask1, mask2).sum()
    union = np.logical_or(mask1, mask2).sum()
    if union == 0:
        return 0
    return inter / union

def mask_to_polygon(mask, epsilon_factor=0.001):
    mask = np.squeeze(mask)
    mask_uint8 = (mask > 0).astype(np.uint8)
    contours, _ = cv2.findContours(mask_uint8, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
    if not contours:
        return []
    largest = max(contours, key=cv2.contourArea)
    if epsilon_factor > 0:
        epsilon = epsilon_factor * cv2.arcLength(largest, True)
        approx = cv2.approxPolyDP(largest, epsilon, True)
    else:
        approx = largest
    points = [[float(p[0][0]), float(p[0][1])] for p in approx]
    return points

def segment_everything(image_bgr, model_path, points_per_side=32, conf_thresh=0.8, nms_thresh=0.5):
    print("Loading SAM3 Model locally...")
    device = "cuda" if torch.cuda.is_available() else "cpu"
    bpe_path = os.path.join(server_path, "bpe_simple_vocab_16e6.txt.gz")
    
    model = build_sam3_image_model(
        bpe_path=bpe_path,
        device=device,
        checkpoint_path=model_path,
    )
    
    if device == "cuda":
        torch.backends.cuda.matmul.allow_tf32 = True
        torch.backends.cudnn.allow_tf32 = True

    processor = Sam3Processor(model, confidence_threshold=conf_thresh, device=device)
    
    # PIL image format
    image_rgb = cv2.cvtColor(image_bgr, cv2.COLOR_BGR2RGB)
    from PIL import Image
    pil_image = Image.fromarray(image_rgb)
    
    print("Computing image embedding...")
    t0 = time.time()
    state = processor.set_image(pil_image)
    print(f"Image embedding done in {time.time() - t0:.2f}s")
    
    grid_points = build_point_grid(points_per_side)
    print(f"Generated {len(grid_points)} grid points for sampling.")
    
    masks = []
    scores = []
    
    t0 = time.time()
    for i, (nx, ny) in enumerate(grid_points):
        if i % 100 == 0:
            print(f"  Processed {i}/{len(grid_points)} points...")
            
        processor.reset_all_prompts(state)
        state = processor.add_point_prompt(point=[nx, ny], label=True, state=state)
        
        if "masks" in state and len(state["masks"]) > 0:
            # Take the mask with the highest score
            best_idx = torch.argmax(state["scores"])
            mask = state["masks"][best_idx].cpu().numpy()
            score = state["scores"][best_idx].item()
            if score > conf_thresh:
                masks.append(mask)
                scores.append(score)
    print(f"Grid prediction done in {time.time() - t0:.2f}s")
    print(f"Found {len(masks)} raw masks.")
    
    if not masks:
        return []
    
    # Simple NMS based on IoU
    print("Applying NMS...")
    order = np.argsort(scores)[::-1]
    keep = []
    
    for idx in order:
        if len(keep) == 0:
            keep.append(idx)
            continue
            
        current_mask = masks[idx]
        overlap = False
        for k in keep:
            iou = mask_iou(current_mask, masks[k])
            if iou > nms_thresh:
                overlap = True
                break
        if not overlap:
            keep.append(idx)
            
    final_masks = [masks[idx] for idx in keep]
    final_scores = [scores[idx] for idx in keep]
    print(f"Kept {len(final_masks)} masks after NMS.")
    
    # Convert to polygons
    results = []
    for m, s in zip(final_masks, final_scores):
        poly = mask_to_polygon(m)
        if poly:
            results.append({"polygon": poly, "score": s})
            
    return results

def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("--input", required=True, help="Input image path")
    parser.add_argument("--output", required=True, help="Output vis image path")
    parser.add_argument("--points", type=int, default=32, help="Points per side")
    parser.add_argument("--conf", type=float, default=0.7, help="Confidence threshold")
    parser.add_argument("--nms", type=float, default=0.7, help="NMS IoU threshold")
    parser.add_argument("--split", action="store_true",
                        help="겹치지 않는 마스크끼리 묶어 여러 장으로 저장")
    args = parser.parse_args()
    
    buf = np.fromfile(args.input, dtype=np.uint8)
    image = cv2.imdecode(buf, cv2.IMREAD_COLOR)
    if image is None:
        print(f"Failed to read {args.input}")
        return
        
    # Shrink image if too large just to make NMS faster
    h, w = image.shape[:2]
    max_dim = 1024
    if max(h, w) > max_dim:
        scale = max_dim / max(h, w)
        image_proc = cv2.resize(image, (int(w * scale), int(h * scale)))
    else:
        image_proc = image.copy()
        
    model_path = SAM3_CHECKPOINT
    
    results = segment_everything(image_proc, model_path, points_per_side=args.points,
                                 conf_thresh=args.conf, nms_thresh=args.nms)
    
    np.random.seed(42)
    colors = [np.random.randint(0, 255, (3,)).tolist() for _ in results]
    H, W = image_proc.shape[:2]
    regions = []
    for res in results:
        region = np.zeros((H, W), dtype=np.uint8)
        cv2.fillPoly(region, [np.array(res["polygon"], dtype=np.int32)], 255)
        regions.append(region)

    if args.split:
        # 겹치지 않는 것끼리 묶어 레이어 분리 (점수 높은 순 greedy first-fit)
        layers = []          # [(누적마스크, [인덱스...])]
        for i, region in enumerate(regions):
            for canvas, members in layers:
                if not np.any(cv2.bitwise_and(canvas, region)):
                    cv2.bitwise_or(canvas, region, canvas)
                    members.append(i)
                    break
            else:
                layers.append((region.copy(), [i]))
        print(f"Split into {len(layers)} non-overlapping layers.")
        groups = [members for _, members in layers]
    else:
        groups = [list(range(len(results)))]

    out_path = Path(args.output)
    for n, members in enumerate(groups, 1):
        vis = image_proc.copy()
        for i in members:
            pts = np.array(results[i]["polygon"], dtype=np.int32)
            color = colors[i]

            overlay = vis.copy()
            cv2.fillPoly(overlay, [pts], color)
            cv2.addWeighted(overlay, 0.4, vis, 0.6, 0, vis)

            # 테두리: 선을 그린 뒤 마스크 내부만 남겨 안쪽 1px로 만듦
            edge = np.zeros((H, W), dtype=np.uint8)
            cv2.polylines(edge, [pts], True, 255, 2, cv2.LINE_8)
            vis[cv2.bitwise_and(edge, regions[i]) > 0] = [int(c * 0.35) for c in color]

        path = (out_path if len(groups) == 1
                else out_path.with_name(f"{out_path.stem}_L{n}{out_path.suffix}"))
        is_success, im_buf_arr = cv2.imencode(out_path.suffix, vis)
        if is_success:
            im_buf_arr.tofile(str(path))
        print(f"Saved {len(members)} masks to {path}")

    try:                                    # 완료 알림음
        import winsound
        winsound.Beep(880, 150)
        winsound.Beep(1175, 250)
    except Exception:
        print("\a", end="")
    
if __name__ == "__main__":
    main()