import io
import numpy as np
from PIL import Image, ImageColor, ImageOps  # <--- اضافه شدن ImageOps برای Crop/Fill راحت‌تر
from pathlib import Path
from typing import Optional, Tuple
import logging

logger = logging.getLogger(__name__)

BASE_DIR   = Path(__file__).resolve().parent.parent
OUTPUTS_DIR = BASE_DIR / "outputs"
OUTPUTS_DIR.mkdir(exist_ok=True)


# ─── Background Remover ───────────────────────────────────
class BackgroundRemover:

    def __init__(self):
        self.session = None
        self._loaded = False

    def _load_model(self):
        if self._loaded:
            return
        logger.info("Loading rembg model (u2net)...")
        try:
            from rembg import new_session
            self.session = new_session("u2net")
            self._loaded = True
            logger.info("rembg u2net loaded OK")
        except Exception as e:
            logger.error(f"rembg load error: {e}")
            raise RuntimeError(f"خطا در بارگذاری مدل: {e}")

    def remove_background(self, image: Image.Image) -> Image.Image:
        self._load_model()
        try:
            from rembg import remove
            if image.mode not in ("RGB", "RGBA"):
                image = image.convert("RGB")
            result = remove(image, session=self.session)
            return result
        except Exception as e:
            logger.error(f"Background removal error: {e}")
            raise


# ─── Image Processor ──────────────────────────────────────
class ImageProcessor:

    def __init__(self):
        self.bg_remover = BackgroundRemover()

    def process(
        self,
        image_data: bytes,
        project_config: dict,
        template_path: Optional[Path] = None
    ) -> bytes:

        # 1. بارگذاری تصویر
        image = Image.open(io.BytesIO(image_data))
        image = self._fix_orientation(image)

        canvas_w  = project_config["canvas_width"]
        canvas_h  = project_config["canvas_height"]
        remove_bg = project_config["remove_background"]
        size_pct  = project_config["product_size_percent"] / 100.0
        bg_color  = project_config.get("background_color", "#FFFFFF")

        # 2. حذف بکگراند یا پر کردن (Fill)
        if remove_bg:
            logger.info("Removing background...")
            image = self.bg_remover.remove_background(image)
            
            # ساخت canvas رنگی
            canvas = self._create_canvas(canvas_w, canvas_h, bg_color)
            # تغییر سایز متناسب با درصد مشخص شده
            product = self._resize_product(image, canvas_w, canvas_h, size_pct)
            # قرار دادن در وسط بوم
            canvas = self._place_on_canvas(canvas, product, canvas_w, canvas_h)
        else:
            logger.info("Background removal skipped. Filling the canvas...")
            # بدون دوربری — تبدیل به RGB
            if image.mode == "RGBA":
                bg = Image.new("RGB", image.size, "white")
                bg.paste(image, mask=image.split()[3])
                image = bg
            else:
                image = image.convert("RGB")
            
            # پر کردن کل بوم (Fill/Cover) به ابعاد دقیق canvas_w و canvas_h
            canvas = ImageOps.fit(image, (canvas_w, canvas_h), Image.LANCZOS)

        # 3. اعمال template (در صورت وجود) روی تصویر نهایی
        if template_path and template_path.exists():
            canvas = self._apply_template(
                canvas, template_path, canvas_w, canvas_h
            )

        # 4. خروجی WebP
        output = io.BytesIO()
        canvas.convert("RGB").save(output, format="WEBP", quality=85)
        return output.getvalue()

    # ─── Private ──────────────────────────────────────────
    def _fix_orientation(self, image: Image.Image) -> Image.Image:
        try:
            return ImageOps.exif_transpose(image)
        except Exception:
            return image

    def _create_canvas(
        self, width: int, height: int, bg_color: str
    ) -> Image.Image:
        try:
            rgba_color = ImageColor.getrgb(bg_color)
            if len(rgba_color) == 3:
                rgba_color = rgba_color + (255,)
            return Image.new("RGBA", (width, height), rgba_color)
        except Exception as e:
            logger.warning(f"Invalid color '{bg_color}', falling back to white. Error: {e}")
            return Image.new("RGBA", (width, height), (255, 255, 255, 255))

    def _resize_product(
        self,
        image: Image.Image,
        canvas_w: int,
        canvas_h: int,
        size_pct: float
    ) -> Image.Image:
        max_w = int(canvas_w * size_pct)
        max_h = int(canvas_h * size_pct)
        img   = image.copy()
        img.thumbnail((max_w, max_h), Image.LANCZOS)
        return img

    def _place_on_canvas(
        self,
        canvas: Image.Image,
        product: Image.Image,
        canvas_w: int,
        canvas_h: int
    ) -> Image.Image:
        pw, ph = product.size
        x = (canvas_w - pw) // 2
        y = (canvas_h - ph) // 2

        if product.mode == "RGBA":
            canvas.paste(product, (x, y), product.split()[3])
        else:
            canvas.paste(product, (x, y))
        return canvas

    def _apply_template(
        self,
        canvas: Image.Image,
        template_path: Path,
        canvas_w: int,
        canvas_h: int
    ) -> Image.Image:
        template = Image.open(template_path).convert("RGBA")
        if template.size != (canvas_w, canvas_h):
            template = template.resize(
                (canvas_w, canvas_h), Image.LANCZOS
            )
        canvas.paste(template, (0, 0), template.split()[3])
        return canvas


# ─── Output Manager ───────────────────────────────────────
class OutputManager:

    @staticmethod
    def save(
        image_data: bytes,
        project_id: str,
        prefix: str,
        index: int
    ) -> Tuple[Path, str]:
        out_dir = OUTPUTS_DIR / project_id
        out_dir.mkdir(parents=True, exist_ok=True)
        filename = f"{prefix}_{index:03d}.webp"
        out_path = out_dir / filename
        with open(out_path, "wb") as f:
            f.write(image_data)
        return out_path, filename

    @staticmethod
    def create_zip(project_id: str, filenames: list) -> bytes:
        import zipfile
        buf     = io.BytesIO()
        out_dir = OUTPUTS_DIR / project_id
        with zipfile.ZipFile(buf, "w", zipfile.ZIP_DEFLATED) as zf:
            for fn in filenames:
                fp = out_dir / fn
                if fp.exists():
                    zf.write(fp, fn)
        return buf.getvalue()


# Singletons
image_processor = ImageProcessor()
output_manager  = OutputManager()