from __future__ import annotations

import os
from pathlib import Path
from typing import Optional

import numpy as np
from PIL import Image


DEFAULT_LAMA_URL = os.environ.get(
    "LAMA_MODEL_URL",
    "https://github.com/Sanster/models/releases/download/add_big_lama/big-lama.pt",
)


def _default_cache_path() -> Path:
    # 放在 python_server/outputs/cache 下，避免污染用户家目录
    here = Path(__file__).resolve()
    cache_dir = here.parents[2] / "outputs" / "cache"
    cache_dir.mkdir(parents=True, exist_ok=True)
    return cache_dir / "big-lama.pt"


def _download(url: str, dst: Path) -> None:
    import requests

    dst.parent.mkdir(parents=True, exist_ok=True)
    resp = requests.get(url, stream=True, timeout=120)
    resp.raise_for_status()
    tmp = dst.with_suffix(dst.suffix + ".tmp")
    with tmp.open("wb") as f:
        for chunk in resp.iter_content(chunk_size=1024 * 1024):
            if not chunk:
                continue
            f.write(chunk)
    tmp.replace(dst)


def _pad_to_mod8(img: np.ndarray) -> tuple[np.ndarray, tuple[int, int]]:
    h, w = img.shape[:2]
    pad_h = (8 - (h % 8)) % 8
    pad_w = (8 - (w % 8)) % 8
    if pad_h == 0 and pad_w == 0:
        return img, (0, 0)
    out = np.pad(img, ((0, pad_h), (0, pad_w), (0, 0)), mode="reflect")
    return out, (pad_h, pad_w)


def _unpad(img: np.ndarray, pad_hw: tuple[int, int]) -> np.ndarray:
    pad_h, pad_w = pad_hw
    if pad_h == 0 and pad_w == 0:
        return img
    h, w = img.shape[:2]
    return img[: h - pad_h, : w - pad_w]


class LaMaTorchscript:
    """
    轻量 LaMa（big-lama.pt torchscript）推理封装。
    - 输入：RGB 图 + L mask（白=需要修补）
    - 输出：RGB 图（同尺寸）
    """

    def __init__(self, device: str = "cpu", model_path: Optional[str] = None, model_url: Optional[str] = None):
        import torch

        self.device = device
        self.model_url = model_url or DEFAULT_LAMA_URL
        p = Path(model_path) if model_path else _default_cache_path()
        if not p.exists():
            _download(self.model_url, p)
        self.model_path = p
        self.model = torch.jit.load(str(self.model_path), map_location=device).eval()

    @staticmethod
    def _norm_img_u8(rgb: np.ndarray) -> np.ndarray:
        # [H,W,3] uint8 -> float32 [0,1]
        return (rgb.astype(np.float32) / 255.0).clip(0.0, 1.0)

    def __call__(self, image: Image.Image, mask: Image.Image) -> Image.Image:
        import torch

        img = np.asarray(image.convert("RGB"), dtype=np.uint8)
        m = np.asarray(mask.convert("L"), dtype=np.uint8)
        m = (m >= 128).astype(np.float32)  # 1=hole

        img_f = self._norm_img_u8(img)
        # pad to multiple of 8
        img_pad, pad_hw = _pad_to_mod8(img_f)
        m_pad, _ = _pad_to_mod8(np.repeat(m[:, :, None], 3, axis=2))
        m1 = m_pad[:, :, 0]  # [H,W]

        # torch: [1,C,H,W]
        it = torch.from_numpy(img_pad).permute(2, 0, 1).unsqueeze(0).to(self.device)
        mt = torch.from_numpy(m1).unsqueeze(0).unsqueeze(0).to(self.device)

        with torch.no_grad():
            out = self.model(it, mt)

        out_np = out[0].permute(1, 2, 0).detach().cpu().numpy()
        out_np = np.clip(out_np * 255.0, 0.0, 255.0).astype(np.uint8)
        out_np = _unpad(out_np, pad_hw)
        return Image.fromarray(out_np, mode="RGB")