# Copyright (c) Microsoft Corporation.
# Licensed under the MIT License.

from collections.abc import Callable
import copy
from dataclasses import dataclass, replace
import datetime
from typing import Literal, TypeGuard, cast

import typechat

from ..knowpro.collections import PropertyTermSet
from ..knowpro.convutils import get_time_range_prompt_section_for_conversation
from ..knowpro.interfaces import (
    DateRange,
    Datetime,
    IConversation,
    KnowledgePropertyName,
    KnowledgeType,
    PropertySearchTerm,
    SearchSelectExpr,
    SearchTerm,
    SearchTermGroup,
    SearchTermGroupTypes,
    Term,
    WhenFilter,
)
from ..storage.memory.propindex import PropertyNames
from ..knowpro.search import (
    ConversationSearchResult,
    SearchOptions,
    SearchQueryExpr,
    has_conversation_result,
    has_conversation_results,
    run_search_query,
)
from ..knowpro.searchlib import create_property_search_term

from .date_time_schema import DateTime, DateTimeRange
from .search_query_schema import (
    ActionTerm,
    EntityTerm,
    SearchExpr,
    SearchFilter,
    SearchQuery,
    VerbsTerm,
)

# APIs for searching with Natural Language.
# Work in progress; frequent improvements/tweaks.


type SearchQueryTranslator = typechat.TypeChatJsonTranslator[SearchQuery]


@dataclass
class LanguageSearchFilter:
    """Type representing the filter options for language search."""

    knowledgeType: KnowledgeType | None = None
    threadDescription: str | None = None
    tags: list[str] | None = None


@dataclass
class LanguageQueryExpr:
    query_text: str  # The text of the query.
    query: SearchQuery  # The structured search query the queryText was translated to.
    # The search query expressions the structured query was compiled to:
    query_expressions: list[SearchQueryExpr]


@dataclass
class LanguageQueryCompileOptions:
    exact_scope: bool = False  # Is fuzzy matching enabled when applying scope?
    verb_scope: bool = True
    term_filter: Callable[[str], bool] | None = None  # To ignore noise terms
    # Debug flags:
    apply_scope: bool = True  # False to turn off scope matching entirely


@dataclass
class LanguageSearchOptions(SearchOptions):
    compile_options: LanguageQueryCompileOptions | None = None
    fallback_rag_options: None = None  # Don't need LanguageSearchRagOptions yet
    model_instructions: list[typechat.PromptSection] | None = None

    def __repr__(self):
        parts = []
        for key in dir(self):
            if not key.startswith("_"):
                value = getattr(self, key)
                if value is not None:
                    parts.append(f"{key}={value!r}")
        return f"{self.__class__.__name__}({', '.join(parts)})"


@dataclass
class LanguageSearchDebugContext:
    # Query returned by the LLM:
    search_query: SearchQuery | None = None
    # What search_query was compiled into:
    search_query_expr: list[SearchQueryExpr] | None = None
    # For each expr in searchQueryExpr, returns if a raw text similarity match was used:
    # TODO: used_similarity_fallback: list[bool] | None = None

    # Values to override the search query
    use_search_query: SearchQuery | None = None
    use_compiled_search_query_exprs: list[SearchQueryExpr] | None = None


# NOTE: Arguments 2 and 3 are reversed compared to the TypeScript version
# for consistency with other similar functions in this file.
async def search_conversation_with_language(
    # TODO: Add comments to the parameters (copy from @param in TS code).
    conversation: IConversation,
    query_translator: SearchQueryTranslator,
    search_text: str,
    options: LanguageSearchOptions | None = None,
    lang_search_filter: LanguageSearchFilter | None = None,
    debug_context: LanguageSearchDebugContext | None = None,
) -> typechat.Result[list[ConversationSearchResult]]:
    options = options or LanguageSearchOptions()
    if debug_context and debug_context.use_compiled_search_query_exprs:
        search_query = debug_context.use_search_query
        search_query_exprs = debug_context.use_compiled_search_query_exprs
    else:
        lang_query_result = await search_query_expr_from_language(
            conversation,
            query_translator,
            search_text,
            options,
            lang_search_filter,
            debug_context,
        )
        if not isinstance(lang_query_result, typechat.Success):
            return lang_query_result
        search_query = lang_query_result.value.query
        search_query_exprs = lang_query_result.value.query_expressions

    if debug_context:
        debug_context.search_query_expr = search_query_exprs
        # TODO: debug_context.used_similarity_fallback = [False] * len(search_query_exprs)

    fallback_query_exprs: list[SearchQueryExpr] | None = None
    if search_query:
        fallback_query_exprs = _compile_fallback_query(
            conversation,
            search_query,
            options.compile_options or LanguageQueryCompileOptions(),
            lang_search_filter,
        )

    search_results: list[ConversationSearchResult] = []
    for i, search_query_expr in enumerate(search_query_exprs):
        fallback_query = fallback_query_exprs[i] if fallback_query_exprs else None
        query_result = await run_search_query(conversation, search_query_expr, options)
        if fallback_query and not has_conversation_results(query_result):
            # Rerun the query but with verb matching turned off for scopes.
            query_result = await run_search_query(
                conversation,
                fallback_query,
                options,
            )
        # TODO: If no matches and fallback enabled... run the raw query. (RAG fallback)
        search_results.extend(query_result)

    return typechat.Success(search_results)


def _compile_fallback_query(
    conversation: IConversation,
    query: SearchQuery,
    compile_options: LanguageQueryCompileOptions,
    lang_search_filter: LanguageSearchFilter | None = None,
) -> list[SearchQueryExpr] | None:
    """
    Scoping queries can be precise. However, there may be random variations in how LLMs
    translate some user utterances into queries... particularly verbs.
    The verbs may not match action verbs actually in the index...
    related terms may not meet the similarity cutoff.
    If configured (compile_options.exact_scope == false),
    we can do a fallback query that does not enforce verb matching.
    This improves recall while still providing a reasonable level of scoping because it
    will still match the action subject and object.
    """
    # If no exact scope... and verbScope is not provided or true,
    # then we can build a fallback query that is more forgiving.
    if compile_options.verb_scope and not compile_options.exact_scope:
        return compile_search_query(
            conversation,
            query,
            replace(compile_options, verb_scope=False),
            lang_search_filter,
        )

    # No fallback query currently possible.
    return None


async def search_query_expr_from_language(
    # TODO: Add comments to the parameters (copy from @param in TS code).
    conversation: IConversation,
    translator: SearchQueryTranslator,
    query_text: str,
    options: LanguageSearchOptions | None = None,
    lang_search_filter: LanguageSearchFilter | None = None,
    debug_context: LanguageSearchDebugContext | None = None,
) -> typechat.Result[LanguageQueryExpr]:
    options = options or LanguageSearchOptions()
    if debug_context and debug_context.use_search_query:
        # If the debug context has a use_search query, use it instead of translating.
        query = debug_context.use_search_query
    else:
        query_result = await search_query_from_language(
            conversation,
            translator,
            query_text,
            options.model_instructions,
        )
        if not isinstance(query_result, typechat.Success):
            return query_result
        query = query_result.value
    if debug_context:
        debug_context.search_query = query
    query_expressions = compile_search_query(
        conversation,
        query,
        options.compile_options,
        lang_search_filter,
    )
    return typechat.Success(
        LanguageQueryExpr(
            query_text,
            query,
            query_expressions,
        )
    )


def compile_search_query(
    conversation: IConversation,
    query: SearchQuery,
    options: LanguageQueryCompileOptions | None = None,
    lang_search_filter: LanguageSearchFilter | None = None,
) -> list[SearchQueryExpr]:
    compiler = SearchQueryCompiler(
        conversation,
        options or LanguageQueryCompileOptions(),
        lang_search_filter,
    )
    return compiler.compile_query(query)


def compile_search_filter(
    conversation: IConversation,
    search_filter: SearchFilter,
    options: LanguageQueryCompileOptions | None = None,
    lang_search_filter: LanguageSearchFilter | None = None,
) -> SearchSelectExpr:
    compiler = SearchQueryCompiler(
        conversation,
        options or LanguageQueryCompileOptions(),
        lang_search_filter,
    )
    return compiler.compile_search_filter(search_filter)


class SearchQueryCompiler:

    def __init__(
        self,
        conversation: IConversation,
        options: LanguageQueryCompileOptions | None = None,
        lang_search_filter: LanguageSearchFilter | None = None,
    ):
        self.conversation = conversation
        self.options = options = options or LanguageQueryCompileOptions()
        self.lang_search_filter = lang_search_filter or LanguageSearchFilter()
        self.exact_scope = options.exact_scope
        self.verb_scope = options.verb_scope
        self.term_filter = options.term_filter
        self.apply_scope = options.apply_scope

        self.entity_terms_added = PropertyTermSet()
        self.dedupe = True

    def compile_query(self, query: SearchQuery) -> list[SearchQueryExpr]:
        query = copy.copy(query)  # Shallow copy so we can modify it
        query_expressions: list[SearchQueryExpr] = []
        for search_expr in query.search_expressions:
            query_expressions.append(self.compile_search_expr(search_expr))
        return query_expressions

    # Every searchExpr has one or more filters.
    # Each filter is compiled into a selectExpr.
    def compile_search_expr(self, search_expr: SearchExpr) -> SearchQueryExpr:
        query_expr = SearchQueryExpr(select_expressions=[])
        if search_expr.filters:
            for filter in search_expr.filters:
                query_expr.select_expressions.append(self.compile_search_filter(filter))
        query_expr.raw_query = search_expr.rewritten_query
        return query_expr

    def compile_search_filter(self, filter: SearchFilter) -> SearchSelectExpr:
        search_term_group = self.compile_term_group(filter)
        when = self.compile_when(filter)
        return SearchSelectExpr(
            search_term_group,
            when,
        )

    def compile_term_group(self, filter: SearchFilter) -> SearchTermGroup:
        term_group = SearchTermGroup(boolean_op="or", terms=[])
        self.entity_terms_added.clear()
        terms = filter.entity_search_terms
        if is_entity_term_list(terms):
            self.compile_entity_terms(terms, term_group)
        if filter.action_search_term:
            # term_group.terms.append("filter.actionSearchTerm, false, true")
            self.compile_action_term_as_search_terms(
                filter.action_search_term, term_group, False
            )
        if filter.search_terms is not None:
            self.compile_search_terms(filter.search_terms, term_group)
        elif len(term_group.terms) == 0:
            # Summary
            term_group.terms.append(create_property_search_term("topic", "*"))
        return term_group

    def compile_when(self, filter: SearchFilter) -> WhenFilter | None:
        when: WhenFilter | None = None
        action_term = filter.action_search_term
        if (
            self.apply_scope
            and action_term is not None
            and self.should_add_scope(action_term)
        ):
            scope_defining_terms = self.compile_scope(
                action_term,
                include_additional_entities=False,
                include_verbs=self.verb_scope if self.verb_scope is not None else True,
            )
            if scope_defining_terms.terms:
                if when is None:
                    when = WhenFilter()
                when.scope_defining_terms = scope_defining_terms
        if filter.time_range is not None:
            if when is None:
                when = WhenFilter()
            when.date_range = date_range_from_datetime_range(filter.time_range)
        return when

    def compile_action_term_as_search_terms(
        self,
        action_term: ActionTerm,
        term_group: SearchTermGroup | None = None,
        use_or_max: bool = True,
    ) -> SearchTermGroup:
        if term_group is None:
            term_group = SearchTermGroup("or")
        action_group = SearchTermGroup("or_max") if use_or_max else term_group
        if action_term.action_verbs is not None:
            for verb in action_term.action_verbs.words:
                self.add_property_term_to_group("topic", verb, action_group)
        if is_entity_term_list(action_term.actor_entities):
            self.compile_entity_terms_as_search_terms(
                action_term.actor_entities, action_group
            )
        if is_entity_term_list(action_term.target_entities):
            self.compile_entity_terms_as_search_terms(
                action_term.target_entities, action_group
            )
        if is_entity_term_list(action_term.additional_entities):
            self.compile_entity_terms_as_search_terms(
                action_term.additional_entities, action_group
            )
        return term_group

    def compile_search_terms(
        self, search_terms: list[str], term_group: SearchTermGroup | None = None
    ) -> SearchTermGroup:
        if term_group is None:
            term_group = SearchTermGroup(boolean_op="or", terms=[])
        for search_term in search_terms:
            term_group.terms.append(SearchTerm(Term(search_term)))
        return term_group

    def compile_entity_terms(
        self,
        entity_terms: list[EntityTerm],
        term_group: SearchTermGroup,
        use_or_max: bool = True,
    ) -> None:
        if use_or_max:
            save_dedupe = self.dedupe
            self.dedupe = False
            for term in entity_terms:
                or_max = SearchTermGroup(
                    boolean_op="or_max",
                    terms=[],
                )
                self.add_entity_term_to_group(term, or_max)
                term_group.terms.append(optimize_or_max(or_max))
            self.dedupe = save_dedupe
        else:
            for term in entity_terms:
                self.add_entity_term_to_group(term, term_group)
        # Also search for topics.
        for term in entity_terms:
            self.add_entity_name_to_group(term, PropertyNames.Topic, term_group)
            if term.facets is not None:
                for facet in term.facets:
                    if facet.facet_value not in (None, "*"):
                        self.add_property_term_to_group(
                            facet.facet_value, "topic", term_group
                        )

    def compile_entity_terms_as_search_terms(
        self,
        entity_terms: list[EntityTerm],
        term_group: SearchTermGroup,
    ) -> None:
        for term in entity_terms:
            self.add_entity_term_as_search_terms_to_group(term, term_group)

    def compile_scope(
        self,
        action_term: ActionTerm,
        include_additional_entities: bool = True,
        include_verbs: bool = True,
    ) -> SearchTermGroup:
        save_dedupe = self.dedupe
        self.dedupe = False

        term_group = self.compile_action_term(action_term, True, include_verbs)
        if include_additional_entities and is_entity_term_list(
            action_term.additional_entities
        ):
            self.add_entity_names_to_group(
                action_term.additional_entities,
                PropertyNames.EntityName,
                term_group,
                self.exact_scope,
            )

        self.dedupe = save_dedupe
        return term_group

    def compile_action_term(
        self,
        action_term: ActionTerm,
        use_and: bool,
        include_verbs: bool,
    ) -> SearchTermGroup:
        save_dedupe = self.dedupe
        self.dedupe = False

        term_group: SearchTermGroup
        if is_entity_term_list(action_term.target_entities):
            term_group = SearchTermGroup("and" if use_and else "or")
            for entity in action_term.target_entities:
                # S.V.O. == Subject, Verb, Object
                svo_term_group = (
                    self.compile_subject_and_verb(action_term)
                    if include_verbs
                    else self.compile_subject(action_term)
                )
                # A target can be the name of an object of an action OR the name of an entity.
                object_term_group = self.compile_object(entity)
                if object_term_group.terms:
                    svo_term_group.terms.append(object_term_group)
                term_group.terms.append(svo_term_group)
            if len(term_group.terms) == 1:
                term_group = cast(SearchTermGroup, term_group.terms[0])
        else:
            term_group = self.compile_subject_and_verb(action_term)

        self.dedupe = save_dedupe
        return term_group

    def compile_subject_and_verb(self, action_term: ActionTerm) -> SearchTermGroup:
        term_group = SearchTermGroup("and")
        self.add_subject_to_group(action_term, term_group)
        term_group = SearchTermGroup("and")
        self.add_subject_to_group(action_term, term_group)
        if action_term.action_verbs is not None:
            self.add_verbs_to_group(action_term.action_verbs, term_group)
        return term_group

    def compile_subject(self, action_term: ActionTerm) -> SearchTermGroup:
        term_group = SearchTermGroup("and")
        self.add_subject_to_group(action_term, term_group)
        return term_group

    def add_subject_to_group(
        self,
        action_term: ActionTerm,
        term_group: SearchTermGroup,
    ) -> None:
        if is_entity_term_list(action_term.actor_entities):
            self.add_entity_names_to_group(
                action_term.actor_entities, PropertyNames.Subject, term_group
            )

    def compile_object(self, entity: EntityTerm) -> SearchTermGroup:
        # A target can be the name of an object of an action OR the name of an entity.
        term_group = SearchTermGroup("or")
        self.add_entity_name_to_group(entity, PropertyNames.Object, term_group)
        self.add_entity_name_to_group(
            entity, PropertyNames.EntityName, term_group, self.exact_scope
        )
        self.add_entity_name_to_group(
            entity, PropertyNames.Topic, term_group, self.exact_scope
        )
        return term_group

    def add_verbs_to_group(
        self,
        verbs: VerbsTerm,
        term_group: SearchTermGroup,
    ) -> None:
        for verb in verbs.words:
            self.add_property_term_to_group("verb", verb, term_group)

    def add_entity_term_as_search_terms_to_group(
        self, entity_term: EntityTerm, term_group: SearchTermGroup
    ) -> None:
        if entity_term.is_name_pronoun:
            return
        self.add_search_term_to_group(entity_term.name, term_group)
        if entity_term.type:
            for type in entity_term.type:
                self.add_search_term_to_group(type, term_group)
        if entity_term.facets:
            for facet in entity_term.facets:
                self.add_search_term_to_group(facet.facet_name, term_group)
                self.add_search_term_to_group(facet.facet_value, term_group)

    def add_search_term_to_group(
        self,
        term: str,
        term_group: SearchTermGroup,
    ) -> None:
        if self.is_searchable_string(term):
            term_group.terms.append(SearchTerm(Term(term)))

    def add_entity_term_to_group(
        self,
        entity_term: EntityTerm,
        term_group: SearchTermGroup,
        exact_match_name=False,
    ) -> None:
        self.add_property_term_to_group(
            PropertyNames.EntityName.value,
            entity_term.name,
            term_group,
            exact_match_name,
        )
        if entity_term.type:
            for type in entity_term.type:
                self.add_property_term_to_group(
                    PropertyNames.EntityType.value, type, term_group
                )
        if entity_term.facets:
            for facet in entity_term.facets:
                name_is_wildcard = facet.facet_name == "*"
                value_is_wildcard = facet.facet_value == "*"
                match name_is_wildcard, value_is_wildcard:
                    case False, False:
                        self.add_property_term_to_group(
                            facet.facet_name,
                            facet.facet_value,
                            term_group,
                        )
                    case False, True:
                        self.add_property_term_to_group(
                            PropertyNames.FacetName.value,
                            facet.facet_name,
                            term_group,
                        )
                    case True, False:
                        self.add_property_term_to_group(
                            PropertyNames.FacetValue.value,
                            facet.facet_value,
                            term_group,
                        )
                    case True, True:
                        pass

    def add_entity_names_to_group(
        self,
        entity_terms: list[EntityTerm],
        property_name: PropertyNames,
        term_group: SearchTermGroup,
        exact_match_value: bool = False,
    ) -> None:
        for entity_term in entity_terms:
            self.add_entity_name_to_group(
                entity_term, property_name, term_group, exact_match_value
            )

    def add_entity_name_to_group(
        self,
        entity_term: EntityTerm,
        property_name: PropertyNames,
        term_group: SearchTermGroup,
        exact_match_value: bool = False,
    ) -> None:
        if not entity_term.is_name_pronoun:
            self.add_property_term_to_group(
                property_name.value,
                entity_term.name,
                term_group,
                exact_match_value,
            )

    def add_search_term_to_groupadd_entity_name_to_group(
        self,
        entity_term: EntityTerm,
        property_name: PropertyNames,
        term_group: SearchTermGroup,
        exact_match_value: bool = False,
    ) -> None:
        if not entity_term.is_name_pronoun:
            self.add_property_term_to_group(
                property_name.value,
                entity_term.name,
                term_group,
                exact_match_value,
            )

    def add_property_term_to_group(
        self,
        property_name: str,
        property_value: str,
        term_group: SearchTermGroup,
        exact_match_value=False,
    ) -> None:
        if not self.is_searchable_string(property_name):
            return
        if not self.is_searchable_string(property_value):
            return
        if self.is_noise_term(property_value):
            return
        # Dedupe any terms already added to the group earlier.
        if not self.dedupe or not self.entity_terms_added.has(
            property_name, property_value
        ):
            search_term = create_property_search_term(
                property_name, property_value, exact_match_value
            )
            term_group.terms.append(search_term)
            self.entity_terms_added.add(property_name, search_term.property_value.term)

    def is_searchable_string(self, value: str) -> bool:
        if not value or value == "*":
            return False
        return self.term_filter is None or self.term_filter(value)

    def is_noise_term(self, value: str) -> bool:
        return value.lower() in ("thing", "object", "concept", "idea", "entity")

    def should_add_scope(self, action_term: ActionTerm) -> bool:
        if not action_term or action_term.is_informational:
            return False
        if self.exact_scope:
            return True
        # If the action has no subject, disable scope.
        return is_entity_term_list(action_term.actor_entities)


# Miscellaneous helper functions.


# A type guard makes a promise to the type checker when it returns True.
def is_entity_term_list(
    terms: list[EntityTerm] | Literal["*"] | None,
) -> TypeGuard[list[EntityTerm]]:
    return isinstance(terms, list)


def optimize_or_max(term_group: SearchTermGroup) -> SearchTermGroupTypes:
    if len(term_group.terms) == 1:
        return term_group.terms[0]
    return term_group


def date_range_from_datetime_range(date_time_range: DateTimeRange) -> DateRange:
    return DateRange(
        start=datetime_from_date_time(date_time_range.start_date),
        end=(
            datetime_from_date_time(date_time_range.stop_date)
            if date_time_range.stop_date
            else None
        ),
    )


def datetime_from_date_time(date_time: DateTime) -> Datetime:
    # We ASSUME that the LLM gave the DateTime in UTC.
    # If it didn't, well, how would we know???
    dt = Datetime(
        year=date_time.date.year,
        month=date_time.date.month,
        day=date_time.date.day,
        hour=date_time.time.hour if date_time.time else 0,
        minute=date_time.time.minute if date_time.time else 0,
        second=date_time.time.seconds if date_time.time else 0,
        tzinfo=datetime.timezone.utc,
    )
    return dt


# TODO: Move to searchquerytranslator.py?
async def search_query_from_language(
    conversation: IConversation,
    translator: SearchQueryTranslator,
    query_text: str,
    model_instructions: list[typechat.PromptSection] | None = None,
) -> typechat.Result[SearchQuery]:
    time_range = await get_time_range_prompt_section_for_conversation(conversation)
    prompt_preamble: list[typechat.PromptSection] = []
    if model_instructions:
        prompt_preamble.extend(model_instructions)
    if time_range:
        prompt_preamble.append(time_range)
    # print("[" * 50)
    # print(translator.schema_str)
    # print("]" * 50)
    return await translator.translate(query_text, prompt_preamble=prompt_preamble)
