"""
Multi-language query builder for handling domain-specific searches.

This module extends the ESQueryBuilder to support multi-language field mappings,
allowing queries to be routed to appropriate language-specific fields while
maintaining a unified external interface.
"""

from typing import Dict, Any, List, Optional
import numpy as np

from config import CustomerConfig, IndexConfig
from query import ParsedQuery
from .es_query_builder import ESQueryBuilder


class MultiLanguageQueryBuilder(ESQueryBuilder):
    """
    Enhanced query builder with multi-language support.

    Handles routing queries to appropriate language-specific fields based on:
    1. Detected query language
    2. Available translations
    3. Domain configuration (language_field_mapping)
    """

    def __init__(
        self,
        config: CustomerConfig,
        index_name: str,
        text_embedding_field: Optional[str] = None,
        image_embedding_field: Optional[str] = None
    ):
        """
        Initialize multi-language query builder.

        Args:
            config: Customer configuration
            index_name: ES index name
            text_embedding_field: Field name for text embeddings
            image_embedding_field: Field name for image embeddings
        """
        self.config = config

        # For default domain, use all fields as fallback
        default_fields = self._get_domain_fields("default")

        super().__init__(
            index_name=index_name,
            match_fields=default_fields,
            text_embedding_field=text_embedding_field,
            image_embedding_field=image_embedding_field
        )

        # Build domain configurations
        self.domain_configs = self._build_domain_configs()

    def _build_domain_configs(self) -> Dict[str, IndexConfig]:
        """Build mapping of domain name to IndexConfig."""
        return {index.name: index for index in self.config.indexes}

    def _get_domain_fields(self, domain_name: str) -> List[str]:
        """Get fields for a specific domain with boost notation."""
        for index in self.config.indexes:
            if index.name == domain_name:
                result = []
                for field_name in index.fields:
                    field = self._get_field_by_name(field_name)
                    if field and field.boost != 1.0:
                        result.append(f"{field_name}^{field.boost}")
                    else:
                        result.append(field_name)
                return result
        return []

    def _get_field_by_name(self, field_name: str):
        """Get field configuration by name."""
        for field in self.config.fields:
            if field.name == field_name:
                return field
        return None

    def build_multilang_query(
        self,
        parsed_query: ParsedQuery,
        query_vector: Optional[np.ndarray] = None,
        filters: Optional[Dict[str, Any]] = None,
        size: int = 10,
        from_: int = 0,
        enable_knn: bool = True,
        knn_k: int = 50,
        knn_num_candidates: int = 200,
        min_score: Optional[float] = None
    ) -> Dict[str, Any]:
        """
        Build ES query with multi-language support.

        Args:
            parsed_query: Parsed query with language info and translations
            query_vector: Query embedding for KNN search
            filters: Additional filters
            size: Number of results
            from_: Offset for pagination
            enable_knn: Whether to use KNN search
            knn_k: K value for KNN
            knn_num_candidates: Number of candidates for KNN
            min_score: Minimum score threshold

        Returns:
            ES query DSL dictionary
        """
        domain = parsed_query.domain
        domain_config = self.domain_configs.get(domain)

        if not domain_config:
            # Fallback to default domain
            domain = "default"
            domain_config = self.domain_configs.get("default")

        if not domain_config:
            # Use original behavior
            return super().build_query(
                query_text=parsed_query.rewritten_query,
                query_vector=query_vector,
                filters=filters,
                size=size,
                from_=from_,
                enable_knn=enable_knn,
                knn_k=knn_k,
                knn_num_candidates=knn_num_candidates,
                min_score=min_score
            )

        print(f"[MultiLangQueryBuilder] Building query for domain: {domain}")
        print(f"[MultiLangQueryBuilder] Detected language: {parsed_query.detected_language}")
        print(f"[MultiLangQueryBuilder] Available translations: {list(parsed_query.translations.keys())}")

        # Build query clause with multi-language support
        query_clause = self._build_multilang_text_query(parsed_query, domain_config)

        es_query = {
            "size": size,
            "from": from_
        }

        # Add filters if provided
        if filters:
            es_query["query"] = {
                "bool": {
                    "must": [query_clause],
                    "filter": self._build_filters(filters)
                }
            }
        else:
            es_query["query"] = query_clause

        # Add KNN search if enabled and vector provided
        if enable_knn and query_vector is not None and self.text_embedding_field:
            knn_clause = {
                "field": self.text_embedding_field,
                "query_vector": query_vector.tolist(),
                "k": knn_k,
                "num_candidates": knn_num_candidates
            }
            es_query["knn"] = knn_clause

        # Add minimum score filter
        if min_score is not None:
            es_query["min_score"] = min_score

        return es_query

    def _build_multilang_text_query(
        self,
        parsed_query: ParsedQuery,
        domain_config: IndexConfig
    ) -> Dict[str, Any]:
        """
        Build text query with multi-language field routing.

        Args:
            parsed_query: Parsed query with language info
            domain_config: Domain configuration

        Returns:
            ES query clause
        """
        if not domain_config.language_field_mapping:
            # No multi-language mapping, use all fields with default analyzer
            fields_with_boost = []
            for field_name in domain_config.fields:
                field = self._get_field_by_name(field_name)
                if field and field.boost != 1.0:
                    fields_with_boost.append(f"{field_name}^{field.boost}")
                else:
                    fields_with_boost.append(field_name)

            return {
                "multi_match": {
                    "query": parsed_query.rewritten_query,
                    "fields": fields_with_boost,
                    "minimum_should_match": "67%",
                    "tie_breaker": 0.9,
                    "boost": domain_config.boost,
                    "_name": f"{domain_config.name}_query"
                }
            }

        # Multi-language mapping exists - build targeted queries
        should_clauses = []
        available_languages = set(domain_config.language_field_mapping.keys())

        # 1. Query in detected language (if it exists in mapping)
        detected_lang = parsed_query.detected_language
        if detected_lang in available_languages:
            target_fields = domain_config.language_field_mapping[detected_lang]
            fields_with_boost = self._apply_field_boosts(target_fields)

            should_clauses.append({
                "multi_match": {
                    "query": parsed_query.rewritten_query,
                    "fields": fields_with_boost,
                    "minimum_should_match": "67%",
                    "tie_breaker": 0.9,
                    "boost": domain_config.boost * 1.5,  # Higher boost for detected language
                    "_name": f"{domain_config.name}_{detected_lang}_query"
                }
            })
            print(f"[MultiLangQueryBuilder] Added query for detected language '{detected_lang}' on fields: {target_fields}")

        # 2. Query in translated languages (only for languages in mapping)
        for lang, translation in parsed_query.translations.items():
            # Only use translations for languages that exist in the mapping
            if lang in available_languages and translation and translation.strip():
                target_fields = domain_config.language_field_mapping[lang]
                fields_with_boost = self._apply_field_boosts(target_fields)

                should_clauses.append({
                    "multi_match": {
                        "query": translation,
                        "fields": fields_with_boost,
                        "minimum_should_match": "67%",
                        "tie_breaker": 0.9,
                        "boost": domain_config.boost,
                        "_name": f"{domain_config.name}_{lang}_translated_query"
                    }
                })
                print(f"[MultiLangQueryBuilder] Added translated query for language '{lang}' on fields: {target_fields}")

        # 3. Fallback: query all fields in mapping if no language-specific query was built
        if not should_clauses:
            print(f"[MultiLangQueryBuilder] No language mapping matched, using all fields from mapping")
            # Use all fields from all languages in the mapping
            all_mapped_fields = []
            for lang_fields in domain_config.language_field_mapping.values():
                all_mapped_fields.extend(lang_fields)
            # Remove duplicates while preserving order
            unique_fields = list(dict.fromkeys(all_mapped_fields))
            fields_with_boost = self._apply_field_boosts(unique_fields)

            should_clauses.append({
                "multi_match": {
                    "query": parsed_query.rewritten_query,
                    "fields": fields_with_boost,
                    "minimum_should_match": "67%",
                    "tie_breaker": 0.9,
                    "boost": domain_config.boost * 0.8,  # Lower boost for fallback
                    "_name": f"{domain_config.name}_fallback_query"
                }
            })

        if len(should_clauses) == 1:
            return should_clauses[0]
        else:
            return {
                "bool": {
                    "should": should_clauses,
                    "minimum_should_match": 1
                }
            }

    def _apply_field_boosts(self, field_names: List[str]) -> List[str]:
        """Apply boost values to field names."""
        result = []
        for field_name in field_names:
            field = self._get_field_by_name(field_name)
            if field and field.boost != 1.0:
                result.append(f"{field_name}^{field.boost}")
            else:
                result.append(field_name)
        return result

    def get_domain_summary(self) -> Dict[str, Any]:
        """Get summary of all configured domains."""
        summary = {}
        for domain_name, domain_config in self.domain_configs.items():
            summary[domain_name] = {
                "label": domain_config.label,
                "fields": domain_config.fields,
                "analyzer": domain_config.analyzer.value,
                "boost": domain_config.boost,
                "has_multilang_mapping": domain_config.language_field_mapping is not None,
                "supported_languages": list(domain_config.language_field_mapping.keys()) if domain_config.language_field_mapping else []
            }
        return summary