#!/usr/bin/env python3
from __future__ import annotations

import json
import math
import os
import sys
from pathlib import Path

import numpy as np
import torch
import torchvision.transforms.functional as TF
from PIL import Image

SCRIPT_DIR = Path(__file__).resolve().parent
BUNDLE_ROOT = SCRIPT_DIR.parent
SAM3_REPO = Path(os.environ.get("SAM3_REPO", BUNDLE_ROOT / "runtime" / "sam3_repo"))
for path in (SCRIPT_DIR, SAM3_REPO):
    if str(path) not in sys.path:
        sys.path.insert(0, str(path))

from sam3_decoder_experiment_lib import SAM3FeatureModel, resize_pad_image_and_mask  # noqa: E402


def bbox_xyxy_to_padded_cxcywh(box, orig_w: int, orig_h: int, image_size: int):
    x1, y1, x2, y2 = [float(v) for v in box]
    scale = float(image_size) / float(max(orig_h, orig_w))
    new_w = orig_w * scale
    new_h = orig_h * scale
    left = (image_size - new_w) / 2.0
    top = (image_size - new_h) / 2.0
    x1p, x2p = x1 * scale + left, x2 * scale + left
    y1p, y2p = y1 * scale + top, y2 * scale + top
    cx = (x1p + x2p) / 2.0 / image_size
    cy = (y1p + y2p) / 2.0 / image_size
    w = max(1.0, x2p - x1p) / image_size
    h = max(1.0, y2p - y1p) / image_size
    return torch.tensor([[[cx, cy, w, h]]], dtype=torch.float32)


def unpad_resize_prediction(prob_512: np.ndarray, orig_w: int, orig_h: int, image_size: int):
    scale = float(image_size) / float(max(orig_h, orig_w))
    new_w = max(1, int(round(orig_w * scale)))
    new_h = max(1, int(round(orig_h * scale)))
    left = (image_size - new_w) // 2
    top = (image_size - new_h) // 2
    crop = prob_512[top : top + new_h, left : left + new_w]
    return np.asarray(Image.fromarray(crop).resize((orig_w, orig_h), Image.BILINEAR))


class SAM3BuscotPredictor:
    def __init__(
        self,
        sam3_checkpoint: str,
        checkpoint_path: str | None = None,
        prompt_type: str = "semantic_text",
        prompt_text: str = "breast tumor",
        image_size: int = 512,
        device: str = "cuda",
        encoder_trainable: str = "frozen",
        lora_rank: int = 8,
        lora_alpha: float = 16,
        threshold: float = 0.5,
    ):
        self.image_size = int(image_size)
        self.threshold = float(threshold)
        self.device = torch.device(device if device == "cuda" and torch.cuda.is_available() else "cpu")
        self.model = SAM3FeatureModel(
            sam3_checkpoint,
            image_size=self.image_size,
            encoder_trainable=encoder_trainable,
            decoder_name="sam3_native",
            prompt_type=prompt_type,
            prompt_text=prompt_text,
            lora_rank=lora_rank,
            lora_alpha=lora_alpha,
        ).to(self.device)
        if checkpoint_path:
            ckpt = torch.load(checkpoint_path, map_location=self.device)
            state = ckpt["model"] if isinstance(ckpt, dict) and "model" in ckpt else ckpt
            cleaned = {}
            for k, v in state.items():
                for prefix in ("module.", "model."):
                    if k.startswith(prefix):
                        k = k[len(prefix) :]
                cleaned[k] = v
            self.model.load_state_dict(cleaned, strict=False)
        self.model.eval()

    @torch.no_grad()
    def predict(self, image_path: str, bbox_xyxy=None):
        image = Image.open(image_path).convert("RGB")
        orig_w, orig_h = image.size
        dummy = np.zeros((orig_h, orig_w), dtype=np.uint8)
        padded, _ = resize_pad_image_and_mask(image, dummy, self.image_size)
        x = TF.to_tensor(padded).unsqueeze(0).to(self.device)
        bbox_prompt = None
        if bbox_xyxy is not None:
            bbox_prompt = bbox_xyxy_to_padded_cxcywh(bbox_xyxy, orig_w, orig_h, self.image_size).to(self.device)
        with torch.amp.autocast(device_type="cuda", dtype=torch.bfloat16, enabled=self.device.type == "cuda"):
            raw = self.model(x, return_details=True, bbox_prompt=bbox_prompt)
        logits = raw["mask_logits"] if isinstance(raw, dict) else raw
        prob_512 = torch.sigmoid(logits)[0, 0].detach().float().cpu().numpy()
        prob = unpad_resize_prediction(prob_512, orig_w, orig_h, self.image_size)
        pred = prob >= self.threshold
        details = {}
        if isinstance(raw, dict):
            for key, value in raw.items():
                if key == "mask_logits":
                    continue
                try:
                    details[key] = float(torch.as_tensor(value).detach().flatten()[0].cpu())
                except Exception:
                    details[key] = str(value)
        return pred.astype(np.uint8), details