# image_captioner/app/captioner.py
from __future__ import annotations
from transformers import pipeline, AutoProcessor
import os
from dataclasses import dataclass
from typing import Optional, Union

from PIL import Image
from transformers import pipeline
from dotenv import load_dotenv

# Load environment variables from .env if present
load_dotenv()

def _get_env(name: str, default: Optional[str] = None) -> str:
    v = os.getenv(name)
    return v if v is not None and v != "" else (default or "")

# ----------------------
# Config (env-driven)
# ----------------------
CAPTION_MODEL        = _get_env("CAPTION_MODEL", "Salesforce/blip-image-captioning-base")
DEVICE               = _get_env("DEVICE", "cpu").lower()        # cpu | mps | cuda
MAX_IMAGE_SIDE       = int(_get_env("MAX_IMAGE_SIDE", "512"))    # px, 0 = no resize
MAX_NEW_TOKENS       = int(_get_env("MAX_NEW_TOKENS", "32"))
NUM_BEAMS            = int(_get_env("NUM_BEAMS", "1"))
REPETITION_PENALTY   = float(_get_env("REPETITION_PENALTY", "1.05"))
PROMPT_PREFIX        = _get_env("PROMPT_PREFIX", "Describe the image clearly and specifically.")

# ----------------------
# Helpers
# ----------------------
def _resize_keep_aspect(img: Image.Image, max_side: int) -> Image.Image:
    if max_side <= 0:
        return img
    w, h = img.size
    m = max(w, h)
    if m <= max_side:
        return img
    if w >= h:
        new_w, new_h = max_side, int(h * (max_side / w))
    else:
        new_h, new_w = max_side, int(w * (max_side / h))
    return img.resize((new_w, new_h), Image.Resampling.LANCZOS)

@dataclass
class CaptionerConfig:
    model_id: str = CAPTION_MODEL
    device: str = DEVICE                # "cpu" | "mps" | "cuda"
    max_image_side: int = MAX_IMAGE_SIDE
    max_new_tokens: int = MAX_NEW_TOKENS
    num_beams: int = NUM_BEAMS
    repetition_penalty: float = REPETITION_PENALTY
    prompt_prefix: str = PROMPT_PREFIX

class ImageCaptioner:
    """
    Lightweight, swappable image captioner.
    - Defaults to a CPU-friendly model.
    - Uses Hugging Face 'image-to-text' pipeline.
    - Device is controlled via env/config; we prefer device_map to avoid accelerate/device conflicts.
    """

    def __init__(self, cfg: CaptionerConfig | None = None):
        self.cfg = cfg or CaptionerConfig()

        # Try to load a FAST processor; fall back to slow if torchvision isn't available
        processor = None
        try:
            processor = AutoProcessor.from_pretrained(self.cfg.model_id, use_fast=True)
        except Exception:
            processor = AutoProcessor.from_pretrained(self.cfg.model_id, use_fast=False)

        tok = getattr(processor, "tokenizer", None)
        # HF versions differ: sometimes it's image_processor, sometimes feature_extractor
        img_proc = getattr(processor, "image_processor", None) or getattr(processor, "feature_extractor", None)

        # On CPU keep it simple; no device_map to avoid accelerate requirements
        self.pipe = pipeline(
            task="image-to-text",
            model=self.cfg.model_id,
            tokenizer=tok,
            **({"image_processor": img_proc} if img_proc is not None else {"feature_extractor": img_proc}),
            device=-1  # CPU
        )

    def generate_caption(
        self,
        image: Union[str, Image.Image],
        prompt_prefix: Optional[str] = None,
        max_new_tokens: Optional[int] = None,
        num_beams: Optional[int] = None,
        repetition_penalty: Optional[float] = None,
    ) -> str:
        """
        Generate a caption for an input image (path or PIL.Image).
        Returns a plain string.
        """
        if isinstance(image, str):
            img = Image.open(image).convert("RGB")
        else:
            img = image.convert("RGB")

        img = _resize_keep_aspect(img, self.cfg.max_image_side)

        prompt = (prompt_prefix if prompt_prefix is not None else self.cfg.prompt_prefix).strip()
        gen_kwargs = {
            "max_new_tokens": int(max_new_tokens or self.cfg.max_new_tokens),
            "num_beams": int(num_beams or self.cfg.num_beams),
            "repetition_penalty": float(repetition_penalty or self.cfg.repetition_penalty),
        }

        # Some pipeline versions accept 'prompt'; others only rely on the model's default behavior.
        # Passing 'prompt' is supported by BLIP family; it’s ignored gracefully by simpler heads.
        try:
            out = self.pipe(img, generate_kwargs=gen_kwargs)
        except TypeError:
            # Fallback if this pipeline signature doesn’t accept 'prompt'
            out = self.pipe(img, generate_kwargs=gen_kwargs)

        # Normalized return: [{"generated_text": "..."}]
        if isinstance(out, list) and out and isinstance(out[0], dict):
            return str(out[0].get("generated_text", "")).strip()
        return str(out).strip()

# Convenience factory (so other modules can import get_captioner())
_captioner_singleton: Optional[ImageCaptioner] = None

def get_captioner() -> ImageCaptioner:
    global _captioner_singleton
    if _captioner_singleton is None:
        _captioner_singleton = ImageCaptioner()
    return _captioner_singleton