from __future__ import annotations

import re
from contextvars import ContextVar, Token
from dataclasses import dataclass
from enum import Enum
from typing import Any

from utils.recall_signals import CROSS_THREAD_SCOPE_HINT_RE


class MemoryRetrievalTier(str, Enum):
    LOCAL = "local"
    RECALL = "recall"
    DEEP = "deep"


@dataclass(frozen=True)
class RetrievalTierSnapshot:
    tier: MemoryRetrievalTier
    max_sessions_cap: int
    reason: str


@dataclass(frozen=True)
class RetrievalTierHistoryResult:
    snapshot: RetrievalTierSnapshot
    history_depth: int
    history_match_index: int | None


_CURRENT_RETRIEVAL_TIER: ContextVar[RetrievalTierSnapshot | None] = ContextVar(
    "mugi_current_retrieval_tier",
    default=None,
)

_DEFAULT_PROFILE = "conservative"
_DEFAULT_HISTORY_DEPTH = 5
_MAX_CLASSIFIER_CHARS = 4000
_MAX_SESSIONS_BY_TIER: dict[MemoryRetrievalTier, int] = {
    MemoryRetrievalTier.LOCAL: 1,
    MemoryRetrievalTier.RECALL: 3,
    MemoryRetrievalTier.DEEP: 5,
}

_DEEP_PATTERNS = (
    re.compile(r"\bsearch\s+(all|across)\s+(our\s+)?(past|previous|older)\s+(chat|chats|conversations?)\b", re.IGNORECASE),
    re.compile(r"\bdeep\s+recall\b", re.IGNORECASE),
    re.compile(r"\bfull\s+history\b", re.IGNORECASE),
    # Apostrophe variants: ASCII + U+2018/U+2019 (typographic quotes from paste/OS).
    re.compile(r"\beverything\s+we['\u2018\u2019]ve\s+discussed\b", re.IGNORECASE),
)

_RECALL_PATTERNS_STRONG = (
    re.compile(r"\bremember\s+when\b", re.IGNORECASE),
    re.compile(r"\bwe\s+talked\s+about\b", re.IGNORECASE),
    re.compile(r"\blast\s+time\b", re.IGNORECASE),
    re.compile(r"\bprevious\s+conversation\b", re.IGNORECASE),
    re.compile(r"\bwhat\s+did\s+i\s+say\b", re.IGNORECASE),
    re.compile(r"\bfrom\s+our\s+earlier\s+conversation\b", re.IGNORECASE),
)

_RECALL_PATTERNS_BALANCED_ONLY = (
    re.compile(r"\bearlier\s+you\s+mentioned\b", re.IGNORECASE),
    re.compile(r"\bfrom\s+before\b", re.IGNORECASE),
    re.compile(r"\bwe\s+covered\s+this\b", re.IGNORECASE),
)

_LOCAL_PATTERNS = (
    # Avoid matching "test" inside hyphenated compounds like stress-test / e2e-test.
    re.compile(
        r"\b(refactor|implement|write|fix|debug)\b|\b(?<!-)test\b",
        re.IGNORECASE,
    ),
    re.compile(r"\b(open|read|edit)\s+this\s+file\b", re.IGNORECASE),
)
# "In this thread" often means *another* thread when the user pastes a bridge id or UUID;
# only treat it as a local-only signal when there are no cross-thread hints.
_LOCAL_THREAD_SCOPE_RE = re.compile(r"\b(in|from)\s+this\s+(thread|chat)\b", re.IGNORECASE)
_SYNTHETIC_USER_RE = re.compile(
    r"^\s*Wake:\s+\w+[\s\S]*?^\s*Local time:\s+.+$",
    re.IGNORECASE | re.MULTILINE,
)


def _normalize_profile(profile: str | None) -> str:
    value = str(profile or "").strip().lower()
    return value if value in {"conservative", "balanced"} else _DEFAULT_PROFILE


def _cap_for_tier(tier: MemoryRetrievalTier) -> int:
    return int(_MAX_SESSIONS_BY_TIER.get(tier, 3))


def clamp_search_sessions_count(value: Any, *, cap: int = 5) -> int:
    try:
        cap_int = int(cap)
    except Exception:
        cap_int = 5
    cap_int = max(1, min(cap_int, 5))
    try:
        requested_int = int(value)
    except Exception:
        requested_int = 3
    requested_int = max(1, requested_int)
    return min(requested_int, cap_int)


def looks_like_synthetic_user_message(text: str | None) -> bool:
    payload = (text or "").strip()
    if not payload:
        return False
    if _SYNTHETIC_USER_RE.search(payload):
        return True
    # Common heartbeat assembly fallback marker from agent_plane.
    if payload.startswith("No pre-loaded workspace context found for this task."):
        return True
    return False


def choose_retrieval_tier(
    latest_user_message: str | None,
    *,
    profile: str | None = None,
    synthetic_user_message: bool = False,
) -> RetrievalTierSnapshot:
    text = (latest_user_message or "").strip()
    norm_profile = _normalize_profile(profile)
    lowered = text.lower()
    if synthetic_user_message:
        tier = MemoryRetrievalTier.LOCAL
        return RetrievalTierSnapshot(tier=tier, max_sessions_cap=_cap_for_tier(tier), reason="synthetic_user_message")

    for pattern in _DEEP_PATTERNS:
        if pattern.search(text):
            tier = MemoryRetrievalTier.DEEP
            return RetrievalTierSnapshot(tier=tier, max_sessions_cap=_cap_for_tier(tier), reason="deep_phrase")

    for pattern in _RECALL_PATTERNS_STRONG:
        if pattern.search(text):
            tier = MemoryRetrievalTier.RECALL
            return RetrievalTierSnapshot(tier=tier, max_sessions_cap=_cap_for_tier(tier), reason="recall_phrase")

    if norm_profile == "balanced":
        for pattern in _RECALL_PATTERNS_BALANCED_ONLY:
            if pattern.search(text):
                tier = MemoryRetrievalTier.RECALL
                return RetrievalTierSnapshot(tier=tier, max_sessions_cap=_cap_for_tier(tier), reason="recall_phrase_balanced")

    for pattern in _LOCAL_PATTERNS:
        if pattern.search(text):
            tier = MemoryRetrievalTier.LOCAL
            return RetrievalTierSnapshot(tier=tier, max_sessions_cap=_cap_for_tier(tier), reason="explicit_local_task")

    if _LOCAL_THREAD_SCOPE_RE.search(text) and not CROSS_THREAD_SCOPE_HINT_RE.search(text):
        tier = MemoryRetrievalTier.LOCAL
        return RetrievalTierSnapshot(tier=tier, max_sessions_cap=_cap_for_tier(tier), reason="explicit_local_task")

    # Tiering-enabled default should preserve today's recall-oriented behavior.
    tier = MemoryRetrievalTier.RECALL
    if not lowered:
        reason = "empty_default_recall"
    else:
        reason = "ambiguous_default_recall"
    return RetrievalTierSnapshot(tier=tier, max_sessions_cap=_cap_for_tier(tier), reason=reason)


def _truncate_for_classifier(text: str, *, max_chars: int) -> str:
    if max_chars <= 0:
        return text
    if len(text) <= max_chars:
        return text
    return text[:max_chars]


def choose_retrieval_tier_from_recent_user_texts(
    texts_newest_first: list[str],
    *,
    profile: str | None = None,
    synthetic_latest: bool = False,
    max_history_messages: int = _DEFAULT_HISTORY_DEPTH,
    max_chars_per_message: int = _MAX_CLASSIFIER_CHARS,
) -> RetrievalTierHistoryResult:
    """Pick a retrieval tier using recent user texts (newest first).

    Raises ValueError when no non-empty text is provided.
    """
    if max_history_messages <= 0:
        max_history_messages = _DEFAULT_HISTORY_DEPTH

    candidates = [str(t).strip() for t in texts_newest_first if str(t).strip()]
    if not candidates:
        raise ValueError("texts_newest_first must contain at least one non-empty user message")
    candidates = candidates[:max_history_messages]

    latest_text = _truncate_for_classifier(candidates[0], max_chars=max_chars_per_message)
    if synthetic_latest:
        snapshot = choose_retrieval_tier(
            latest_text,
            profile=profile,
            synthetic_user_message=True,
        )
        return RetrievalTierHistoryResult(snapshot=snapshot, history_depth=1, history_match_index=0)

    for idx, text in enumerate(candidates):
        snapshot = choose_retrieval_tier(
            _truncate_for_classifier(text, max_chars=max_chars_per_message),
            profile=profile,
            synthetic_user_message=False,
        )
        if snapshot.reason not in {"ambiguous_default_recall", "empty_default_recall"}:
            return RetrievalTierHistoryResult(
                snapshot=snapshot,
                history_depth=len(candidates),
                history_match_index=idx,
            )

    fallback = choose_retrieval_tier(
        latest_text,
        profile=profile,
        synthetic_user_message=False,
    )
    return RetrievalTierHistoryResult(
        snapshot=fallback,
        history_depth=len(candidates),
        history_match_index=None,
    )


def tier_skills_prompt_suffix(tier: MemoryRetrievalTier) -> str:
    if tier == MemoryRetrievalTier.LOCAL:
        return (
            "- Current retrieval tier is local: prioritize this thread, workspace files, and saved workflows before searching past threads.\n"
        )
    return (
        "- Current retrieval tier supports cross-thread continuity: use search_sessions when needed, following Session Search guidance below.\n"
    )


def apply_search_sessions_tier_cap(args: dict[str, Any], cap: int) -> dict[str, Any]:
    out = dict(args or {})
    out["max_sessions"] = clamp_search_sessions_count(out.get("max_sessions", 3), cap=cap)
    return out


def set_current_retrieval_tier(snapshot: RetrievalTierSnapshot | None) -> Token:
    return _CURRENT_RETRIEVAL_TIER.set(snapshot)


def reset_current_retrieval_tier(token: Token) -> None:
    _CURRENT_RETRIEVAL_TIER.reset(token)


def get_current_retrieval_tier() -> RetrievalTierSnapshot | None:
    return _CURRENT_RETRIEVAL_TIER.get()
