"""
Local Vision Model — CPU-based image classification for Canifa fashion search.

Uses google/vit-base-patch16-224 (~340MB) for ImageNet-level tag extraction,
then maps labels to Canifa-specific fashion categories via curated heuristics.

Singleton pattern ensures the model loads once and is reused across requests.
"""

import logging
import re
from PIL import Image

logger = logging.getLogger(__name__)


# ─── Fashion Category Mapping ───────────────────────────────────────
# Maps ImageNet labels → Canifa product categories (Vietnamese)

CATEGORY_RULES = [
    # Áo
    (r"\b(t-shirt|tee|jersey)\b",                    "áo thun"),
    (r"\b(polo)\b",                                   "áo polo"),
    (r"\b(shirt|blouse)\b",                           "áo sơ mi"),
    (r"\b(sweater|pullover|sweatshirt|cardigan)\b",   "áo len / sweater"),
    (r"\b(hoodie|hooded)\b",                          "áo hoodie"),
    (r"\b(jacket|blazer|windbreaker)\b",              "áo khoác"),
    (r"\b(coat|overcoat|trench)\b",                   "áo khoác dài"),
    (r"\b(vest|waistcoat)\b",                         "áo gile"),
    (r"\b(tank.?top|sleeveless)\b",                   "áo ba lỗ"),
    # Quần
    (r"\b(jean|denim)\b",                             "quần jeans"),
    (r"\b(trouser|pant|chino|slack)\b",               "quần âu / kaki"),
    (r"\b(short|bermuda)\b",                          "quần short"),
    (r"\b(jogger|sweatpant|track.?pant)\b",           "quần jogger"),
    (r"\b(legging)\b",                                "quần legging"),
    # Váy / Đầm
    (r"\b(dress|gown|sundress)\b",                    "váy / đầm"),
    (r"\b(skirt|miniskirt)\b",                        "chân váy"),
    # Đồ bộ / Đồ ngủ
    (r"\b(pajama|nightgown|robe)\b",                  "đồ ngủ / đồ bộ"),
    (r"\b(suit|tuxedo)\b",                            "bộ vest"),
    # Phụ kiện
    (r"\b(cap|hat|beanie|beret)\b",                   "mũ / nón"),
    (r"\b(scarf|muffler|shawl)\b",                    "khăn"),
    (r"\b(sock|stocking)\b",                          "tất / vớ"),
    (r"\b(bag|backpack|handbag|purse|tote)\b",        "túi / balo"),
    (r"\b(belt|suspender)\b",                         "thắt lưng"),
    (r"\b(shoe|sneaker|boot|sandal|loafer|slipper)\b", "giày dép"),
    (r"\b(sunglasses|glasses|spectacles)\b",          "kính"),
    (r"\b(watch|wristwatch)\b",                       "đồng hồ"),
    (r"\b(tie|bow.?tie|necktie)\b",                   "cà vạt"),
    (r"\b(glove|mitten)\b",                           "găng tay"),
]

COLOR_RULES = [
    (r"\bred\b",           "đỏ"),
    (r"\bblack\b",         "đen"),
    (r"\bwhite\b",         "trắng"),
    (r"\bblue\b",          "xanh dương"),
    (r"\bnavy\b",          "xanh navy"),
    (r"\bgreen\b",         "xanh lá"),
    (r"\byellow\b",        "vàng"),
    (r"\bpink\b",          "hồng"),
    (r"\borange\b",        "cam"),
    (r"\bpurple|violet\b", "tím"),
    (r"\bbrown|tan\b",     "nâu"),
    (r"\bgr[ae]y\b",       "xám"),
    (r"\bbeige|cream\b",   "be / kem"),
    (r"\bkhaki\b",         "kaki"),
    (r"\bmaroon|burgundy\b", "đỏ đậm"),
    (r"\bcoral\b",         "san hô"),
    (r"\bteal|turquoise\b", "xanh ngọc"),
    (r"\bgold\b",          "vàng gold"),
    (r"\bsilver\b",        "bạc"),
    (r"\bindigo\b",        "chàm"),
]

STYLE_RULES = [
    (r"\b(suit|tuxedo|blazer|formal)\b",              "formal"),
    (r"\b(sport|athletic|gym|running|yoga)\b",        "thể thao"),
    (r"\b(casual|everyday|relaxed)\b",                "casual"),
    (r"\b(vintage|retro)\b",                          "vintage"),
    (r"\b(streetwear|street|urban|hip.?hop)\b",       "streetwear"),
    (r"\b(elegant|luxury|premium|silk)\b",            "thanh lịch"),
    (r"\b(outdoor|hiking|camping)\b",                 "outdoor"),
]

# Minimum confidence threshold — below this, tags are considered noise
MIN_CONFIDENCE = 0.05


# ─── Model Class ─────────────────────────────────────────────────────

class LocalVisionModel:
    """
    Ultra-lightweight local Vision Model running strictly on CPU.
    Uses google/vit-base-patch16-224 (rất nhẹ, ~340MB) to extract tags from images.
    If not installed, transformers will auto-download from HuggingFace cache.
    """
    _instance = None
    _classifier = None

    def __new__(cls):
        if cls._instance is None:
            cls._instance = super(LocalVisionModel, cls).__new__(cls)
            cls._instance._init_model()
        return cls._instance

    def _init_model(self):
        try:
            from transformers import pipeline
            logger.info("Khởi tạo Local Vision Model (google/vit-base-patch16-224) trên CPU...")
            # Force CPU with device=-1
            self._classifier = pipeline(
                task="image-classification",
                model="google/vit-base-patch16-224",
                device=-1
            )
            logger.info("Đã khởi tạo thành công Local Vision Model.")
        except ImportError:
            logger.error(
                "Chưa cài đặt thư viện 'transformers' và 'torch'. "
                "Vui lòng chạy: pip install transformers torch torchvision"
            )
            self._classifier = None
        except Exception as e:
            logger.error("Lỗi khởi tạo Vision Model: %s", e)
            self._classifier = None

    def analyze_image(self, image_path: str) -> dict:
        """
        Analyze image at given path and extract fashion-specific features.

        Returns:
            dict with keys: success, features, confidence
            features contains: raw_labels, category, color, style, all_categories
        """
        if not self._classifier:
            return {"error": "Vision model chưa sẵn sàng. Vui lòng kiểm tra logs."}

        try:
            # 1. Load image
            img = Image.open(image_path).convert("RGB")

            # 2. Classify via model
            predictions = self._classifier(img)

            # 3. Filter by confidence threshold
            strong_predictions = [
                p for p in predictions if p["score"] >= MIN_CONFIDENCE
            ]
            if not strong_predictions:
                strong_predictions = predictions[:1]  # Keep at least top-1

            top_tags = [p["label"] for p in strong_predictions[:5]]
            tags_str = " ".join(top_tags).lower()

            # 4. Extract features using curated rules
            category = self._match_first(tags_str, CATEGORY_RULES, "unknown")
            color = self._match_first(tags_str, COLOR_RULES, "unknown")
            style = self._match_first(tags_str, STYLE_RULES, "casual")

            # Collect ALL matching categories (for multi-tag search)
            all_categories = self._match_all(tags_str, CATEGORY_RULES)

            extracted_features = {
                "raw_labels": top_tags,
                "category": category,
                "color": color,
                "style": style,
                "all_categories": all_categories,
            }

            return {
                "success": True,
                "features": extracted_features,
                "confidence": strong_predictions[0]["score"],
            }
        except Exception as e:
            logger.error("Lỗi khi phân tích ảnh: %s", e)
            return {"error": str(e)}

    @staticmethod
    def _match_first(text: str, rules: list, default: str) -> str:
        """Return the first matching value from rules list."""
        for pattern, value in rules:
            if re.search(pattern, text, re.IGNORECASE):
                return value
        return default

    @staticmethod
    def _match_all(text: str, rules: list) -> list[str]:
        """Return all matching values from rules list."""
        matches = []
        for pattern, value in rules:
            if re.search(pattern, text, re.IGNORECASE):
                matches.append(value)
        return matches


# Singleton instance — shared across the application
vision_model = LocalVisionModel()
