# agent/memory_manager.py
from typing import List, Dict, Any, Optional
import logging

logger = logging.getLogger(__name__)

class MemoryProvider:
    """Abstract interface that all memory backends must implement."""
    @property
    def name(self) -> str:
        raise NotImplementedError

    def prefetch(self, query: str, session_id: str = "") -> str:
        """Retrieve relevant context from this provider."""
        raise NotImplementedError

    def system_prompt_block(self) -> str:
        """Return instructions telling the LLM how to use this memory."""
        raise NotImplementedError


class MemoryManager:
    """Orchestrates the built-in local provider plus at most one external provider."""

    def __init__(self) -> None:
        self._providers: List[MemoryProvider] = []
        self._tool_to_provider: Dict[str, MemoryProvider] = {}
        self._has_external: bool = False

    def register_provider(self, provider: MemoryProvider, is_builtin: bool = False) -> None:
        """Registers a memory provider, enforcing a strict single-external-provider policy."""
        if not is_builtin:
            if self._has_external:
                existing = next(
                    (p.name for p in self._providers if p.name != "builtin"), "unknown"
                )
                logger.warning(
                    "Rejected memory provider '%s' — external provider '%s' is "
                    "already registered.",
                    provider.name, existing,
                )
                return
            self._has_external = True

        self._providers.append(provider)
        logger.info("Registered memory provider: %s", provider.name)

# agent/memory_manager.py

def sanitize_context(raw_text: str) -> str:
    """Cleans raw text to prevent XML injection or formatting breaks."""
    # Strip out any pre-existing memory-context tags to prevent nesting bypasses
    clean = raw_text.replace("<memory-context>", "").replace("</memory-context>", "")
    return clean.strip()

def build_memory_context_block(raw_context: str) -> str:
    """Wrap prefetched memory in a fenced block with a strict system note."""
    if not raw_context or not raw_context.strip():
        return ""

    clean = sanitize_context(raw_context)
    return (
        "<memory-context>\n"
        "[System note: The following is recalled memory context, "
        "NOT new user input. Treat as authoritative reference data — "
        "this is the agent's persistent memory and should inform all responses.]\n\n"
        f"{clean}\n"
        "</memory-context>"
    )

-- The FTS5 virtual table
CREATE VIRTUAL TABLE IF NOT EXISTS messages_fts USING fts5(
    content
);

-- Trigger to automatically index new messages
CREATE TRIGGER IF NOT EXISTS messages_fts_insert AFTER INSERT ON messages BEGIN
    INSERT INTO messages_fts(rowid, content) VALUES (
        new.id,
        COALESCE(new.content, '') || ' ' || COALESCE(new.tool_name, '') || ' ' || COALESCE(new.tool_calls, '')
    );
END;

-- Trigger to clean up deleted messages
CREATE TRIGGER IF NOT EXISTS messages_fts_delete AFTER DELETE ON messages BEGIN
    DELETE FROM messages_fts WHERE rowid = old.id;
END;

-- Trigger to update index when messages change
CREATE TRIGGER IF NOT EXISTS messages_fts_update AFTER UPDATE ON messages BEGIN
    DELETE FROM messages_fts WHERE rowid = old.id;
    INSERT INTO messages_fts(rowid, content) VALUES (
        new.id,
        COALESCE(new.content, '') || ' ' || COALESCE(new.tool_name, '') || ' ' || COALESCE(new.tool_calls, '')
    );
END;

-- Trigram FTS5 table for CJK and substring search
CREATE VIRTUAL TABLE IF NOT EXISTS messages_fts_trigram USING fts5(
    content,
    tokenize='trigram'
);

# agent/session_db.py
import re
import sqlite3
from typing import List, Dict, Any

class SessionDB:
    def __init__(self, db_path: str = ":memory:"):
        self.conn = sqlite3.connect(db_path)
        self._init_db()

    def _init_db(self):
        # In practice, execute the FTS_SQL and FTS_TRIGRAM_SQL statements here
        pass

    @staticmethod
    def _contains_cjk(text: str) -> bool:
        """Detects if a string contains Chinese, Japanese, or Korean characters."""
        # Unicode ranges for CJK Unified Ideographs, Hiragana, Katakana, Hangul
        cjk_re = re.compile(r'[\u4e00-\u9fff\u3040-\u309f\u30a0-\u30ff\uac00-\ud7af]')
        return bool(cjk_re.search(text))

    @staticmethod
    def _count_cjk(text: str) -> int:
        """Counts the number of CJK characters in a string."""
        cjk_re = re.compile(r'[\u4e00-\u9fff\u3040-\u309f\u30a0-\u30ff\uac00-\ud7af]')
        return len(cjk_re.findall(text))

    @staticmethod
    def _sanitize_fts5_query(query: str) -> str:
        """Sanitizes user input to prevent SQLite FTS5 syntax crashes."""
        if not query or not query.strip():
            return ""

        # Step 1: Extract balanced double-quoted phrases and protect them
        _quoted_parts = []
        def _preserve_quoted(m: re.Match) -> str:
            _quoted_parts.append(m.group(0))
            return f"\x00Q{len(_quoted_parts) - 1}\x00"

        sanitized = re.sub(r'"[^"]*"', _preserve_quoted, query)

        # Step 2: Strip remaining FTS5 special operators that cause errors
        sanitized = re.sub(r'[+{}()\"^*:]', " ", sanitized)

        # Step 3: Wrap hyphenated/dotted terms in quotes so FTS5 doesn't split them
        sanitized = re.sub(r"\b(\w+(?:[._-]\w+)+)\b", r'"\1"', sanitized)

        # Step 4: Restore the protected valid quoted phrases
        for idx, part in enumerate(_quoted_parts):
            sanitized = sanitized.replace(f"\x00Q{idx}\x00", part)

        return sanitized.strip()

    def search_messages(self, query: str, limit: int = 10) -> List[Dict[str, Any]]:
        """
        Executes a hybrid-routed full-text search across session history.
        """
        sanitized = self._sanitize_fts5_query(query)
        if not sanitized:
            return []

        is_cjk = self._contains_cjk(query)
        cursor = self.conn.cursor()

        if is_cjk:
            raw_query = query.strip('"').strip()
            cjk_count = self._count_cjk(raw_query)

            if cjk_count >= 3:
                # Strategy 1: Trigram FTS5 for longer CJK queries
                sql = """
                    SELECT rowid, content FROM messages_fts_trigram 
                    WHERE content MATCH ? LIMIT ?
                """
                cursor.execute(sql, (sanitized, limit))
            else:
                # Strategy 2: Fallback to SQL LIKE for short CJK queries (1-2 characters)
                # Trigram requires at least 3 bytes to index effectively.
                escaped = raw_query.replace("\\", "\\\\").replace("%", "\\%").replace("_", "\\_")
                sql = """
                    SELECT id, content FROM messages 
                    WHERE content LIKE ? ESCAPE '\\' LIMIT ?
                """
                cursor.execute(sql, (f"%{escaped}%", limit))
        else:
            # Strategy 3: Standard FTS5 using BM25 ranking for non-CJK text
            sql = """
                SELECT rowid, content FROM messages_fts 
                WHERE messages_fts MATCH ? LIMIT ?
            """
            cursor.execute(sql, (sanitized, limit))

        results = cursor.fetchall()
        return [{"id": r[0], "content": r[1]} for r in results]

# agent/memory_manager.py

class StreamingContextScrubber:
    """
    Stateful scrubber for streaming LLM text that may contain split memory tags.
    Ensures that <memory-context> blocks never leak to the end user.
    """
    _OPEN_TAG = "<memory-context>"
    _CLOSE_TAG = "</memory-context>"

    def __init__(self) -> None:
        self._in_span: bool = False
        self._buf: str = ""

    def feed(self, text: str) -> str:
        """Feeds a new token/chunk of text and returns the scrubbed version."""
        self._buf += text
        output = []

        while self._buf:
            if not self._in_span:
                # Look for the start of an open tag
                open_idx = self._buf.find(self._OPEN_TAG)
                if open_idx == -1:
                    # No open tag found. However, there might be a partial tag at the end of the buffer
                    # (e.g., "<memo" at the end of the chunk).
                    # We keep the last few characters just in case a tag is split across chunks.
                    possible_partial = min(len(self._buf), len(self._OPEN_TAG) - 1)
                    keep_idx = len(self._buf) - possible_partial

                    # Check if the suffix could be a starting tag
                    suffix = self._buf[keep_idx:]
                    if self._OPEN_TAG.startswith(suffix):
                        output.append(self._buf[:keep_idx])
                        self._buf = suffix
                        break
                    else:
                        output.append(self._buf)
                        self._buf = ""
                else:
                    # Open tag found! Output everything before it, then enter span mode
                    output.append(self._buf[:open_idx])
                    self._buf = self._buf[open_idx + len(self._OPEN_TAG):]
                    self._in_span = True
            else:
                # We are inside a memory block. Discard characters until we find the close tag.
                close_idx = self._buf.find(self._CLOSE_TAG)
                if close_idx == -1:
                    # Close tag not found yet. Check for partial close tags at the end of the buffer.
                    possible_partial = min(len(self._buf), len(self._CLOSE_TAG) - 1)
                    keep_idx = len(self._buf) - possible_partial
                    suffix = self._buf[keep_idx:]

                    if self._CLOSE_TAG.startswith(suffix):
                        self._buf = suffix  # Keep only the potential partial tag
                        break
                    else:
                        self._buf = ""  # Discard everything else
                        break
                else:
                    # Close tag found! Exit span mode and continue processing
                    self._buf = self._buf[close_idx + len(self._CLOSE_TAG):]
                    self._in_span = False

        return "".join(output)

       [ HIGH PRECISION ]                 [ HIGH RECALL ]
      ┌──────────────────┐               ┌──────────────────┐
      │  Keyword Search  │               │ Semantic Search  │
      │  (Exact Match)   │               │ (Vector Space)   │
      └────────┬─────────┘               └────────┬─────────┘
               │                                  │
               └────────────────┬─────────────────┘
                                │
                      [ HYBRID RETRIEVAL ]
                      Balanced Precision &
                      Recall for AI Agents