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

import bisect
from collections.abc import Callable, Iterable, Iterator
from dataclasses import dataclass, field
import heapq
import math
import sys
from typing import cast, Set  # Set is an alias for builtin set

from .interfaces import (
    ICollection,
    IMessage,
    IMessageCollection,
    ISemanticRefCollection,
    Knowledge,
    KnowledgeType,
    MessageOrdinal,
    ScoredMessageOrdinal,
    ScoredSemanticRefOrdinal,
    SemanticRef,
    SemanticRefOrdinal,
    Term,
    TextRange,
)


@dataclass
class Match[T]:
    value: T
    score: float
    hit_count: int
    related_score: float
    related_hit_count: int


# TODO: sortMatchesByRelevance


class MatchAccumulator[T]:
    def __init__(self):
        self._matches: dict[T, Match[T]] = {}

    def __len__(self) -> int:
        return len(self._matches)

    def __iter__(self) -> Iterator[Match[T]]:
        return iter(self._matches.values())

    def __contains__(self, value: T) -> bool:
        return value in self._matches

    def get_match(self, value: T) -> Match[T] | None:
        return self._matches.get(value)

    def set_match(self, match: Match[T]) -> None:
        self._matches[match.value] = match

    # TODO: Maybe make the callers call clear_matches()?
    def set_matches(self, matches: Iterable[Match[T]], *, clear: bool = False) -> None:
        if clear:
            self.clear_matches()
        for match in matches:
            self.set_match(match)

    def get_max_hit_count(self) -> int:
        count = 0
        for match in self._matches.values():
            count = max(count, match.hit_count)
        return count

    # TODO: Rename to add_exact if we ever add add_related
    def add(self, value: T, score: float, is_exact_match: bool = True) -> None:
        existing_match = self.get_match(value)
        if existing_match is not None:
            if is_exact_match:
                existing_match.hit_count += 1
                existing_match.score += score
            else:
                existing_match.related_hit_count += 1
                existing_match.related_score += score
        else:
            if is_exact_match:
                self.set_match(
                    Match(
                        value,
                        hit_count=1,
                        score=score,
                        related_hit_count=0,
                        related_score=0.0,
                    )
                )
            else:
                self.set_match(
                    Match(
                        value,
                        hit_count=1,
                        score=0.0,
                        related_hit_count=1,
                        related_score=score,
                    )
                )

    def add_union(self, other: "MatchAccumulator[T]") -> None:
        """Add matches from another collection of matches."""
        for other_match in other:
            existing_match = self.get_match(other_match.value)
            if existing_match is None:
                self.set_match(other_match)
            else:
                self.combine_matches(existing_match, other_match)

    def intersect(
        self, other: "MatchAccumulator[T]", intersection: "MatchAccumulator[T]"
    ) -> "MatchAccumulator[T]":
        """Intersect with another collection of matches."""
        for self_match in self:
            other_match = other.get_match(self_match.value)
            if other_match is not None:
                self.combine_matches(self_match, other_match)
                intersection.set_match(self_match)
        return intersection

    def combine_matches(self, match: Match[T], other: Match[T]) -> None:
        """Combine the other match into the first."""
        match.hit_count += other.hit_count
        match.score += other.score
        match.related_hit_count += other.related_hit_count
        match.related_score += other.related_score

    def calculate_total_score(
        self, scorer: Callable[[Match[T]], None] | None = None
    ) -> None:
        if scorer is None:
            scorer = add_smooth_related_score_to_match_score
        for match in self:
            scorer(match)

    def get_sorted_by_score(self, min_hit_count: int | None = None) -> list[Match[T]]:
        """Get matches sorted by score"""
        if len(self._matches) == 0:
            return []
        matches = [*self._matches_with_min_hit_count(min_hit_count)]
        matches.sort(key=lambda m: m.score, reverse=True)
        return matches

    def get_top_n_scoring(
        self,
        max_matches: int | None = None,
        min_hit_count: int | None = None,
    ) -> list[Match[T]]:
        """Get the top N scoring matches."""
        if not self._matches:
            return []
        if max_matches and max_matches > 0:
            top_list = TopNList[T](max_matches)
            for match in self._matches_with_min_hit_count(min_hit_count):
                top_list.push(match.value, match.score)
            ranked = top_list.by_rank()
            return [self._matches[match.item] for match in ranked]
        else:
            return self.get_sorted_by_score(min_hit_count)

    def get_with_hit_count(self, min_hit_count: int) -> list[Match[T]]:
        """Get matches with a minimum hit count."""
        return list(self.matches_with_min_hit_count(min_hit_count))

    def get_matches(
        self, predicate: Callable[[Match[T]], bool] | None = None
    ) -> Iterator[Match[T]]:
        """Iterate over all matches."""
        if predicate is None:
            return iter(self._matches.values())
        else:
            return filter(predicate, self._matches.values())

    def get_matched_values(self) -> Iterator[T]:
        """Iterate over all matched values."""
        return iter(self._matches)

    def clear_matches(self):
        self._matches.clear()

    def select_top_n_scoring(
        self,
        max_matches: int | None = None,
        min_hit_count: int | None = None,
    ) -> int:
        """Retain only the top N matches sorted by score."""
        top_n = self.get_top_n_scoring(max_matches, min_hit_count)
        self.set_matches(top_n, clear=True)
        return len(top_n)

    def select_with_hit_count(self, min_hit_count: int) -> int:
        """Retain only matches with a minimum hit count."""
        matches = self.get_with_hit_count(min_hit_count)
        self.set_matches(matches, clear=True)
        return len(matches)

    def _matches_with_min_hit_count(
        self, min_hit_count: int | None
    ) -> Iterable[Match[T]]:
        """Get matches with a minimum hit count"""
        if min_hit_count is not None and min_hit_count > 0:
            return self.get_matches(lambda m: m.hit_count >= min_hit_count)
        else:
            return self._matches.values()

    def matches_with_min_hit_count(
        self, min_hit_count: int | None
    ) -> Iterable[Match[T]]:
        if min_hit_count is not None and min_hit_count > 0:
            return filter(lambda m: m.hit_count >= min_hit_count, self.get_matches())
        else:
            return self._matches.values()


def get_smooth_score(
    total_score: float,
    hit_count: int,
) -> float:
    """See the long comment in collections.ts for an explanation."""
    if hit_count > 0:
        if hit_count == 1:
            return total_score
        avg = total_score / hit_count
        smooth_avg = math.log(hit_count + 1) * avg
        return smooth_avg
    else:
        return 0.0


def add_smooth_related_score_to_match_score[T](match: Match[T]) -> None:
    """Add the smooth related score to the match score."""
    if match.related_hit_count > 0:
        # Related term matches can be noisy and duplicative.
        # See the comment on getSmoothScore  in collections.ts.
        smooth_related_score = get_smooth_score(
            match.related_score, match.related_hit_count
        )
        match.score += smooth_related_score


def smooth_match_score[T](match: Match[T]) -> None:
    if match.hit_count > 0:
        match.score = get_smooth_score(match.score, match.hit_count)


type KnowledgePredicate[T: Knowledge] = Callable[[T], bool]


class SemanticRefAccumulator(MatchAccumulator[SemanticRefOrdinal]):
    def __init__(self, search_term_matches: set[str] = set()):
        super().__init__()
        self.search_term_matches = search_term_matches

    def add_term_matches(
        self,
        search_term: Term,
        scored_refs: Iterable[ScoredSemanticRefOrdinal] | None,
        is_exact_match: bool,
        *,
        weight: float | None = None,
    ) -> None:
        """Add term matches to the accumulator"""
        if scored_refs is not None:
            if weight is None:
                weight = search_term.weight
                if weight is None:
                    weight = 1.0
            for scored_ref in scored_refs:
                self.add(
                    scored_ref.semantic_ref_ordinal,
                    scored_ref.score * weight,
                    is_exact_match,
                )
            self.search_term_matches.add(search_term.text)

    def add_term_matches_if_new(
        self,
        search_term: Term,
        scored_refs: Iterable[ScoredSemanticRefOrdinal] | None,
        is_exact_match: bool,
        weight: float | None = None,
    ) -> None:
        """Add term matches if they are new."""
        if scored_refs is not None:
            if weight is None:
                weight = search_term.weight
                if weight is None:
                    weight = 1.0
            for scored_ref in scored_refs:
                if scored_ref.semantic_ref_ordinal not in self:
                    self.add(
                        scored_ref.semantic_ref_ordinal,
                        scored_ref.score * weight,
                        is_exact_match,
                    )
            self.search_term_matches.add(search_term.text)

    async def get_semantic_refs(
        self,
        semantic_refs: ISemanticRefCollection,
        predicate: Callable[[SemanticRef], bool],
    ) -> list[SemanticRef]:
        result = []
        for match in self:
            semantic_ref = await semantic_refs.get_item(match.value)
            if predicate is None or predicate(semantic_ref):
                result.append(semantic_ref)
        return result

    def get_matches_of_type[T: Knowledge](
        self,
        semantic_refs: list[SemanticRef],
        knowledgeType: KnowledgeType,
        predicate: KnowledgePredicate[T] | None = None,
    ) -> Iterable[Match[SemanticRefOrdinal]]:
        for match in self:
            semantic_ref = semantic_refs[match.value]
            if predicate is None or predicate(cast(T, semantic_ref.knowledge)):
                yield match

    async def group_matches_by_type(
        self,
        semantic_refs: ISemanticRefCollection,
    ) -> dict[KnowledgeType, "SemanticRefAccumulator"]:
        groups: dict[KnowledgeType, SemanticRefAccumulator] = {}
        for match in self:
            semantic_ref = await semantic_refs.get_item(match.value)
            group = groups.get(semantic_ref.knowledge.knowledge_type)
            if group is None:
                group = SemanticRefAccumulator()
                group.search_term_matches = self.search_term_matches
                groups[semantic_ref.knowledge.knowledge_type] = group
            group.set_match(match)
        return groups

    async def get_matches_in_scope(
        self,
        semantic_refs: ISemanticRefCollection,
        ranges_in_scope: "TextRangesInScope",
    ) -> "SemanticRefAccumulator":
        accumulator = SemanticRefAccumulator(self.search_term_matches)
        for match in self:
            if ranges_in_scope.is_range_in_scope(
                (await semantic_refs.get_item(match.value)).range
            ):
                accumulator.set_match(match)
        return accumulator

    def add_union(self, other: "MatchAccumulator[SemanticRefOrdinal]") -> None:
        """Add matches from another SemanticRefAccumulator."""
        assert isinstance(other, SemanticRefAccumulator)
        super().add_union(other)
        self.search_term_matches.update(other.search_term_matches)

    def intersect(
        self,
        other: MatchAccumulator[SemanticRefOrdinal],
        intersection: MatchAccumulator[SemanticRefOrdinal] | None = None,
    ) -> "SemanticRefAccumulator":
        """Intersect with another SemanticRefAccumulator."""
        assert isinstance(other, SemanticRefAccumulator)
        if intersection is None:
            intersection = SemanticRefAccumulator()
        else:
            assert isinstance(intersection, SemanticRefAccumulator)
        super().intersect(other, intersection)
        if len(intersection) > 0:
            intersection.search_term_matches.update(self.search_term_matches)
            intersection.search_term_matches.update(other.search_term_matches)
        return intersection

    def to_scored_semantic_refs(self) -> list[ScoredSemanticRefOrdinal]:
        """Convert the accumulator to a list of scored semantic references."""
        return [
            ScoredSemanticRefOrdinal(
                semantic_ref_ordinal=match.value,
                score=match.score,
            )
            for match in self.get_sorted_by_score()
        ]


class MessageAccumulator(MatchAccumulator[MessageOrdinal]):
    def __init__(self, matches: list[Match[MessageOrdinal]] | None = None):
        super().__init__()
        if matches:
            self.set_matches(matches)

    def add(
        self, value: MessageOrdinal, score: float, is_exact_match: bool = True
    ) -> None:
        match = self.get_match(value)
        if match is None:
            match = Match(value, score, 1, 0.0, 0)
            self.set_match(match)
        elif score > match.score:
            match.score = score
            # TODO: Question(Guido->Umesh): Why not increment hit_count always?
            match.hit_count += 1

    # TODO: add_messages_from_locations

    def add_messages_for_semantic_ref(
        self,
        semantic_ref: SemanticRef,
        score: float,
    ) -> None:
        message_ordinal_start = semantic_ref.range.start.message_ordinal
        if semantic_ref.range.end is not None:
            message_ordinal_end = semantic_ref.range.end.message_ordinal
            for message_ordinal in range(
                message_ordinal_start, message_ordinal_end + 1
            ):
                self.add(message_ordinal, score)
        else:
            self.add(message_ordinal_start, score)

    def add_scored_matches(self, scored_ordinals: list[ScoredMessageOrdinal]) -> None:
        """Add scored message ordinals to the accumulator."""
        for scored_ordinal in scored_ordinals:
            self.add(scored_ordinal.message_ordinal, scored_ordinal.score)

    def intersect(
        self,
        other: MatchAccumulator[MessageOrdinal],
        intersection: MatchAccumulator[MessageOrdinal] | None = None,
    ) -> "MessageAccumulator":
        if intersection is None:
            intersection = MessageAccumulator()
        else:
            assert isinstance(intersection, MessageAccumulator)
        super().intersect(other, intersection)
        return intersection

    def smooth_scores(self) -> None:
        for match in self:
            smooth_match_score(match)

    def to_scored_message_ordinals(self) -> list[ScoredMessageOrdinal]:
        sorted_matches = self.get_sorted_by_score()
        return [ScoredMessageOrdinal(m.value, m.score) for m in sorted_matches]

    async def select_messages_in_budget(
        self, messages: IMessageCollection, max_chars_in_budget: int
    ) -> None:
        """Select messages that fit within the character budget."""
        scored_matches = self.get_sorted_by_score()
        ranked_ordinals = [m.value for m in scored_matches]
        message_count_in_budget = await get_count_of_messages_in_char_budget(
            messages, ranked_ordinals, max_chars_in_budget
        )
        self.clear_matches()
        if message_count_in_budget > 0:
            scored_matches = scored_matches[:message_count_in_budget]
            self.set_matches(scored_matches)

    @staticmethod
    def from_scored_ordinals(
        ordinals: list[ScoredMessageOrdinal] | None,
    ) -> "MessageAccumulator":
        """Create a MessageAccumulator from scored ordinals."""
        accumulator = MessageAccumulator()
        if ordinals and len(ordinals) > 0:
            accumulator.add_scored_matches(ordinals)
        return accumulator


# TODO: intersectScoredMessageOrdinals


@dataclass
class TextRangeCollection(Iterable[TextRange]):
    _ranges: list[TextRange]

    def __init__(
        self,
        ranges: list[TextRange] | None = None,
        ensure_sorted: bool = False,
    ) -> None:
        if ensure_sorted:
            self._ranges = []
            if ranges:
                self.add_ranges(ranges)
        else:
            self._ranges = ranges if ranges is not None else []

    def __len__(self) -> int:
        return len(self._ranges)

    def __iter__(self) -> Iterator[TextRange]:
        return iter(self._ranges)

    def get_ranges(self) -> list[TextRange]:
        return self._ranges  # TODO: Maybe return a copy?

    def add_range(self, text_range: TextRange) -> bool:
        # This assumes TextRanges are totally ordered.
        pos = bisect.bisect_left(self._ranges, text_range)
        if pos < len(self._ranges) and self._ranges[pos] == text_range:
            return False
        self._ranges.insert(pos, text_range)
        return True

    def add_ranges(self, text_ranges: "list[TextRange] | TextRangeCollection") -> None:
        if isinstance(text_ranges, list):
            for text_range in text_ranges:
                self.add_range(text_range)
        else:
            assert isinstance(text_ranges, TextRangeCollection)
            for text_range in text_ranges._ranges:
                self.add_range(text_range)

    def is_in_range(self, inner_range: TextRange) -> bool:
        if len(self._ranges) == 0:
            return False
        i = bisect.bisect_left(self._ranges, inner_range)
        for outer_range in self._ranges[i:]:
            if outer_range.start > inner_range.start:
                break
            if inner_range in outer_range:
                return True
        return False


@dataclass
class TextRangesInScope:
    text_ranges: list[TextRangeCollection] | None = None

    def add_text_ranges(
        self,
        ranges: TextRangeCollection,
    ) -> None:
        if self.text_ranges is None:
            self.text_ranges = []
        self.text_ranges.append(ranges)

    def is_range_in_scope(self, inner_range: TextRange) -> bool:
        if self.text_ranges is not None:
            # Since outer ranges come from a set of range selectors, they may overlap, or may not agree.
            # Outer ranges allowed by say a date range selector... may not be allowed by a tag selector.
            # We have a very simple impl: we don't intersect/union ranges yet.
            # Instead, we ensure that the inner range is not rejected by any outer ranges.
            for outer_ranges in self.text_ranges:
                if not outer_ranges.is_in_range(inner_range):
                    return False
        return True


@dataclass
class TermSet:
    """A collection of terms with support for adding, updating, and retrieving terms."""

    terms: dict[str, Term]

    def __init__(self, terms: list[Term] | None = None):
        self.terms = {}
        self.add_or_union(terms)

    def __len__(self) -> int:
        """Return the number of terms in the set."""
        return len(self.terms)

    def add(self, term: Term) -> bool:
        """Add a term to the set if it doesn't already exist."""
        if term.text in self.terms:
            return False
        self.terms[term.text] = term
        return True

    def add_or_union(self, terms: Term | list[Term] | None) -> None:
        """Add a term or merge a list of terms into the set."""
        if terms is None:
            return
        if isinstance(terms, list):
            for term in terms:
                self.add_or_union(term)
        else:
            existing_term = self.terms.get(terms.text)
            if existing_term:
                existing_score = existing_term.weight or 0
                new_score = terms.weight or 0
                if new_score > existing_score:
                    existing_term.weight = new_score
            else:
                self.terms[terms.text] = terms

    def get(self, term: str | Term) -> Term | None:
        """Retrieve a term by its text."""
        return self.terms.get(term if isinstance(term, str) else term.text)

    def get_weight(self, term: Term) -> float | None:
        """Retrieve the weight of a term."""
        t = self.terms.get(term.text)
        return t.weight if t is not None else None

    def __contains__(self, term: Term) -> bool:
        """Check if a term exists in the set."""
        return term.text in self.terms

    def remove(self, term: Term):
        """Remove a term from the set, if present."""
        self.terms.pop(term.text, None)

    def clear(self):
        """Clear all terms from the set."""
        self.terms.clear()

    def values(self) -> list[Term]:
        """Retrieve all terms in the set."""
        return list(self.terms.values())


@dataclass
class PropertyTermSet:
    """A collection of property terms with support for adding, checking, and clearing."""

    terms: dict[str, Term] = field(default_factory=dict[str, Term])

    def add(self, property_name: str, property_value: Term) -> None:
        """Add a property term to the set."""
        key = self._make_key(property_name, property_value)
        if key not in self.terms:
            self.terms[key] = property_value

    def has(self, property_name: str, property_value: Term | str) -> bool:
        """Check if a property term exists in the set."""
        key = self._make_key(property_name, property_value)
        return key in self.terms

    def clear(self) -> None:
        """Clear all property terms from the set."""
        self.terms.clear()

    def _make_key(self, property_name: str, property_value: Term | str) -> str:
        """Create a unique key for a property term."""
        value = (
            property_value if isinstance(property_value, str) else property_value.text
        )
        return f"{property_name}:{value}"


# TODO: unionArrays
# TODO: union
# TODO: addToSet
# TODO: setUnion
# TODO: setIntersect
# TODO: getBatches


@dataclass
class Scored[T]:
    item: T
    score: float

    def __lt__(self, other: "Scored[T]") -> bool:
        return self.score < other.score

    def __gt__(self, other: "Scored[T]") -> bool:
        return self.score > other.score

    def __le__(self, other: "Scored[T]") -> bool:
        return self.score <= other.score

    def __ge__(self, other: "Scored[T]") -> bool:
        return self.score >= other.score


# Implementation change compared to TS version: Use heapq; no sentinel.
# API change: pop/top are not properties.
class TopNCollection[T]:
    """A collection that maintains the top N items based on their scores."""

    def __init__(self, max_count: int):
        self._max_count = max_count
        self._heap: list[Scored[T]] = []

    def __len__(self) -> int:
        return len(self._heap)

    def reset(self) -> None:
        self._heap = []

    def pop(self) -> Scored[T]:
        return heapq.heappop(self._heap)

    def top(self) -> Scored[T]:
        return self._heap[0]

    def push(self, item: T, score: float) -> None:
        if len(self._heap) < self._max_count:
            heapq.heappush(self._heap, Scored(item, score))
        else:
            heapq.heappushpop(self._heap, Scored(item, score))

    def by_rank(self) -> list[Scored[T]]:
        return sorted(self._heap, reverse=True)

    def values_by_rank(self) -> list[T]:
        return [item.item for item in self.by_rank()]


class TopNList[T](TopNCollection[T]):
    """Alias for TopNCollection."""


class TopNListAll[T](TopNList[T]):
    """A Top N list for N = infinity (approximated by sys.maxsize)."""

    def __init__(self):
        super().__init__(sys.maxsize)


def get_top_k[T](
    scored_items: Iterable[Scored[T]],
    top_k: int,
) -> list[Scored[T]]:
    """A function to get the top K of an unsorted list of scored items."""
    top_n_list = TopNCollection[T](top_k)
    for scored_item in scored_items:
        top_n_list.push(scored_item.item, scored_item.score)
    return top_n_list.by_rank()


def add_to_set[T](
    set: Set[T],
    values: Iterable[T],
) -> None:
    """Add values to a set."""
    set.update(values)


def get_message_char_count(message: IMessage) -> int:
    """Get the character count of a message."""
    total = 0
    for chunk in message.text_chunks:
        total += len(chunk)
    return total


async def get_count_of_messages_in_char_budget(
    messages: IMessageCollection,
    message_ordinals: Iterable[MessageOrdinal],
    max_chars_in_budget: int,
) -> int:
    """Get the count of messages that fit within the character budget."""
    i = 0
    total_char_count = 0
    for message_ordinal in message_ordinals:
        message = await messages.get_item(message_ordinal)
        message_char_count = get_message_char_count(message)
        if message_char_count + total_char_count > max_chars_in_budget:
            break
        total_char_count += message_char_count
        i += 1
    return i


# Storage Collections - moved here to avoid circular imports
class MemoryCollection[T, TOrdinal: int](ICollection[T, TOrdinal]):
    """A generic in-memory (non-persistent) collection class."""

    def __init__(self, items: list[T] | None = None):
        self.items: list[T] = items or []

    async def size(self) -> int:
        return len(self.items)

    def __aiter__(self):
        """Return an async iterator over the collection."""
        return self._async_iterator()

    async def _async_iterator(self):
        """Async generator that yields items from the collection."""
        for item in self.items:
            yield item

    async def get_item(self, arg: int) -> T:
        """Retrieve an item by its ordinal."""
        return self.items[arg]

    async def get_slice(self, start: int, stop: int) -> list[T]:
        """Retrieve a slice of items."""
        return self.items[start:stop]

    async def get_multiple(self, arg: list[TOrdinal]) -> list[T]:
        """Retrieve multiple items by their ordinals."""
        return [await self.get_item(ordinal) for ordinal in arg]

    @property
    def is_persistent(self) -> bool:
        return False

    async def append(self, item: T) -> None:
        """Append items to the collection."""
        self.items.append(item)

    async def extend(self, items: Iterable[T]) -> None:
        """Extend the collection with multiple items."""
        self.items.extend(items)


class MemorySemanticRefCollection(MemoryCollection[SemanticRef, SemanticRefOrdinal]):
    """A collection of semantic references."""


class MemoryMessageCollection[TMessage: IMessage](
    MemoryCollection[TMessage, MessageOrdinal]
):
    """A collection of messages."""


# Alias for backward compatibility
SemanticRefCollection = MemorySemanticRefCollection
