"""
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, Set, 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
    intent_texts: Dict[str, str]


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,
        tokenizer_getter: Optional[Callable[[], Any]] = None,
    ) -> None:
        self.registry = registry
        self._text_encoder_getter = text_encoder_getter
        self._tokenizer_getter = tokenizer_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

        query_texts = self._build_query_texts(parsed_query, style_profile)
        query_vector = self._get_query_vector(parsed_query)
        tokenizer = self._get_tokenizer()

        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,
                query_texts=query_texts,
                query_vector=query_vector,
                tokenizer=tokenizer,
            )
            if decision is None:
                continue

            self._apply_decision_to_source(source, decision)
            if decision.rerank_suffix:
                hit["_style_rerank_suffix"] = decision.rerank_suffix

            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

    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 _get_text_encoder(self) -> Any:
        if self._text_encoder_getter is None:
            return None
        return self._text_encoder_getter()

    def _get_tokenizer(self) -> Any:
        if self._tokenizer_getter is None:
            return None
        return self._tokenizer_getter()

    @staticmethod
    def _fallback_sku_text(sku: Dict[str, Any]) -> str:
        parts = []
        for field_name in ("option1_value", "option2_value", "option3_value"):
            value = str(sku.get(field_name) or "").strip()
            if value:
                parts.append(value)
        return " ".join(parts)

    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]:
        candidates: List[_SkuCandidate] = []
        for index, sku in enumerate(skus):
            fallback_text = self._fallback_sku_text(sku)
            intent_texts: Dict[str, str] = {}
            for intent_type, field_name in resolved_dimensions.items():
                if field_name:
                    value = str(sku.get(field_name) or "").strip()
                    intent_texts[intent_type] = value or fallback_text
                else:
                    intent_texts[intent_type] = fallback_text

            selection_parts: List[str] = []
            seen = set()
            for value in intent_texts.values():
                normalized = normalize_query_text(value)
                if not normalized or normalized in seen:
                    continue
                seen.add(normalized)
                selection_parts.append(str(value).strip())

            selection_text = " ".join(selection_parts).strip() or fallback_text
            candidates.append(
                _SkuCandidate(
                    index=index,
                    sku_id=str(sku.get("sku_id") or ""),
                    sku=sku,
                    selection_text=selection_text,
                    intent_texts=intent_texts,
                )
            )
        return candidates

    @staticmethod
    def _is_direct_match(
        candidate: _SkuCandidate,
        query_texts: Sequence[str],
    ) -> bool:
        if not candidate.intent_texts or not query_texts:
            return False
        for value in candidate.intent_texts.values():
            normalized_value = normalize_query_text(value)
            if not normalized_value:
                return False
            if not any(normalized_value in query_text for query_text in query_texts):
                return False
        return True

    def _is_generalized_match(
        self,
        candidate: _SkuCandidate,
        style_profile: StyleIntentProfile,
        tokenizer: Any,
    ) -> bool:
        if not candidate.intent_texts:
            return False

        for intent_type, value in candidate.intent_texts.items():
            definition = self.registry.get_definition(intent_type)
            if definition is None:
                return False
            matched_canonicals = definition.match_text(value, tokenizer=tokenizer)
            if not matched_canonicals.intersection(style_profile.get_canonical_values(intent_type)):
                return False
        return True

    def _select_by_embedding(
        self,
        candidates: Sequence[_SkuCandidate],
        query_vector: Optional[np.ndarray],
    ) -> Tuple[Optional[_SkuCandidate], Optional[float]]:
        if not candidates:
            return None, None
        text_encoder = self._get_text_encoder()
        if query_vector is None or text_encoder is None:
            return candidates[0], None

        unique_texts = list(
            dict.fromkeys(
                normalize_query_text(candidate.selection_text)
                for candidate in candidates
                if normalize_query_text(candidate.selection_text)
            )
        )
        if not unique_texts:
            return candidates[0], None

        vectors = text_encoder.encode(unique_texts, priority=1)
        vector_map: Dict[str, np.ndarray] = {}
        for key, vector in zip(unique_texts, vectors):
            if vector is None:
                continue
            vector_map[key] = np.asarray(vector, dtype=np.float32)

        best_candidate: Optional[_SkuCandidate] = None
        best_score: Optional[float] = None
        query_vector_array = np.asarray(query_vector, dtype=np.float32)
        for candidate in candidates:
            normalized_text = normalize_query_text(candidate.selection_text)
            candidate_vector = vector_map.get(normalized_text)
            if candidate_vector is None:
                continue
            score = float(np.inner(query_vector_array, candidate_vector))
            if best_score is None or score > best_score:
                best_candidate = candidate
                best_score = score

        return best_candidate or candidates[0], best_score

    def _select_for_source(
        self,
        source: Dict[str, Any],
        *,
        style_profile: StyleIntentProfile,
        query_texts: Sequence[str],
        query_vector: Optional[np.ndarray],
        tokenizer: Any,
    ) -> Optional[SkuSelectionDecision]:
        skus = source.get("skus")
        if not isinstance(skus, list) or not skus:
            return None

        resolved_dimensions = self._resolve_dimensions(source, style_profile)
        candidates = self._build_candidates(skus, resolved_dimensions)
        if not candidates:
            return None

        direct_matches = [candidate for candidate in candidates if self._is_direct_match(candidate, query_texts)]
        if len(direct_matches) == 1:
            chosen = direct_matches[0]
            return self._build_decision(chosen, resolved_dimensions, matched_stage="direct")

        generalized_matches: List[_SkuCandidate] = []
        if not direct_matches:
            generalized_matches = [
                candidate
                for candidate in candidates
                if self._is_generalized_match(candidate, style_profile, tokenizer)
            ]
            if len(generalized_matches) == 1:
                chosen = generalized_matches[0]
                return self._build_decision(chosen, resolved_dimensions, matched_stage="generalized")

        embedding_pool = direct_matches or generalized_matches or candidates
        chosen, similarity_score = self._select_by_embedding(embedding_pool, query_vector)
        if chosen is None:
            return None
        stage = "embedding_from_matches" if direct_matches or generalized_matches else "embedding_from_all"
        return self._build_decision(
            chosen,
            resolved_dimensions,
            matched_stage=stage,
            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