from pathlib import Path
import os
import urllib.request

import numpy as np
import torch
from PIL import Image, ImageDraw, ImageFilter

MODEL_CACHE = None
FLORENCE_MODEL = None
FLORENCE_PROCESSOR = None
LAMA_MODEL = None


def _get_device():
    return torch.device("cuda" if torch.cuda.is_available() else "cpu")


def _get_model_dtype():
    if not torch.cuda.is_available():
        return torch.float32
    try:
        if torch.cuda.is_bf16_supported():
            return torch.bfloat16
    except Exception:
        pass
    return torch.float16


def _get_writable_cache():
    candidates = []
    if os.getenv("WATERMARK_CACHE_DIR"):
        candidates.append(Path(os.environ["WATERMARK_CACHE_DIR"]))
    candidates += [
        Path("/data/.cache/watermark"),
        Path("/tmp/watermark"),
        Path.home() / ".cache" / "watermark",
    ]
    for p in candidates:
        try:
            p.mkdir(parents=True, exist_ok=True)
            test = p / ".write_test"
            test.write_text("ok", encoding="utf-8")
            try:
                test.unlink()
            except Exception:
                pass
            print(f"[Watermark] Cache directory: {p}")
            return p
        except (PermissionError, OSError):
            pass
    raise RuntimeError("No writable watermark cache directory available.")


def _cache_dir():
    global MODEL_CACHE
    if MODEL_CACHE is None:
        MODEL_CACHE = _get_writable_cache()
    return MODEL_CACHE


def _patch_florence_compatibility():
    """Compatibility patches for Florence-2 remote code + newer Transformers."""
    try:
        from transformers import PretrainedConfig, PreTrainedModel

        if not hasattr(PretrainedConfig, "forced_bos_token_id"):
            PretrainedConfig.forced_bos_token_id = None
            print("[Florence] Applied forced_bos_token_id compatibility patch.")

        if not hasattr(PreTrainedModel, "_supports_sdpa"):
            PreTrainedModel._supports_sdpa = False
            print("[Florence] Applied _supports_sdpa compatibility patch.")
    except Exception as e:
        print(f"[Florence] Compatibility patch skipped: {e}")


def _load_models():
    global FLORENCE_MODEL, FLORENCE_PROCESSOR, LAMA_MODEL

    if (
        FLORENCE_MODEL is not None
        and FLORENCE_PROCESSOR is not None
        and LAMA_MODEL is not None
    ):
        return

    _patch_florence_compatibility()

    from transformers import AutoProcessor, AutoModelForCausalLM

    cache = _cache_dir()
    model_id = os.getenv(
        "FLORENCE_MODEL_ID",
        "microsoft/Florence-2-large",
    )
    device = _get_device()
    dtype = _get_model_dtype()

    print(f"[Florence] Loading model: {model_id}")
    print(f"[Florence] device={device}, dtype={dtype}")

    FLORENCE_PROCESSOR = AutoProcessor.from_pretrained(
        model_id,
        cache_dir=str(cache / "hf"),
        trust_remote_code=True,
    )

    model_kwargs = {
        "cache_dir": str(cache / "hf"),
        "trust_remote_code": True,
        "torch_dtype": dtype,
    }

    # Force eager attention so Transformers does not try to negotiate
    # SDPA with older Florence remote-code revisions.
    try:
        model_kwargs["attn_implementation"] = "eager"
    except Exception:
        pass

    FLORENCE_MODEL = AutoModelForCausalLM.from_pretrained(
        model_id,
        **model_kwargs,
    ).to(device).eval()

    # Patch the actual custom model instance as well.
    try:
        if not hasattr(FLORENCE_MODEL, "_supports_sdpa"):
            FLORENCE_MODEL._supports_sdpa = False
    except Exception:
        pass

    try:
        actual_dtype = next(FLORENCE_MODEL.parameters()).dtype
        print(f"[Florence] model dtype={actual_dtype}")
    except Exception:
        pass

    lama_path = cache / "big-lama.pt"
    if not lama_path.exists():
        url = os.getenv(
            "LAMA_URL",
            "https://github.com/Sanster/models/releases/download/add_big_lama/big-lama.pt",
        )
        print("[LaMa] Downloading model...")
        urllib.request.urlretrieve(url, lama_path)

    print("[LaMa] Loading model...")
    LAMA_MODEL = torch.jit.load(
        str(lama_path),
        map_location=device,
    ).eval()

    print("[Watermark] All watermark models loaded.")


def _boxes(image, prompt):
    """Detect watermark bounding boxes using Florence-2."""
    _load_models()

    if image is None:
        return []

    image = image.convert("RGB")
    device = _get_device()

    try:
        model_dtype = next(FLORENCE_MODEL.parameters()).dtype
    except Exception:
        model_dtype = _get_model_dtype()

    task = "<OPEN_VOCABULARY_DETECTION>"
    text = task + (prompt or "watermark")

    inputs = FLORENCE_PROCESSOR(
        text=text,
        images=image,
        return_tensors="pt",
    )

    prepared = {}
    for key, value in inputs.items():
        if hasattr(value, "to"):
            value = value.to(device)
        prepared[key] = value

    # Florence vision tower Conv2D requires input dtype to match model weights.
    if "pixel_values" in prepared:
        prepared["pixel_values"] = prepared["pixel_values"].to(
            device=device,
            dtype=model_dtype,
        )

    if "input_ids" in prepared:
        prepared["input_ids"] = prepared["input_ids"].long()
    if "attention_mask" in prepared:
        prepared["attention_mask"] = prepared["attention_mask"].long()

    print(
        f"[Florence] model_dtype={model_dtype}, "
        f"pixel_dtype={prepared.get('pixel_values').dtype if 'pixel_values' in prepared else None}, "
        f"device={device}"
    )

    generation_kwargs = {
        "input_ids": prepared.get("input_ids"),
        "pixel_values": prepared.get("pixel_values"),
        "max_new_tokens": 1024,
        "num_beams": 3,
        "do_sample": False,
        "use_cache": False,
    }
    generation_kwargs = {
        k: v for k, v in generation_kwargs.items() if v is not None
    }

    with torch.inference_mode():
        generated = FLORENCE_MODEL.generate(**generation_kwargs)

    decoded = FLORENCE_PROCESSOR.batch_decode(
        generated,
        skip_special_tokens=False,
    )[0]

    print("[Florence] decoded:", decoded[:500])

    parsed = FLORENCE_PROCESSOR.post_process_generation(
        decoded,
        task=task,
        image_size=image.size,
    )
    data = parsed.get(task, parsed)

    if not isinstance(data, dict):
        return []

    raw_boxes = data.get("bboxes", []) or []
    valid_boxes = []

    for box in raw_boxes:
        if not isinstance(box, (list, tuple)) or len(box) != 4:
            continue
        try:
            x1, y1, x2, y2 = map(float, box)
        except (TypeError, ValueError):
            continue

        x1 = max(0.0, min(float(image.width), x1))
        y1 = max(0.0, min(float(image.height), y1))
        x2 = max(0.0, min(float(image.width), x2))
        y2 = max(0.0, min(float(image.height), y2))

        if x2 <= x1 or y2 <= y1:
            continue

        valid_boxes.append([x1, y1, x2, y2])

    print(f"[Florence] detected {len(valid_boxes)} bounding boxes")
    return valid_boxes


def _mask(image, boxes, max_pct):
    mask = Image.new("L", image.size, 0)
    draw = ImageDraw.Draw(mask)
    total = image.width * image.height

    for box in boxes or []:
        if len(box) != 4:
            continue

        x1, y1, x2, y2 = map(float, box)
        x1 = max(0, min(image.width, x1))
        x2 = max(0, min(image.width, x2))
        y1 = max(0, min(image.height, y1))
        y2 = max(0, min(image.height, y2))

        area = max(0, x2 - x1) * max(0, y2 - y1)

        if total and area / total * 100 <= max_pct:
            draw.rectangle((x1, y1, x2, y2), fill=255)

    # Small expansion around the detected watermark.
    return mask.filter(ImageFilter.MaxFilter(7))


def _to_lama_tensors(image, mask, device, multiple=8):
    """Convert an image/mask crop to LaMa tensors with safe dimensions."""
    image = image.convert("RGB")
    mask = mask.convert("L")

    w, h = image.size
    pad_w = (multiple - (w % multiple)) % multiple
    pad_h = (multiple - (h % multiple)) % multiple

    image_np = np.asarray(image, dtype=np.float32) / 255.0
    mask_np = np.asarray(mask, dtype=np.float32) / 255.0

    if pad_w or pad_h:
        # Edge padding keeps the context natural and avoids black borders.
        image_np = np.pad(
            image_np,
            ((0, pad_h), (0, pad_w), (0, 0)),
            mode="edge",
        )
        mask_np = np.pad(
            mask_np,
            ((0, pad_h), (0, pad_w)),
            mode="constant",
            constant_values=0,
        )

    image_t = torch.from_numpy(image_np).permute(2, 0, 1).unsqueeze(0).to(device)
    image_t = image_t * 2.0 - 1.0

    mask_t = torch.from_numpy(mask_np).unsqueeze(0).unsqueeze(0).to(device)
    mask_t = (mask_t > 0.5).float()

    return image_t, mask_t, (w, h)


def _run_lama(image, mask):
    """Run LaMa on a crop and return a crop-sized PIL image."""
    device = _get_device()
    image_t, mask_t, original_size = _to_lama_tensors(
        image,
        mask,
        device,
        multiple=8,
    )

    with torch.inference_mode():
        out = LAMA_MODEL(image_t, mask_t)

    if isinstance(out, (tuple, list)):
        out = out[0]

    out = (
        (out.clamp(-1, 1) + 1.0) / 2.0
    ).squeeze(0).permute(1, 2, 0).float().cpu().numpy()

    restored = Image.fromarray(
        (out * 255).clip(0, 255).astype(np.uint8)
    )

    # Remove any safety padding added for LaMa.
    restored = restored.crop(
        (0, 0, original_size[0], original_size[1])
    )
    return restored


def _lama(image, mask, boxes=None):
    """
    Run LaMa only on watermark regions.

    Important V2/V3 fixes:
      1. Pad crop dimensions to a multiple of 8 to avoid LaMa tensor-size
         mismatches such as 245 vs 248.
      2. Composite the LaMa result through the actual mask. The old code pasted
         the entire restored crop back into the frame, which could create a
         bright/white rectangular patch around the watermark.
      3. Keep the original pixels outside the mask untouched.
    """
    _load_models()

    image = image.convert("RGB")
    mask = mask.convert("L")

    if not boxes:
        restored = _run_lama(image, mask)
        blend_mask = mask.filter(ImageFilter.GaussianBlur(1.0))
        return Image.composite(restored, image, blend_mask)

    result = image.copy()
    pad = max(16, int(min(image.size) * 0.02))

    for box in boxes:
        if not isinstance(box, (list, tuple)) or len(box) != 4:
            continue

        x1, y1, x2, y2 = map(float, box)
        left = max(0, int(x1) - pad)
        top = max(0, int(y1) - pad)
        right = min(image.width, int(x2) + pad)
        bottom = min(image.height, int(y2) + pad)

        if right <= left or bottom <= top:
            continue

        crop = result.crop((left, top, right, bottom))
        crop_mask = mask.crop((left, top, right, bottom))

        if not np.asarray(crop_mask).any():
            continue

        restored = _run_lama(crop, crop_mask)

        # Feather only the transition edge. Do NOT paste the entire LaMa crop.
        blend_mask = crop_mask.filter(ImageFilter.GaussianBlur(1.2))
        blended = Image.composite(
            restored,
            crop,
            blend_mask,
        )

        result.paste(
            blended,
            (left, top),
        )

    return result


def preview(image, max_bbox_percent=30, detection_prompt="watermark"):
    if image is None:
        raise ValueError("Please upload an image.")

    boxes = _boxes(image, detection_prompt)
    out = image.convert("RGB").copy()
    draw = ImageDraw.Draw(out)

    for box in boxes:
        if len(box) == 4:
            draw.rectangle(
                tuple(map(float, box)),
                outline="red",
                width=4,
            )

    return out


def remove_image(
    image,
    max_bbox_percent=30,
    detection_prompt="watermark",
    transparent=False,
):
    if image is None:
        raise ValueError("Please upload an image.")

    boxes = _boxes(image, detection_prompt)
    mask = _mask(image, boxes, max_bbox_percent)

    if not boxes:
        return image.convert("RGBA" if transparent else "RGB"), mask

    result = _lama(image, mask, boxes=boxes)

    if transparent:
        result = result.convert("RGBA")
        result.putalpha(mask.point(lambda p: 255 - p))

    return result, mask