"""Retrieval helpers for memory ranking."""
from __future__ import annotations

from difflib import SequenceMatcher
import re
from typing import TYPE_CHECKING

if TYPE_CHECKING:
    from logos_adk.memory import MemoryEntry

VALID_RETRIEVALS = {"keyword", "hybrid", "semantic"}


def tokenize(text: str) -> list[str]:
    """Tokenize text for retrieval scoring."""
    return [
        token
        for token in re.findall(r"[a-zA-Z0-9_]+", text.lower())
        if len(token) >= 3
    ]


def retrieval_queries(query: str, retrieval: str) -> list[str]:
    """Expand a user query into backend search queries."""
    if retrieval == "keyword":
        return [query]

    queries = [query]
    queries.extend(tokenize(query))
    return list(dict.fromkeys(q for q in queries if q))


def score_entry(query: str, entry: MemoryEntry, retrieval: str) -> float:
    """Score a memory entry for a query under a retrieval mode."""
    query_lower = query.lower()
    content_lower = entry.content.lower()
    query_tokens = set(tokenize(query))
    entry_tokens = set(tokenize(entry.content))

    overlap = len(query_tokens & entry_tokens)
    overlap_ratio = overlap / len(query_tokens) if query_tokens else 0.0
    phrase_bonus = 1.0 if query_lower and query_lower in content_lower else 0.0
    substring_bonus = 0.0
    if query_tokens:
        substring_bonus = sum(token in content_lower for token in query_tokens) / len(query_tokens)

    if retrieval == "keyword":
        return phrase_bonus * 10 + substring_bonus * 3 + overlap_ratio

    lexical_similarity = SequenceMatcher(None, query_lower, content_lower).ratio()
    if retrieval == "hybrid":
        return phrase_bonus * 8 + overlap_ratio * 5 + substring_bonus * 2 + lexical_similarity

    return phrase_bonus * 6 + overlap_ratio * 4 + substring_bonus * 2 + lexical_similarity * 3
