"""Sandboxed Python execution using RestrictedPython.

This module is split into focused components:
- MessageDatabase: Database connection and queries
- LlmOrchestrator: Recursive LLM sub-calls
- RestrictedExecutor: Sandboxed code execution
- Sandbox: Facade composing the above
"""

from __future__ import annotations

import sqlite3
import re
import json
import sys
from collections import Counter, defaultdict
from datetime import datetime, timedelta
from io import StringIO
from typing import Any, TYPE_CHECKING

from RestrictedPython import compile_restricted, safe_globals
from RestrictedPython.Eval import default_guarded_getiter, default_guarded_getitem
from RestrictedPython.Guards import guarded_iter_unpack_sequence, safer_getattr

from .prompts import SUB_QUERY_PROMPT

# Apple Core Foundation epoch: January 1, 2001 (in Unix seconds)
MAC_EPOCH_OFFSET = 978307200


def cf_to_datetime(timestamp: int) -> datetime:
    """Convert Apple Core Foundation timestamp to datetime.

    Apple timestamps are in nanoseconds since January 1, 2001.
    """
    # Convert nanoseconds to seconds and adjust for Mac epoch
    unix_seconds = (timestamp / 1e9) + MAC_EPOCH_OFFSET
    return datetime.fromtimestamp(unix_seconds)


def unix_seconds_to_cf_nanoseconds(unix_seconds: float) -> int:
    """Convert Unix seconds to Core Foundation nanoseconds since January 1, 2001."""
    return int((unix_seconds - MAC_EPOCH_OFFSET) * 1e9)


if TYPE_CHECKING:
    from .rlm import ExecutionContext


class SandboxError(Exception):
    """Raised when sandbox execution fails."""
    pass


class MessageDatabase:
    """Database connection and query handling.

    Responsibilities:
    - Read-only database connection management
    - SQL query execution with validation
    - GUID tracking for benchmark mode
    - High-level message access helpers
    """

    # SQL keywords that indicate dangerous operations
    DANGEROUS_KEYWORDS: frozenset[str] = frozenset([
        "DROP", "DELETE", "INSERT", "UPDATE", "CREATE", "ALTER", "ATTACH"
    ])

    # Reserved SQL keywords that shouldn't be treated as table aliases
    RESERVED_ALIASES: frozenset[str] = frozenset([
        "where", "join", "left", "right", "inner", "outer",
        "on", "order", "group", "limit", "fts"
    ])

    def __init__(self, db_path: str) -> None:
        self.db_path = db_path
        self._conn: sqlite3.Connection | None = None
        self._retrieved_guids: list[str] = []

    def _get_connection(self) -> sqlite3.Connection:
        """Get read-only database connection."""
        if self._conn is None:
            self._conn = sqlite3.connect(
                f"file:{self.db_path}?mode=ro",
                uri=True,
                check_same_thread=False,
            )
            self._conn.row_factory = sqlite3.Row
        return self._conn

    def _strip_sql_literals_and_comments(self, sql: str) -> str:
        """Strip SQL string literals and comments for safer validation."""
        chars = list(sql)
        i = 0
        mode: str | None = None
        length = len(chars)

        while i < length:
            if mode is None:
                if sql.startswith("--", i):
                    chars[i] = " "
                    if i + 1 < length:
                        chars[i + 1] = " "
                    i += 2
                    mode = "line_comment"
                    continue
                if sql.startswith("/*", i):
                    chars[i] = " "
                    if i + 1 < length:
                        chars[i + 1] = " "
                    i += 2
                    mode = "block_comment"
                    continue

                char = sql[i]
                if char == "'":
                    chars[i] = " "
                    i += 1
                    mode = "single"
                    continue
                if char == '"':
                    chars[i] = " "
                    i += 1
                    mode = "double"
                    continue
                if char == "`":
                    chars[i] = " "
                    i += 1
                    mode = "backtick"
                    continue
                if char == "[":
                    chars[i] = " "
                    i += 1
                    mode = "bracket"
                    continue

                i += 1
                continue

            if mode == "line_comment":
                if sql[i] == "\n":
                    mode = None
                    i += 1
                else:
                    chars[i] = " "
                    i += 1
                continue

            if mode == "block_comment":
                if sql.startswith("*/", i):
                    chars[i] = " "
                    if i + 1 < length:
                        chars[i + 1] = " "
                    i += 2
                    mode = None
                else:
                    chars[i] = " "
                    i += 1
                continue

            if mode == "single":
                if sql[i] == "'":
                    if i + 1 < length and sql[i + 1] == "'":
                        chars[i] = " "
                        chars[i + 1] = " "
                        i += 2
                    else:
                        chars[i] = " "
                        i += 1
                        mode = None
                else:
                    chars[i] = " "
                    i += 1
                continue

            if mode == "double":
                if sql[i] == '"':
                    if i + 1 < length and sql[i + 1] == '"':
                        chars[i] = " "
                        chars[i + 1] = " "
                        i += 2
                    else:
                        chars[i] = " "
                        i += 1
                        mode = None
                else:
                    chars[i] = " "
                    i += 1
                continue

            if mode == "backtick":
                if sql[i] == "`":
                    chars[i] = " "
                    i += 1
                    mode = None
                else:
                    chars[i] = " "
                    i += 1
                continue

            if mode == "bracket":
                if sql[i] == "]":
                    chars[i] = " "
                    i += 1
                    mode = None
                else:
                    chars[i] = " "
                    i += 1
                continue

        return "".join(chars)

    def _validate_sql(self, sql: str) -> None:
        """Validate SQL is safe to execute.

        Raises:
            SandboxError: If SQL contains dangerous operations
        """
        sanitized = self._strip_sql_literals_and_comments(sql)
        sql_upper = sanitized.strip().upper()
        if not sql_upper.startswith("SELECT"):
            raise SandboxError("Only SELECT queries are allowed")

        for keyword in self.DANGEROUS_KEYWORDS:
            if re.search(rf"\b{keyword}\b", sql_upper):
                raise SandboxError(f"Dangerous SQL keyword not allowed: {keyword}")

    def _inject_guid_if_needed(self, sql: str) -> str:
        """Auto-inject guid column for tracking if querying messages table.

        Skips injection for DISTINCT, GROUP BY, and aggregate queries since
        adding guid would change their semantics.
        """
        sql_upper = sql.upper()
        sanitized = self._strip_sql_literals_and_comments(sql)
        sanitized_upper = sanitized.upper()

        if 'GUID' in sql_upper or 'MESSAGES' not in sql_upper:
            return sql

        # Skip injection for queries where adding guid changes semantics
        if re.search(r'\bDISTINCT\b', sanitized_upper):
            return sql
        if re.search(r'\bGROUP\s+BY\b', sanitized_upper):
            return sql
        if re.search(r'\bHAVING\b', sanitized_upper):
            return sql
        # Skip for aggregate-only queries (COUNT, SUM, AVG, MIN, MAX)
        if re.search(r'\b(COUNT|SUM|AVG|MIN|MAX)\s*\(', sanitized_upper):
            return sql

        match = re.match(r'(\s*SELECT\s+)(.*?)(\s+FROM\s+)', sql, re.IGNORECASE | re.DOTALL)
        if not match:
            return sql

        prefix, columns, from_clause = match.groups()
        rest = sql[match.end():]

        # Detect messages table alias in JOINs
        join_alias_match = re.search(
            r'JOIN\s+messages\s+(?:AS\s+)?(\w+)\s+ON',
            sql,
            re.IGNORECASE
        )
        if join_alias_match:
            alias = join_alias_match.group(1).lower()
            if alias not in self.RESERVED_ALIASES:
                return f"{prefix}{alias}.guid, {columns}{from_clause}{rest}"

        # Detect direct FROM messages alias
        from_alias_match = re.search(
            r'FROM\s+messages\s+(?:AS\s+)?(\w+)',
            sql,
            re.IGNORECASE
        )
        if from_alias_match:
            alias = from_alias_match.group(1).lower()
            if alias not in self.RESERVED_ALIASES:
                return f"{prefix}{alias}.guid, {columns}{from_clause}{rest}"

        # FTS-only query without JOIN
        if 'MESSAGES_FTS' in sql_upper and 'JOIN' not in sql_upper:
            return sql

        return f"{prefix}messages.guid, {columns}{from_clause}{rest}"

    def _track_guids(self, results: list[dict[str, Any]], columns: list[str]) -> None:
        """Track GUIDs from results for benchmark mode."""
        guid_key = next((k for k in columns if k.lower() == 'guid'), None)
        if guid_key:
            for row in results:
                guid = row.get(guid_key)
                if guid and guid not in self._retrieved_guids:
                    self._retrieved_guids.append(guid)

    def get_retrieved_guids(self) -> list[str]:
        """Return the list of GUIDs retrieved during query execution."""
        return self._retrieved_guids.copy()

    def query(self, sql: str) -> list[dict[str, Any]]:
        """Execute a read-only SQL query and return results as list of dicts."""
        self._validate_sql(sql)

        modified_sql = self._inject_guid_if_needed(sql)
        conn = self._get_connection()
        cursor = conn.execute(modified_sql)
        columns = [desc[0] for desc in cursor.description]
        results = [dict(zip(columns, row)) for row in cursor.fetchall()]

        self._track_guids(results, columns)
        return results

    def search(
        self,
        terms: str,
        limit: int = 1000,
        sender: str | None = None,
        chat: str | None = None,
        days: int | None = None,
    ) -> list[dict[str, Any]]:
        """Search messages using FTS with optional filters."""
        import time

        conn = self._get_connection()
        conditions = ["messages_fts MATCH ?"]
        params: list[Any] = [terms]

        if sender:
            conditions.append("m.sender LIKE ?")
            params.append(f"%{sender}%")

        if chat:
            conditions.append("m.chat_name LIKE ?")
            params.append(f"%{chat}%")

        if days is not None:
            cutoff = unix_seconds_to_cf_nanoseconds(time.time() - (days * 86400))
            conditions.append("m.timestamp >= ?")
            params.append(cutoff)

        params.append(limit)

        sql = f"""
            SELECT m.guid, m.text, m.sender, m.chat_name, m.chat_id,
                   m.timestamp, m.is_from_me, m.is_group
            FROM messages_fts
            JOIN messages m ON messages_fts.rowid = m.rowid
            WHERE {' AND '.join(conditions)}
            ORDER BY m.timestamp DESC
            LIMIT ?
        """

        cursor = conn.execute(sql, params)
        columns = [desc[0] for desc in cursor.description]
        results = [dict(zip(columns, row)) for row in cursor.fetchall()]

        self._track_guids(results, columns)
        return results

    def get_messages(self, limit: int | None = 10000) -> list[dict[str, Any]]:
        """Get messages in bulk for exploration and processing."""
        conn = self._get_connection()

        if limit is None:
            cursor = conn.execute("""
                SELECT guid, text, sender, chat_name, chat_id,
                       timestamp, is_from_me, is_group
                FROM messages
                ORDER BY timestamp DESC
            """)
        else:
            cursor = conn.execute("""
                SELECT guid, text, sender, chat_name, chat_id,
                       timestamp, is_from_me, is_group
                FROM messages
                ORDER BY timestamp DESC
                LIMIT ?
            """, (limit,))

        columns = [desc[0] for desc in cursor.description]
        results = [dict(zip(columns, row)) for row in cursor.fetchall()]

        self._track_guids(results, columns)
        return results

    def get_stats(self) -> dict[str, Any]:
        """Get high-level statistics about the message database."""
        conn = self._get_connection()

        cursor = conn.execute("SELECT COUNT(*) FROM messages")
        total_messages = cursor.fetchone()[0]

        cursor = conn.execute("SELECT COUNT(DISTINCT sender) FROM messages")
        unique_senders = cursor.fetchone()[0]

        cursor = conn.execute("SELECT COUNT(DISTINCT chat_id) FROM messages")
        unique_chats = cursor.fetchone()[0]

        cursor = conn.execute("SELECT MIN(timestamp), MAX(timestamp) FROM messages")
        row = cursor.fetchone()
        oldest_timestamp = row[0]
        newest_timestamp = row[1]

        cursor = conn.execute("""
            SELECT sender, COUNT(*) as count
            FROM messages
            GROUP BY sender
            ORDER BY count DESC
            LIMIT 10
        """)
        top_senders = [{'sender': r[0], 'count': r[1]} for r in cursor.fetchall()]

        cursor = conn.execute("""
            SELECT chat_name, chat_id, COUNT(*) as count
            FROM messages
            GROUP BY chat_id
            ORDER BY count DESC
            LIMIT 10
        """)
        top_chats = [
            {'chat_name': r[0] or r[1], 'chat_id': r[1], 'count': r[2]}
            for r in cursor.fetchall()
        ]

        return {
            'total_messages': total_messages,
            'unique_senders': unique_senders,
            'unique_chats': unique_chats,
            'oldest_timestamp': oldest_timestamp,
            'newest_timestamp': newest_timestamp,
            'top_senders': top_senders,
            'top_chats': top_chats,
        }

    def get_chats(self) -> list[dict[str, Any]]:
        """Get list of all chats with metadata."""
        conn = self._get_connection()

        cursor = conn.execute("""
            SELECT
                chat_id,
                chat_name,
                COUNT(*) as message_count,
                MAX(timestamp) as last_message,
                MAX(is_group) as is_group
            FROM messages
            GROUP BY chat_id
            ORDER BY last_message DESC
        """)

        return [
            {
                'chat_id': r[0],
                'chat_name': r[1] or r[0],
                'message_count': r[2],
                'last_message_timestamp': r[3],
                'is_group': bool(r[4]),
            }
            for r in cursor.fetchall()
        ]

    def get_recent(self, days: int = 30, limit: int = 10000) -> list[dict[str, Any]]:
        """Get recent messages within a time window."""
        import time
        cutoff = unix_seconds_to_cf_nanoseconds(time.time() - (days * 86400))

        conn = self._get_connection()
        cursor = conn.execute("""
            SELECT guid, text, sender, chat_name, chat_id,
                   timestamp, is_from_me, is_group
            FROM messages
            WHERE timestamp >= ?
            ORDER BY timestamp DESC
            LIMIT ?
        """, (cutoff, limit))

        columns = [desc[0] for desc in cursor.description]
        results = [dict(zip(columns, row)) for row in cursor.fetchall()]

        self._track_guids(results, columns)
        return results

    def sample(self, n: int = 1000) -> list[dict[str, Any]]:
        """Get a random sample of messages."""
        conn = self._get_connection()

        cursor = conn.execute("""
            SELECT guid, text, sender, chat_name, chat_id,
                   timestamp, is_from_me, is_group
            FROM messages
            ORDER BY RANDOM()
            LIMIT ?
        """, (n,))

        columns = [desc[0] for desc in cursor.description]
        results = [dict(zip(columns, row)) for row in cursor.fetchall()]

        self._track_guids(results, columns)
        return results

    def get_chat_messages(self, chat_id: str, limit: int = 1000) -> list[dict[str, Any]]:
        """Get messages from a specific chat."""
        conn = self._get_connection()

        cursor = conn.execute("""
            SELECT guid, text, sender, chat_name, chat_id,
                   timestamp, is_from_me, is_group
            FROM messages
            WHERE chat_id = ?
            ORDER BY timestamp DESC
            LIMIT ?
        """, (chat_id, limit))

        columns = [desc[0] for desc in cursor.description]
        results = [dict(zip(columns, row)) for row in cursor.fetchall()]

        self._track_guids(results, columns)
        return results

    def close(self) -> None:
        """Close the database connection."""
        if self._conn:
            self._conn.close()
            self._conn = None


class LlmOrchestrator:
    """Handles recursive LLM sub-calls within sandbox execution.

    Responsibilities:
    - Execute llm_query() sub-calls with depth/call limits
    - Batch processing of multiple queries
    - Context truncation for token limits
    """

    MAX_CONTEXT_CHARS: int = 8000
    SUB_CALL_MAX_TOKENS: int = 2048

    def __init__(self, context: ExecutionContext) -> None:
        self._context = context
        from .modelrelay import resolve_model
        self._model = resolve_model(None)

    def llm_query(self, prompt: str, context_data: str = "") -> str:
        """Execute a recursive LLM sub-query.

        Raises:
            SandboxError: If max depth or max calls exceeded
        """
        from .modelrelay import chat_completion

        if self._context.depth >= self._context.max_depth:
            raise SandboxError(
                f"Max recursion depth ({self._context.max_depth}) exceeded. "
                "Consider breaking your query into smaller parts."
            )

        if self._context.calls_made >= self._context.max_calls:
            raise SandboxError(
                f"Max sub-calls ({self._context.max_calls}) exceeded. "
                "Consider reducing the number of chunks or increasing --max-calls."
            )

        self._context.calls_made += 1
        self._context.depth += 1
        try:
            truncated_context = context_data[:self.MAX_CONTEXT_CHARS] if context_data else ""

            messages = [
                {"role": "system", "content": SUB_QUERY_PROMPT},
                {"role": "user", "content": f"{prompt}\n\nContext:\n{truncated_context}" if truncated_context else prompt},
            ]

            if self._context.verbose:
                print(f"\n  [Sub-call {self._context.calls_made}] {prompt[:80]}...", file=sys.stderr)

            response, usage = chat_completion(
                messages, model=self._model, max_tokens=self.SUB_CALL_MAX_TOKENS, return_usage=True
            )

            # Track token usage
            self._context.total_tokens += usage.total_tokens

            if self._context.verbose:
                print(f"  [Sub-call {self._context.calls_made}] Response: {response[:100]}... ({usage.input_tokens}+{usage.output_tokens} tokens)", file=sys.stderr)

            return response
        finally:
            self._context.depth -= 1

    def batch_llm_query(self, prompts: list[str], contexts: list[str] | None = None) -> list[str]:
        """Execute multiple llm_query calls sequentially."""
        contexts = contexts or [""] * len(prompts)
        if len(contexts) != len(prompts):
            raise SandboxError("prompts and contexts must have the same length")

        return [self.llm_query(p, c) for p, c in zip(prompts, contexts)]

    @staticmethod
    def chunk_data(data: list[Any], chunk_size: int) -> list[list[Any]]:
        """Split data into chunks for parallel processing."""
        return [data[i:i + chunk_size] for i in range(0, len(data), chunk_size)]


class RestrictedExecutor:
    """Handles sandboxed Python code execution.

    Responsibilities:
    - Compile code with RestrictedPython
    - Build safe execution namespace
    - Execute code with captured stdout
    """

    # Safe builtins allowed in sandbox
    SAFE_BUILTINS: dict[str, Any] = {
        "True": True,
        "False": False,
        "None": None,
        "abs": abs,
        "all": all,
        "any": any,
        "bool": bool,
        "dict": dict,
        "enumerate": enumerate,
        "filter": filter,
        "float": float,
        "frozenset": frozenset,
        "int": int,
        "isinstance": isinstance,
        "len": len,
        "list": list,
        "map": map,
        "max": max,
        "min": min,
        "range": range,
        "reversed": reversed,
        "round": round,
        "set": set,
        "sorted": sorted,
        "str": str,
        "sum": sum,
        "tuple": tuple,
        "zip": zip,
        "hasattr": hasattr,
    }

    def __init__(
        self,
        db: MessageDatabase,
        llm: LlmOrchestrator | None = None,
    ) -> None:
        self._db = db
        self._llm = llm

    def _build_globals(self, output: StringIO) -> dict[str, Any]:
        """Build the restricted globals namespace."""
        def safe_print(*args: Any, **kwargs: Any) -> None:
            kwargs["file"] = output
            print(*args, **kwargs)

        def guarded_write(obj: Any) -> Any:
            if isinstance(obj, (dict, list)):
                return obj
            raise SandboxError(f"Cannot assign items on {type(obj).__name__}")

        restricted_globals = safe_globals.copy()

        # RestrictedPython guards
        restricted_globals["_getiter_"] = default_guarded_getiter
        restricted_globals["_getitem_"] = default_guarded_getitem
        restricted_globals["_iter_unpack_sequence_"] = guarded_iter_unpack_sequence
        restricted_globals["_getattr_"] = safer_getattr
        restricted_globals["_write_"] = guarded_write

        # Safe builtins with print
        builtins = dict(self.SAFE_BUILTINS)
        builtins["print"] = safe_print
        builtins["getattr"] = safer_getattr
        restricted_globals["__builtins__"] = builtins

        # Safe modules
        restricted_globals["re"] = re
        restricted_globals["json"] = json
        restricted_globals["Counter"] = Counter
        restricted_globals["defaultdict"] = defaultdict
        restricted_globals["datetime"] = datetime
        restricted_globals["timedelta"] = timedelta
        restricted_globals["cf_to_datetime"] = cf_to_datetime

        # Database functions
        restricted_globals["query"] = self._db.query
        restricted_globals["get_messages"] = self._db.get_messages
        restricted_globals["get_stats"] = self._db.get_stats
        restricted_globals["get_chats"] = self._db.get_chats
        restricted_globals["search"] = self._db.search
        restricted_globals["get_recent"] = self._db.get_recent
        restricted_globals["get_chat_messages"] = self._db.get_chat_messages
        restricted_globals["sample"] = self._db.sample

        # LLM functions (if orchestrator provided)
        if self._llm:
            restricted_globals["llm_query"] = self._llm.llm_query
            restricted_globals["chunk_data"] = self._llm.chunk_data
            restricted_globals["batch_llm_query"] = self._llm.batch_llm_query

        # Print handler for RestrictedPython
        class PrintHandler:
            def __init__(self, _getattr_: Any = None) -> None:
                self._getattr_ = _getattr_

            def _call_print(self, *args: Any, **kwargs: Any) -> None:
                text = " ".join(str(arg) for arg in args)
                end = kwargs.get("end", "\n")
                output.write(text + end)

            def __call__(self) -> str:
                return ""

        restricted_globals["_print_"] = PrintHandler

        return restricted_globals

    def execute(self, code: str) -> str:
        """Execute Python code in a restricted environment.

        Returns the captured stdout output.

        Raises:
            SandboxError: If compilation or execution fails
        """
        if "__import__" in code or "import " in code:
            raise SandboxError("Imports are not allowed in this environment. Use provided helpers instead.")

        output = StringIO()
        restricted_globals = self._build_globals(output)

        # Compile with RestrictedPython
        try:
            compiled_code = compile_restricted(
                code,
                filename="<rlm>",
                mode="exec",
            )
        except SyntaxError as e:
            raise SandboxError(f"Syntax error in code: {e}")

        if hasattr(compiled_code, 'errors') and compiled_code.errors:
            raise SandboxError(f"Compilation errors: {compiled_code.errors}")

        if hasattr(compiled_code, 'code'):
            compiled_code = compiled_code.code

        if compiled_code is None:
            raise SandboxError("Compilation failed: no code generated")

        # Execute the RestrictedPython-compiled bytecode
        # This is safe: RestrictedPython blocks dangerous operations at compile time
        try:
            exec_locals: dict[str, Any] = {}
            # Get the real exec from builtins (RestrictedPython usage pattern)
            builtins_module = __builtins__ if isinstance(__builtins__, dict) else vars(__builtins__)
            builtin_exec = builtins_module.get("exec")
            if builtin_exec is None:
                raise SandboxError("Could not find exec builtin")
            builtin_exec(compiled_code, restricted_globals, exec_locals)
        except SandboxError:
            raise
        except Exception as e:
            raise SandboxError(f"Execution error: {type(e).__name__}: {e}")

        return output.getvalue()


class Sandbox:
    """Facade for sandboxed Python execution with database access and LLM calls.

    Composes MessageDatabase, LlmOrchestrator, and RestrictedExecutor.
    """

    def __init__(self, db_path: str, context: ExecutionContext | None = None) -> None:
        self._db = MessageDatabase(db_path)
        self._llm = LlmOrchestrator(context) if context else None
        self._executor = RestrictedExecutor(self._db, self._llm)
        self._context = context

    def execute(self, code: str) -> str:
        """Execute Python code in the sandbox."""
        return self._executor.execute(code)

    def get_retrieved_guids(self) -> list[str]:
        """Return the list of GUIDs retrieved during query execution."""
        return self._db.get_retrieved_guids()

    def close(self) -> None:
        """Close the database connection."""
        self._db.close()

    # Expose database methods directly for convenience
    def query(self, sql: str) -> list[dict[str, Any]]:
        """Execute a read-only SQL query."""
        return self._db.query(sql)

    def search(
        self,
        terms: str,
        limit: int = 1000,
        sender: str | None = None,
        chat: str | None = None,
        days: int | None = None,
    ) -> list[dict[str, Any]]:
        """Search messages using FTS."""
        return self._db.search(terms, limit, sender, chat, days)

    def get_messages(self, limit: int | None = 10000) -> list[dict[str, Any]]:
        """Get messages in bulk."""
        return self._db.get_messages(limit)

    def get_stats(self) -> dict[str, Any]:
        """Get database statistics."""
        return self._db.get_stats()

    def get_chats(self) -> list[dict[str, Any]]:
        """Get list of all chats."""
        return self._db.get_chats()

    def get_recent(self, days: int = 30, limit: int = 10000) -> list[dict[str, Any]]:
        """Get recent messages."""
        return self._db.get_recent(days, limit)

    def sample(self, n: int = 1000) -> list[dict[str, Any]]:
        """Get random sample of messages."""
        return self._db.sample(n)

    def get_chat_messages(self, chat_id: str, limit: int = 1000) -> list[dict[str, Any]]:
        """Get messages from a specific chat."""
        return self._db.get_chat_messages(chat_id, limit)

    # Expose LLM methods if available
    def llm_query(self, prompt: str, context_data: str = "") -> str:
        """Execute a recursive LLM sub-query."""
        if self._llm is None:
            raise SandboxError("llm_query requires an ExecutionContext")
        return self._llm.llm_query(prompt, context_data)

    @staticmethod
    def chunk_data(data: list[Any], chunk_size: int) -> list[list[Any]]:
        """Split data into chunks."""
        return LlmOrchestrator.chunk_data(data, chunk_size)

    def batch_llm_query(self, prompts: list[str], contexts: list[str] | None = None) -> list[str]:
        """Execute multiple llm_query calls."""
        if self._llm is None:
            raise SandboxError("batch_llm_query requires an ExecutionContext")
        return self._llm.batch_llm_query(prompts, contexts)
