import sys
import os
import numpy as np
import torch
from pathlib import Path
from PIL import Image


class Sen2SRInference:
    """
    Inference wrapper for WEO-SAS/sen2sr — a Sentinel-2 super-resolution model
    hosted on Hugging Face (https://huggingface.co/WEO-SAS/sen2sr).

    The model's `main` branch expects 4-channel RGBN input (B04, B03, B02, B08)
    as a float32 NumPy array in [0, 1] with shape (C, H, W) and returns a
    4× upscaled output of the same channel order.

    This wrapper accepts standard 3-channel RGB PIL images and adapts them
    by treating the R channel as a proxy for NIR (B08), since single-image
    workflows won't generally have a separate NIR band.
    """

    MODEL_REPO = "WEO-SAS/sen2sr"

    def __init__(self, device=None, local_dir=None):

        if device is None:
            device = "cuda" if torch.cuda.is_available() else "cpu"

        self.device = str(device)
        self.local_dir = local_dir
        self.model = None

        self._load_model()

    # ------------------------------------------------------------------
    # Model loading
    # ------------------------------------------------------------------

    def _load_model(self):
        """Download and load the sen2sr model from Hugging Face Hub."""

        try:
            from huggingface_hub import snapshot_download
        except ImportError:
            raise ImportError(
                "huggingface_hub is required for Sen2SR. "
                "Install it with: pip install huggingface-hub"
            )

        try:
            from safetensors import safe_open  # noqa: F401 — ensure available
        except ImportError:
            raise ImportError(
                "safetensors is required for Sen2SR. "
                "Install it with: pip install safetensors"
            )

        print("Sen2SR: downloading / verifying model from Hugging Face Hub …")

        if self.local_dir is None:
            self.local_dir = snapshot_download(
                repo_id=self.MODEL_REPO,
                revision="main",
            )

        print(f"Sen2SR: model files at {self.local_dir}")

        # The model repo exposes a `model.py` with a `Model` class.
        if str(self.local_dir) not in sys.path:
            sys.path.insert(0, str(self.local_dir))

        from model import Model  # type: ignore  # loaded dynamically

        self.model = Model(local_dir=str(self.local_dir))

        print("Sen2SR: model loaded successfully.")

    # ------------------------------------------------------------------
    # Inference
    # ------------------------------------------------------------------

    @torch.no_grad()
    def predict(self, image) -> Image.Image:
        """
        Run 4× super-resolution on a PIL RGB image.

        Parameters
        ----------
        image : PIL.Image.Image | np.ndarray
            Input satellite image (RGB, any uint8 size).

        Returns
        -------
        PIL.Image.Image
            4× upscaled RGB image.
        """

        if self.model is None:
            raise RuntimeError("Sen2SR model is not loaded.")

        # ── Normalise input ────────────────────────────────────────────
        if not isinstance(image, Image.Image):
            image = Image.fromarray(image)

        image = image.convert("RGB")

        img_np = np.array(image).astype(np.float32) / 255.0  # (H, W, 3)

        # ── Build RGBN array: (4, H, W) ───────────────────────────────
        # Model expects: B04(R), B03(G), B02(B), B08(NIR).
        # Approximate NIR with the Red channel as the best available proxy.
        r = img_np[:, :, 0]
        g = img_np[:, :, 1]
        b = img_np[:, :, 2]
        nir_proxy = r

        rgbn = np.stack([r, g, b, nir_proxy], axis=0)  # (4, H, W)

        # ── Pad input to a multiple of patch_size=128 ─────────────────
        # The HardConstraint FFT low-pass mask is baked at patch_size=128
        # (output size 512). Both the direct path (image ≤ 128) and the
        # tiled predict_large path need H and W to be exact multiples of
        # 128 so every tile lands on a 128×128 boundary.
        PATCH = 128
        _, orig_h, orig_w = rgbn.shape

        pad_h = (PATCH - orig_h % PATCH) % PATCH
        pad_w = (PATCH - orig_w % PATCH) % PATCH

        if pad_h > 0 or pad_w > 0:
            rgbn = np.pad(
                rgbn,
                ((0, 0), (0, pad_h), (0, pad_w)),
                mode="reflect",
            )

        # ── Run model ─────────────────────────────────────────────────
        sr_rgbn = self.model.predict(rgbn)  # (4, H_padded*4, W_padded*4)

        # ── Crop back to the true 4× output size ──────────────────────
        out_h = orig_h * 4
        out_w = orig_w * 4
        sr_rgbn = sr_rgbn[:, :out_h, :out_w]

        # ── Convert to uint8 RGB PIL image ────────────────────────────
        sr_rgb = sr_rgbn[:3, :, :]           # (3, H*4, W*4)
        sr_rgb = np.clip(sr_rgb, 0.0, 1.0)
        sr_rgb = (sr_rgb * 255.0).round().astype(np.uint8)
        sr_rgb = sr_rgb.transpose(1, 2, 0)   # (H*4, W*4, 3)

        return Image.fromarray(sr_rgb, mode="RGB")


# -----------------------------------------------------------------------
# Convenience function
# -----------------------------------------------------------------------

def run_sen2sr(image, device=None, local_dir=None):
    model = Sen2SRInference(device=device, local_dir=local_dir)
    return model.predict(image)