"""
Stylist Engine — Pre-computes fashion matching for all products.
Reads fashion_rules.json → scores pairs → writes ai_matches to Postgres.

Usage:
  from worker.stylist_engine import StylistEngine
  engine = StylistEngine()
  result = engine.run_for_code("5TS25S021-SR079")   # 1 product
  engine.run_batch(limit=None)                        # all products
"""

import json
import logging
import os
import time
from typing import Optional

from common.pool_wrapper import get_pooled_connection_compat
from common.starrocks_connection import get_db_connection

logger = logging.getLogger(__name__)

RULES_PATH = os.path.join(os.path.dirname(__file__), "fashion_rules.json")
TABLE = "dashboard_canifa.ultra_descriptions"
SR_TABLE = "shared_source.magento_product_dimension_with_text_embedding"


# ═══ Batch state (used by API for polling) ═══
BATCH_STATE = {
    "is_running": False,
    "total": 0,
    "done": 0,
    "errors": 0,
    "current_code": None,
    "started_at": None,
    "elapsed_s": None,
}


class StylistEngine:
    """Core fashion matching engine."""

    def __init__(self):
        with open(RULES_PATH, "r", encoding="utf-8") as f:
            self.rules = json.load(f)
        self._catalog: list[dict] = []  # cached product catalog

    # ──────────────────────────────────────────
    # Public API
    # ──────────────────────────────────────────

    def run_for_code(self, code: str) -> dict:
        """Compute ai_matches for a single product code and save to DB."""
        catalog = self._get_catalog()
        source = next((p for p in catalog if p["code"] == code), None)
        if not source:
            logger.warning("[Stylist] Product not found in catalog: %s", code)
            return {}

        matches = self._compute_matches(source, catalog)
        self._save_matches(code, matches)
        logger.info("[Stylist] ✅ %s → %d matches across occasions", code, sum(
            len(items) for occ in matches.values() for items in occ.values()
        ))
        return matches

    def run_batch(self, codes: Optional[list[str]] = None) -> int:
        """Batch run for list of codes (or all if None). Updates BATCH_STATE."""
        global BATCH_STATE
        catalog = self._get_catalog(force_reload=True)

        targets = codes or [p["code"] for p in catalog]
        BATCH_STATE.update({
            "is_running": True,
            "total": len(targets),
            "done": 0,
            "errors": 0,
            "current_code": None,
            "started_at": time.time(),
            "elapsed_s": None,
        })

        count = 0
        for code in targets:
            BATCH_STATE["current_code"] = code
            try:
                source = next((p for p in catalog if p["code"] == code), None)
                if not source:
                    BATCH_STATE["errors"] += 1
                    continue
                matches = self._compute_matches(source, catalog)
                self._save_matches(code, matches)
                count += 1
                BATCH_STATE["done"] += 1
            except Exception as e:
                logger.error("[Stylist] Error for %s: %s", code, e)
                BATCH_STATE["errors"] += 1

        elapsed = round(time.time() - BATCH_STATE["started_at"], 1)
        BATCH_STATE.update({
            "is_running": False,
            "elapsed_s": elapsed,
            "current_code": None,
        })
        logger.info("[Stylist] ✅ Batch done: %d/%d in %.1fs", count, len(targets), elapsed)
        return count

    # ──────────────────────────────────────────
    # Core Matching Logic
    # ──────────────────────────────────────────

    def _compute_matches(self, source: dict, catalog: list[dict]) -> dict:
        """Compute matches for one source product across all occasions + roles."""
        rules = self.rules
        occasions = rules["occasions"]
        roles = rules["roles"]
        max_items = rules["role_max_items"]

        # Accumulate scored candidates per occasion+role
        # Structure: {occasion: {role: [scored_item, ...]}}
        buckets: dict[str, dict[str, list[dict]]] = {
            occ: {role: [] for role in roles} for occ in occasions
        }

        for target in catalog:
            if target["code"] == source["code"]:
                continue

            # Hard filter: same gender/age
            if not self._pass_hard_filter(source, target):
                continue

            # Determine role of target product
            role = self._get_role(target)
            if not role:
                continue

            for occ in occasions:
                score = self._score(source, target, occ)
                if score < rules.get("min_score", 35):
                    continue

                reason = self._build_reason(source, target, occ, score)
                buckets[occ][role].append({
                    "code": target["code"],
                    "name": target["name"],
                    "color": target.get("color", ""),
                    "product_line": target.get("product_line", ""),
                    "image": target.get("image", ""),
                    "score": score,
                    "reason": reason,
                })

        # Pick top N per occasion+role, sort by score
        result: dict[str, dict] = {}
        for occ in occasions:
            occ_data: dict[str, list] = {}
            has_any = False
            for role in roles:
                top = sorted(buckets[occ][role], key=lambda x: -x["score"])
                top = self._deduplicate(top, max_items[role])
                if top:
                    occ_data[role] = top
                    has_any = True
            if has_any:
                result[occ] = occ_data

        return result

    # ──────────────────────────────────────────
    # Scoring Sub-functions
    # ──────────────────────────────────────────

    def _pass_hard_filter(self, src: dict, tgt: dict) -> bool:
        """Return False if the pair should be excluded."""
        rules = self.rules

        # Exclude blacklisted product lines
        excluded = [pl.lower() for pl in rules.get("exclude_product_lines", [])]
        if tgt.get("product_line", "").lower() in excluded:
            return False

        # Gender must match (or one of them is unisex)
        sg = (src.get("gender") or "").lower()
        tg = (tgt.get("gender") or "").lower()
        if sg and tg and sg not in ("unisex",) and tg not in ("unisex",) and sg != tg:
            return False

        # Age group must match (kids vs adults)
        sa = (src.get("age_group") or "").lower()
        ta = (tgt.get("age_group") or "").lower()
        if sa and ta:
            src_kid = any(k in sa for k in ["kid", "bé", "boy", "girl", "trẻ em"])
            tgt_kid = any(k in ta for k in ["kid", "bé", "boy", "girl", "trẻ em"])
            if src_kid != tgt_kid:
                return False

        return True

    def _get_role(self, product: dict) -> Optional[str]:
        """Determine which role this product plays in an outfit."""
        pl = product.get("product_line", "")
        return self.rules["product_line_to_role"].get(pl)

    def _score(self, src: dict, tgt: dict, occasion: str) -> int:
        """Compute compatibility score (0-100) for src+tgt pair in given occasion."""
        rules = self.rules
        weights = rules["score_weights"]
        total = 0

        # 1. Color score (0-30) - Adjusted to use new max from weights
        color_score = self._color_score(src.get("color", ""), tgt.get("color", ""))
        total += int(color_score * weights["color"] / 30)

        # 2. Style score (0-25)
        style_score = self._style_score(
            src.get("style_tags", []),
            tgt.get("style_tags", []),
            occasion,
        )
        total += int(style_score * weights["style"] / 25)

        # 3. Occasion overlap score (0-20)
        occ_score = self._occasion_score(
            src.get("occasion_tags", []),
            tgt.get("occasion_tags", []),
            occasion,
        )
        total += int(occ_score * weights["occasion"] / 20)

        # 4. Role bonus — Bottom gets priority over accessories (0-15)
        role = self._get_role(tgt)
        role_priority = rules.get("role_priority", {})
        role_bonus = max(0, 16 - role_priority.get(role, 3) * 5)
        total += int(role_bonus * weights["role"] / 15)

        # 5. Material score (0-10) - New dimension based on season compatibility
        if "material" in weights:
            material_score = self._material_score(src, tgt)
            total += int(material_score * weights["material"] / 10)

        return min(total, 100)

    def _color_score(self, src_color: str, tgt_color: str) -> int:
        """Score color compatibility using group matrix (0-30), fallback to exact key matrix."""
        rules = self.rules
        color_keys = rules["color_keys"]
        
        src_key = self._detect_color_key(src_color, color_keys)
        tgt_key = self._detect_color_key(tgt_color, color_keys)

        # Retrieve groups if available (Neutral, Light, Dark)
        color_groups = rules.get("color_groups", {})
        src_group = color_groups.get(src_key) if src_key else None
        tgt_group = color_groups.get(tgt_key) if tgt_key else None

        # Primary logic: Framework Excel styling (Neutral vs Light vs Dark)
        if src_group and tgt_group:
            group_matrix = rules.get("color_group_matrix", {})
            return group_matrix.get(src_group, {}).get(tgt_group, rules.get("default_color_score", 12))

        # Fallback: specific color matching (matrix legacy)
        if src_key and tgt_key:
            matrix = rules["color_matrix"]
            row = matrix.get(src_key, {})
            return row.get(tgt_key, rules.get("default_color_score", 12))
            
        return rules.get("default_color_score", 12)

    def _detect_color_key(self, color_str: str, color_keys: dict) -> Optional[str]:
        """Map a master_color string to a canonical color key."""
        c = color_str.lower()
        for key, variants in color_keys.items():
            if any(v in c for v in variants):
                return key
        return None

    def _material_score(self, src: dict, tgt: dict) -> int:
        """Score material based on target's materials matching the required season (0-10)."""
        rules = self.rules
        tgt_mat = tgt.get("material_tags", [])
        tgt_season = tgt.get("season_tags", [])
        src_season = src.get("season_tags", [])
        
        # Determine the target season. Combine source season and target season tags for a broader match.
        target_seasons = set(s.lower() for s in tgt_season + src_season)
        
        # If no seasons are provided, default to medium positive compat
        if not target_seasons and not tgt_mat:
            return 5 # Neutral score if missing

        score = 0
        mat_season_matrix = rules.get("material_season_compat", {})
        
        for m in tgt_mat:
            # Map material string to a canonical material key
            matched_key = None
            for key in mat_season_matrix.keys():
                if key in m.lower():
                    matched_key = key
                    break
            
            if matched_key:
                compat_seasons = mat_season_matrix[matched_key]
                # If target season matches compatibility, full points (10/10)
                if any(s in compat_seasons for s in target_seasons):
                    score = max(score, 10)
                else:
                    # Found material but season doesn't ideally match, partial points
                    score = max(score, 4)
            else:
                score = max(score, 5) # Unrecognized material
                
        # If we loop and didn't find any actionable material score, default to 5
        return score if score > 0 else 5

    def _style_score(self, src_styles: list, tgt_styles: list, occasion: str) -> int:
        """Score style compatibility (0-25) with occasion boost."""
        rules = self.rules
        compat = rules["style_compat"]
        boost_styles = rules.get("occasion_style_boost", {}).get(occasion, [])
        best = rules.get("default_style_score", 10)

        for ss in (src_styles or ["Basic"]):
            row = compat.get(ss, {})
            for ts in (tgt_styles or ["Basic"]):
                v = row.get(ts, rules.get("default_style_score", 10))
                if ts in boost_styles:
                    v = min(v + 5, 25)
                best = max(best, v)
        return best

    def _occasion_score(self, src_occ: list, tgt_occ: list, occasion: str) -> int:
        """Score occasion overlap (0-20)."""
        keywords = self.rules.get("occasion_keywords", {}).get(occasion, [])
        src_text = " ".join(src_occ).lower()
        tgt_text = " ".join(tgt_occ).lower()
        both_match = any(kw in src_text for kw in keywords) and any(kw in tgt_text for kw in keywords)
        either_match = any(kw in src_text for kw in keywords) or any(kw in tgt_text for kw in keywords)
        if both_match:
            return 20
        if either_match:
            return 12
        return 6

    def _build_reason(self, src: dict, tgt: dict, occasion: str, score: int) -> str:
        """Generate Vietnamese reason string."""
        parts = []
        sc = self._detect_color_key(src.get("color", ""), self.rules["color_keys"])
        tc = self._detect_color_key(tgt.get("color", ""), self.rules["color_keys"])
        if sc and tc:
            parts.append(f"phối màu {sc.lower()} – {tc.lower()} hài hòa")

        occ_label = self.rules.get("occasion_labels", {}).get(occasion, occasion)
        if occ_label:
            parts.append(f"phù hợp {occ_label.lower()}")

        role = self._get_role(tgt)
        if role:
            role_label = self.rules.get("role_labels", {}).get(role, role)
            parts.append(f"cân đối dáng {role_label.lower()}")

        return ", ".join(parts) if parts else f"Phù hợp phối cùng (điểm {score})"

    def _deduplicate(self, items: list[dict], max_n: int) -> list[dict]:
        """Remove near-duplicates logic (Enforce product diversity + color diversity)."""
        seen_lines: set[str] = set()
        seen_colors: set[str] = set()
        result = []
        for item in items:
            pl = item.get("product_line", item["code"])
            # Canonicalize color slightly for filtering
            color_val = self._detect_color_key(item.get("color", ""), self.rules["color_keys"]) or item.get("color", "")
            
            # Prioritize inserting if we haven't seen this product line OR color yet
            if pl not in seen_lines or color_val not in seen_colors or len(result) < max_n // 2:
                result.append(item)
                seen_lines.add(pl)
                if color_val:
                    seen_colors.add(color_val)
                    
            if len(result) >= max_n:
                break
        return result

    # ──────────────────────────────────────────
    # Data Loading
    # ──────────────────────────────────────────

    def _get_catalog(self, force_reload: bool = False) -> list[dict]:
        """Load product catalog from StarRocks + description_data from Postgres."""
        if self._catalog and not force_reload:
            return self._catalog

        logger.info("[Stylist] Loading catalog...")
        sr_db = get_db_connection()

        # Fetch product basics from StarRocks
        rows = sr_db.execute_query(f"""
            SELECT
                magento_ref_code AS code,
                ANY_VALUE(product_name) AS name,
                ANY_VALUE(master_color) AS color,
                ANY_VALUE(product_line_vn) AS product_line,
                ANY_VALUE(gender_by_product) AS gender,
                ANY_VALUE(product_image_url_thumbnail) AS image
            FROM {SR_TABLE}
            WHERE magento_ref_code IS NOT NULL
              AND product_line_vn IS NOT NULL
            GROUP BY magento_ref_code
        """)

        # Fetch style/occasion tags from Postgres ultra_descriptions
        pg_conn = get_pooled_connection_compat()
        pg_cur = pg_conn.cursor()
        pg_cur.execute(f"""
            SELECT magento_ref_code, description_data
            FROM {TABLE}
            WHERE magento_ref_code IS NOT NULL
        """)
        desc_map: dict[str, dict] = {}
        for r in pg_cur.fetchall():
            code, dd = r
            if isinstance(dd, str):
                try:
                    dd = json.loads(dd)
                except Exception:
                    dd = {}
            if isinstance(dd, dict):
                desc_map[code] = dd
        pg_cur.close()
        pg_conn.close()

        # Build catalog with enriched tags
        catalog = []
        for row in rows:
            code = row["code"]
            dd = desc_map.get(code, {})
            catalog.append({
                "code": code,
                "name": row["name"] or "",
                "color": row["color"] or "",
                "product_line": row["product_line"] or "",
                "gender": row["gender"] or "",
                "age_group": row["gender"] or "",   # re-use gender field for kid/adult
                "image": row["image"] or "",
                "style_tags": self._extract_tags(dd.get("phong_cach", ""), dd.get("tags", "")),
                "occasion_tags": self._extract_occasion(dd.get("dip_mac", ""), dd.get("mua", "")),
                "material_tags": self._extract_tags(dd.get("vat_lieu", "")),
                "season_tags": self._extract_tags(dd.get("mua", "")),
            })

        self._catalog = catalog
        logger.info("[Stylist] Catalog loaded: %d products", len(catalog))
        return catalog

    def _extract_tags(self, *fields: str) -> list[str]:
        """Parse comma/dot/pipe separated style tags."""
        tags = []
        for field in fields:
            if field:
                for sep in [",", "·", "|", ";"]:
                    field = field.replace(sep, ",")
                tags.extend([t.strip() for t in field.split(",") if t.strip()])
        return tags

    def _extract_occasion(self, *fields: str) -> list[str]:
        """Parse occasion strings."""
        return self._extract_tags(*fields)

    # ──────────────────────────────────────────
    # DB Write
    # ──────────────────────────────────────────

    def _save_matches(self, code: str, matches: dict) -> bool:
        """Save ai_matches JSONB to ultra_descriptions."""
        conn = None
        try:
            conn = get_pooled_connection_compat()
            cur = conn.cursor()
            data_json = json.dumps(matches, ensure_ascii=False)
            cur.execute(
                f"""UPDATE {TABLE}
                    SET ai_matches = %s::jsonb, updated_at = NOW()
                    WHERE magento_ref_code = %s OR internal_ref_code = %s""",
                (data_json, code, code),
            )
            if cur.rowcount == 0:
                logger.warning("[Stylist] No row updated for code: %s", code)
            conn.commit()
            cur.close()
            return True
        except Exception as e:
            if conn:
                conn.rollback()
            logger.error("[Stylist] DB save error for %s: %s", code, e)
            return False
        finally:
            if conn:
                conn.close()

    # ── Public aliases for score-test API ───────
    def _load_catalog(self) -> list[dict]:
        """Alias for _get_catalog (used by score-test API)."""
        return self._get_catalog()

    def _score_breakdown(self, src: dict, tgt: dict) -> dict:
        """Return per-criterion score breakdown for a source+target pair."""
        rules = self.rules
        weights = rules["score_weights"]

        # Color
        raw_color = self._color_score(src.get("color", ""), tgt.get("color", ""))
        color_pts = int(raw_color * weights["color"] / 30)

        # Style (use first occasion as context)
        occ = (rules.get("occasions") or ["hang_ngay"])[0]
        raw_style = self._style_score(src.get("style_tags", []), tgt.get("style_tags", []), occ)
        style_pts = int(raw_style * weights["style"] / 25)

        # Occasion
        raw_occ = self._occasion_score(src.get("occasion_tags", []), tgt.get("occasion_tags", []), occ)
        occ_pts = int(raw_occ * weights["occasion"] / 20)

        # Role
        role = self._get_role(tgt)
        role_priority = rules.get("role_priority", {})
        role_bonus = max(0, 16 - role_priority.get(role, 3) * 5)
        role_pts = int(role_bonus * weights["role"] / 15)
        
        # Material
        material_pts = 0
        raw_material = 0
        if "material" in weights:
            raw_material = self._material_score(src, tgt)
            material_pts = int(raw_material * weights["material"] / 10)

        sc = self._detect_color_key(src.get("color", ""), rules["color_keys"])
        tc = self._detect_color_key(tgt.get("color", ""), rules["color_keys"])
        
        cgroups = rules.get("color_groups", {})
        sg = cgroups.get(sc, "?") if sc else "?"
        tg = cgroups.get(tc, "?") if tc else "?"

        return {
            "color":    {"score": color_pts,  "max": weights["color"],    "reason": f"[{sg}] {sc or '?'} → [{tg}] {tc or '?'} (raw={raw_color})"},
            "style":    {"score": style_pts,  "max": weights["style"],    "reason": f"style compat raw={raw_style}"},
            "occasion": {"score": occ_pts,    "max": weights["occasion"], "reason": f"occasion overlap raw={raw_occ}"},
            "role":     {"score": role_pts,   "max": weights["role"],     "reason": f"role={role}, priority={role_priority.get(role,3)}"},
            "material": {"score": material_pts, "max": weights.get("material", 10), "reason": f"material season align raw={raw_material}"},
        }
