"""
  SKU selection for style-intent-aware search results.
  """
  
  from __future__ import annotations
  
  from dataclasses import dataclass, field

  from typing import Any, Callable, Dict, List, Optional, Sequence, Tuple

  
  import numpy as np
  
  from query.style_intent import StyleIntentProfile, StyleIntentRegistry
  from query.tokenization import normalize_query_text
  
  
  @dataclass(frozen=True)
  class SkuSelectionDecision:
      selected_sku_id: Optional[str]
      rerank_suffix: str
      selected_text: str
      matched_stage: str
      similarity_score: Optional[float] = None
      resolved_dimensions: Dict[str, Optional[str]] = field(default_factory=dict)
  
      def to_dict(self) -> Dict[str, Any]:
          return {
              "selected_sku_id": self.selected_sku_id,
              "rerank_suffix": self.rerank_suffix,
              "selected_text": self.selected_text,
              "matched_stage": self.matched_stage,
              "similarity_score": self.similarity_score,
              "resolved_dimensions": dict(self.resolved_dimensions),
          }
  
  
  @dataclass
  class _SkuCandidate:
      index: int
      sku_id: str
      sku: Dict[str, Any]
      selection_text: str

      normalized_selection_text: str
      intent_values: Dict[str, str]

      normalized_intent_values: Dict[str, str]

  
  
  @dataclass
  class _SelectionContext:
      query_texts: Tuple[str, ...]
      matched_terms_by_intent: Dict[str, Tuple[str, ...]]
      query_vector: Optional[np.ndarray]
      text_match_cache: Dict[Tuple[str, str], bool] = field(default_factory=dict)
      selection_vector_cache: Dict[str, Optional[np.ndarray]] = field(default_factory=dict)
      similarity_cache: Dict[str, Optional[float]] = field(default_factory=dict)

  
  
  class StyleSkuSelector:
      """Selects the best SKU for an SPU based on detected style intent."""
  
      def __init__(
          self,
          registry: StyleIntentRegistry,
          *,
          text_encoder_getter: Optional[Callable[[], Any]] = None,

      ) -> None:
          self.registry = registry
          self._text_encoder_getter = text_encoder_getter

  
      def prepare_hits(
          self,
          es_hits: List[Dict[str, Any]],
          parsed_query: Any,
      ) -> Dict[str, SkuSelectionDecision]:
          decisions: Dict[str, SkuSelectionDecision] = {}
          style_profile = getattr(parsed_query, "style_intent_profile", None)
          if not isinstance(style_profile, StyleIntentProfile) or not style_profile.is_active:
              return decisions
  

          selection_context = self._build_selection_context(parsed_query, style_profile)

  
          for hit in es_hits:
              source = hit.get("_source")
              if not isinstance(source, dict):
                  continue
  
              decision = self._select_for_source(
                  source,
                  style_profile=style_profile,

                  selection_context=selection_context,

              )
              if decision is None:
                  continue
  

              if decision.rerank_suffix:
                  hit["_style_rerank_suffix"] = decision.rerank_suffix

              else:
                  hit.pop("_style_rerank_suffix", None)

  
              doc_id = hit.get("_id")
              if doc_id is not None:
                  decisions[str(doc_id)] = decision
  
          return decisions
  
      def apply_precomputed_decisions(
          self,
          es_hits: List[Dict[str, Any]],
          decisions: Dict[str, SkuSelectionDecision],
      ) -> None:
          if not es_hits or not decisions:
              return
  
          for hit in es_hits:
              doc_id = hit.get("_id")
              if doc_id is None:
                  continue
              decision = decisions.get(str(doc_id))
              if decision is None:
                  continue
              source = hit.get("_source")
              if not isinstance(source, dict):
                  continue
              self._apply_decision_to_source(source, decision)
              if decision.rerank_suffix:
                  hit["_style_rerank_suffix"] = decision.rerank_suffix

              else:
                  hit.pop("_style_rerank_suffix", None)

  
      def _build_query_texts(
          self,
          parsed_query: Any,
          style_profile: StyleIntentProfile,
      ) -> List[str]:
          texts = [variant.normalized_text for variant in style_profile.query_variants if variant.normalized_text]
          if texts:
              return list(dict.fromkeys(texts))
  
          fallbacks: List[str] = []
          for value in (
              getattr(parsed_query, "original_query", None),
              getattr(parsed_query, "query_normalized", None),
              getattr(parsed_query, "rewritten_query", None),
          ):
              normalized = normalize_query_text(value)
              if normalized:
                  fallbacks.append(normalized)
          translations = getattr(parsed_query, "translations", {}) or {}
          if isinstance(translations, dict):
              for value in translations.values():
                  normalized = normalize_query_text(value)
                  if normalized:
                      fallbacks.append(normalized)
          return list(dict.fromkeys(fallbacks))
  
      def _get_query_vector(self, parsed_query: Any) -> Optional[np.ndarray]:
          query_vector = getattr(parsed_query, "query_vector", None)
          if query_vector is not None:
              return np.asarray(query_vector, dtype=np.float32)
  
          text_encoder = self._get_text_encoder()
          if text_encoder is None:
              return None
  
          query_text = (
              getattr(parsed_query, "rewritten_query", None)
              or getattr(parsed_query, "query_normalized", None)
              or getattr(parsed_query, "original_query", None)
          )
          if not query_text:
              return None
  
          vectors = text_encoder.encode([query_text], priority=1)
          if vectors is None or len(vectors) == 0 or vectors[0] is None:
              return None
          return np.asarray(vectors[0], dtype=np.float32)
  

      def _build_selection_context(
          self,
          parsed_query: Any,
          style_profile: StyleIntentProfile,
      ) -> _SelectionContext:
          matched_terms_by_intent: Dict[str, List[str]] = {}
          for intent in style_profile.intents:
              normalized_term = normalize_query_text(intent.matched_term)
              if not normalized_term:
                  continue
              matched_terms = matched_terms_by_intent.setdefault(intent.intent_type, [])
              if normalized_term not in matched_terms:
                  matched_terms.append(normalized_term)
  
          return _SelectionContext(
              query_texts=tuple(self._build_query_texts(parsed_query, style_profile)),
              matched_terms_by_intent={
                  intent_type: tuple(terms)
                  for intent_type, terms in matched_terms_by_intent.items()
              },
              query_vector=self._get_query_vector(parsed_query),
          )
  

      def _get_text_encoder(self) -> Any:
          if self._text_encoder_getter is None:
              return None
          return self._text_encoder_getter()
  

      def _resolve_dimensions(
          self,
          source: Dict[str, Any],
          style_profile: StyleIntentProfile,
      ) -> Dict[str, Optional[str]]:
          option_names = {
              "option1_value": normalize_query_text(source.get("option1_name")),
              "option2_value": normalize_query_text(source.get("option2_name")),
              "option3_value": normalize_query_text(source.get("option3_name")),
          }
          resolved: Dict[str, Optional[str]] = {}
          for intent in style_profile.intents:
              if intent.intent_type in resolved:
                  continue
              aliases = set(intent.dimension_aliases or self.registry.get_dimension_aliases(intent.intent_type))
              matched_field = None
              for field_name, option_name in option_names.items():
                  if option_name and option_name in aliases:
                      matched_field = field_name
                      break
              resolved[intent.intent_type] = matched_field
          return resolved
  
      def _build_candidates(
          self,
          skus: List[Dict[str, Any]],
          resolved_dimensions: Dict[str, Optional[str]],
      ) -> List[_SkuCandidate]:

          if not resolved_dimensions or any(not field_name for field_name in resolved_dimensions.values()):
              return []
  

          candidates: List[_SkuCandidate] = []
          for index, sku in enumerate(skus):

              intent_values: Dict[str, str] = {}

              normalized_intent_values: Dict[str, str] = {}

              for intent_type, field_name in resolved_dimensions.items():

                  if not field_name:
                      continue

                  raw = str(sku.get(field_name) or "").strip()
                  intent_values[intent_type] = raw
                  normalized_intent_values[intent_type] = normalize_query_text(raw)

  
              selection_parts: List[str] = []

              norm_parts: List[str] = []
              seen: set[str] = set()
              for intent_type, raw in intent_values.items():
                  nv = normalized_intent_values[intent_type]
                  if not nv or nv in seen:

                      continue

                  seen.add(nv)
                  selection_parts.append(raw)
                  norm_parts.append(nv)

              selection_text = " ".join(selection_parts).strip()

              normalized_selection_text = " ".join(norm_parts).strip()

              candidates.append(
                  _SkuCandidate(
                      index=index,
                      sku_id=str(sku.get("sku_id") or ""),
                      sku=sku,
                      selection_text=selection_text,

                      normalized_selection_text=normalized_selection_text,

                      intent_values=intent_values,

                      normalized_intent_values=normalized_intent_values,

                  )
              )
          return candidates
  
      @staticmethod

      def _empty_decision(
          resolved_dimensions: Dict[str, Optional[str]],
          matched_stage: str,
      ) -> SkuSelectionDecision:
          return SkuSelectionDecision(
              selected_sku_id=None,
              rerank_suffix="",
              selected_text="",
              matched_stage=matched_stage,
              resolved_dimensions=dict(resolved_dimensions),
          )
  
      def _is_text_match(

          self,

          intent_type: str,
          value: str,
          selection_context: _SelectionContext,

          *,
          normalized_value: Optional[str] = None,

      ) -> bool:

          if normalized_value is None:
              normalized_value = normalize_query_text(value)

          if not normalized_value:

              return False
  

          cache_key = (intent_type, normalized_value)
          cached = selection_context.text_match_cache.get(cache_key)
          if cached is not None:
              return cached
  
          matched_terms = selection_context.matched_terms_by_intent.get(intent_type, ())
          has_term_match = any(term in normalized_value for term in matched_terms if term)
          query_contains_value = any(
              normalized_value in query_text
              for query_text in selection_context.query_texts
          )
          matched = bool(has_term_match or query_contains_value)
          selection_context.text_match_cache[cache_key] = matched
          return matched
  
      def _find_first_text_match(
          self,
          candidates: Sequence[_SkuCandidate],
          selection_context: _SelectionContext,
      ) -> Optional[_SkuCandidate]:
          for candidate in candidates:
              if candidate.intent_values and all(

                  self._is_text_match(
                      intent_type,
                      value,
                      selection_context,
                      normalized_value=candidate.normalized_intent_values[intent_type],
                  )

                  for intent_type, value in candidate.intent_values.items()
              ):
                  return candidate
          return None

  
      def _select_by_embedding(
          self,
          candidates: Sequence[_SkuCandidate],

          selection_context: _SelectionContext,

      ) -> Tuple[Optional[_SkuCandidate], Optional[float]]:
          if not candidates:
              return None, None
          text_encoder = self._get_text_encoder()

          if selection_context.query_vector is None or text_encoder is None:
              return None, None

  
          unique_texts = list(
              dict.fromkeys(

                  candidate.normalized_selection_text

                  for candidate in candidates

                  if candidate.normalized_selection_text
                  and candidate.normalized_selection_text not in selection_context.selection_vector_cache

              )
          )

          if unique_texts:
              vectors = text_encoder.encode(unique_texts, priority=1)
              for key, vector in zip(unique_texts, vectors):
                  selection_context.selection_vector_cache[key] = (
                      np.asarray(vector, dtype=np.float32) if vector is not None else None
                  )

  
          best_candidate: Optional[_SkuCandidate] = None
          best_score: Optional[float] = None

          query_vector_array = np.asarray(selection_context.query_vector, dtype=np.float32)

          for candidate in candidates:

              normalized_text = candidate.normalized_selection_text
              if not normalized_text:
                  continue
  
              score = selection_context.similarity_cache.get(normalized_text)
              if score is None:
                  candidate_vector = selection_context.selection_vector_cache.get(normalized_text)
                  if candidate_vector is None:
                      selection_context.similarity_cache[normalized_text] = None
                      continue
                  score = float(np.inner(query_vector_array, candidate_vector))
                  selection_context.similarity_cache[normalized_text] = score
  
              if score is None:

                  continue

              if best_score is None or score > best_score:
                  best_candidate = candidate
                  best_score = score
  

          return best_candidate, best_score

  
      def _select_for_source(
          self,
          source: Dict[str, Any],
          *,
          style_profile: StyleIntentProfile,

          selection_context: _SelectionContext,

      ) -> Optional[SkuSelectionDecision]:
          skus = source.get("skus")
          if not isinstance(skus, list) or not skus:
              return None
  
          resolved_dimensions = self._resolve_dimensions(source, style_profile)

          if not resolved_dimensions or any(not field_name for field_name in resolved_dimensions.values()):
              return self._empty_decision(resolved_dimensions, matched_stage="unresolved")
  

          candidates = self._build_candidates(skus, resolved_dimensions)
          if not candidates:

              return self._empty_decision(resolved_dimensions, matched_stage="no_candidates")

          text_match = self._find_first_text_match(candidates, selection_context)
          if text_match is not None:
              return self._build_decision(text_match, resolved_dimensions, matched_stage="text")

          chosen, similarity_score = self._select_by_embedding(candidates, selection_context)

          if chosen is None:

              return self._empty_decision(resolved_dimensions, matched_stage="no_match")

          return self._build_decision(
              chosen,
              resolved_dimensions,

              matched_stage="embedding",

              similarity_score=similarity_score,
          )
  
      @staticmethod
      def _build_decision(
          candidate: _SkuCandidate,
          resolved_dimensions: Dict[str, Optional[str]],
          *,
          matched_stage: str,
          similarity_score: Optional[float] = None,
      ) -> SkuSelectionDecision:
          return SkuSelectionDecision(
              selected_sku_id=candidate.sku_id or None,
              rerank_suffix=str(candidate.selection_text or "").strip(),
              selected_text=str(candidate.selection_text or "").strip(),
              matched_stage=matched_stage,
              similarity_score=similarity_score,
              resolved_dimensions=dict(resolved_dimensions),
          )
  
      @staticmethod
      def _apply_decision_to_source(source: Dict[str, Any], decision: SkuSelectionDecision) -> None:
          skus = source.get("skus")
          if not isinstance(skus, list) or not skus or not decision.selected_sku_id:
              return
  
          selected_index = None
          for index, sku in enumerate(skus):
              if str(sku.get("sku_id") or "") == decision.selected_sku_id:
                  selected_index = index
                  break
          if selected_index is None:
              return
  
          selected_sku = skus.pop(selected_index)
          skus.insert(0, selected_sku)
  
          image_src = selected_sku.get("image_src") or selected_sku.get("imageSrc")
          if image_src:
              source["image_url"] = image_src