"""External tool result store — append-only JSONL with filelock.

Stores large tool results externally so they don't bloat the conversation context.
Results are replaced inline with compact references (ID + summary), and the agent
can retrieve full results on demand via MCP tools.
"""

from __future__ import annotations

import json
import logging
import re
from dataclasses import asdict, dataclass
from datetime import UTC, datetime, timedelta
from enum import Enum
from pathlib import Path
from typing import Any

from ulid import ulid

log = logging.getLogger("cadecoder")

# ---------------------------------------------------------------------------
# Thresholds (bytes, measured via sys.getsizeof for O(1) checks)
# ---------------------------------------------------------------------------

# Per-result: externalize if a single result exceeds this size in bytes (~4k tokens)
PER_RESULT_BYTE_THRESHOLD = 16_000

# Aggregate per-turn: if total bytes across all results in a turn exceed this,
# externalize everything above TINY threshold
AGGREGATE_BYTE_THRESHOLD = 64_000

# Results smaller than this are always kept inline (even in aggregate mode)
TINY_RESULT_BYTES = 1_024

# Tools whose results must NEVER be externalized — doing so creates infinite
# retrieval loops (agent retrieves → result externalized → agent retrieves again → ...)
NEVER_EXTERNALIZE_TOOLS: frozenset[str] = frozenset(
    {
        "Local_RetrieveToolResult",
        "Local_SearchRecentContext",
        "Local_ToolSchema",
        "Local_PinContext",
    }
)

# Tools whose results should be compressed to extract only essential info.
# SelectTools returns full tool schemas that bloat context — we keep only
# tool names, brief descriptions, and the query_id.
COMPRESS_SCHEMA_TOOLS: frozenset[str] = frozenset(
    {
        "Arcade_SelectTools",
    }
)

# ---------------------------------------------------------------------------
# Data structures
# ---------------------------------------------------------------------------


@dataclass
class ToolResultRecord:
    """One JSONL line per externalized result."""

    record_id: str
    tool_call_id: str
    tool_name: str
    thread_id: str
    timestamp: str  # ISO 8601
    content: str
    content_chars: int
    content_token_estimate: int
    summary: str
    status: str  # "success" | "error"
    turn_sequence: int  # index within the turn
    name: str = ""  # human-readable label for UserContext records

    def to_json_line(self) -> str:
        """Serialize to a single JSON line (no embedded newlines)."""
        return json.dumps(asdict(self), ensure_ascii=False, separators=(",", ":"))

    @classmethod
    def from_json_line(cls, line: str) -> ToolResultRecord:
        """Deserialize from a JSON line."""
        data = json.loads(line)
        data.setdefault("name", "")  # backward compat
        return cls(**data)


@dataclass
class ToolResultRef:
    """Compact reference that replaces full content in messages."""

    record_id: str
    tool_name: str
    content_chars: int
    status: str
    summary: str

    def to_inline_text(self) -> str:
        """Produce the compact string that goes into the message content."""
        return (
            f"[EXTERNALIZED_RESULT id={self.record_id} "
            f"tool={self.tool_name} chars={self.content_chars} "
            f"status={self.status}]\n"
            f"Summary: {self.summary}\n"
            f'Use Local_RetrieveToolResult(record_id="{self.record_id}") '
            f"to read the full result."
        )

    @staticmethod
    def is_externalized(content: str) -> bool:
        """Check whether a tool result content string is an externalized reference."""
        return content.startswith("[EXTERNALIZED_RESULT ")

    @staticmethod
    def extract_record_id(content: str) -> str | None:
        """Extract the record_id from an externalized reference string."""
        m = re.search(r"id=(\S+)", content)
        return m.group(1) if m else None


class RetrieveMode(str, Enum):
    """Retrieval strategy for externalized tool results."""

    FULL = "full"
    SEARCH = "search"
    SUMMARIZE = "summarize"
    LINES = "lines"
    JSON = "json"
    LIST = "list"


# ---------------------------------------------------------------------------
# Summary generation (adapted from ui/display.py for plain text)
# ---------------------------------------------------------------------------


def _generate_summary(content: str, tool_name: str, status: str, max_len: int = 120) -> str:
    """Generate a brief plain-text summary for the LLM."""
    if not content:
        return f"(empty result from {tool_name})"
    if status == "error":
        # Extract HTTP status if present
        for code in ("404", "401", "403", "500", "502", "503"):
            if code in content[:200]:
                return f"Error {code} from {tool_name}"
        preview = content.replace("\n", " ").strip()[:max_len]
        return f"Error: {preview}"

    # Try JSON
    try:
        data = json.loads(content)
        return _summarize_json(data, tool_name, max_len)
    except (json.JSONDecodeError, TypeError):
        pass

    # Plain text fallback
    lines = content.strip().splitlines()
    first_line = lines[0].strip()[:max_len] if lines else ""
    if len(lines) > 1:
        return f"{first_line} (+{len(lines) - 1} more lines)"
    return first_line if first_line else "(empty result)"


def _summarize_json(data: Any, tool_name: str, max_len: int) -> str:
    """Summarize a parsed JSON structure, extracting meaningful fields."""
    if isinstance(data, dict):
        preview = _extract_preview_fields(data)
        if preview:
            return preview[:max_len]
        keys = list(data.keys())
        if len(keys) <= 5:
            return f"Object with keys: {', '.join(keys)}"
        return f"Object with {len(keys)} keys: {', '.join(keys[:5])}, ..."
    if isinstance(data, list):
        if not data:
            return "Empty list"
        if isinstance(data[0], dict):
            # Summarize first few items to give the LLM useful context
            item_previews = []
            for item in data[:5]:
                preview = _extract_preview_fields(item)
                if preview:
                    item_previews.append(preview)
            if item_previews:
                suffix = f" (+{len(data) - 5} more)" if len(data) > 5 else ""
                return f"{len(data)} items: " + "; ".join(item_previews) + suffix
            return f"{len(data)} items (objects)"
        return f"{len(data)} items"
    return str(data)[:max_len]


# Fields to extract for human-readable previews, in priority order.
# Each group is tried until one produces output.
_PREVIEW_FIELD_GROUPS: list[list[str]] = [
    # Email-like
    ["subject", "from", "date"],
    ["subject", "sender", "date"],
    ["subject", "from_email", "received_at"],
    # Message-like
    ["title", "author", "created_at"],
    ["name", "description"],
    ["title", "body"],
    # Generic
    ["id", "name", "status"],
    ["id", "type", "state"],
]


def _extract_preview_fields(obj: dict, max_val_len: int = 60) -> str:
    """Extract a human-readable preview from a dict by looking for common fields."""
    for field_group in _PREVIEW_FIELD_GROUPS:
        parts = []
        for field in field_group:
            # Case-insensitive field lookup
            val = _get_field_ci(obj, field)
            if val is not None:
                val_str = str(val).replace("\n", " ").strip()
                if len(val_str) > max_val_len:
                    val_str = val_str[: max_val_len - 3] + "..."
                parts.append(f"{field}={val_str}")
        if len(parts) >= 2:  # Need at least 2 fields to be useful
            return ", ".join(parts)
    return ""


def _get_field_ci(obj: dict, field: str) -> Any:
    """Case-insensitive dict field lookup."""
    if field in obj:
        return obj[field]
    field_lower = field.lower()
    for key, val in obj.items():
        if key.lower() == field_lower:
            return val
    return None


# ---------------------------------------------------------------------------
# SelectTools schema compression
# ---------------------------------------------------------------------------


def compress_select_tools_result(content: str) -> str:
    """Compress Arcade_SelectTools result to keep only essential tool info.

    SelectTools returns full tool schemas (with parameter definitions) that
    can be very large.  We extract tool names, brief descriptions, and the
    query_id, discarding the verbose parameter schemas.

    Returns a compact plain-text summary suitable for inline context.
    """
    try:
        data = json.loads(content)
    except (json.JSONDecodeError, TypeError):
        # Not JSON — return a truncated version
        return content[:500] + ("..." if len(content) > 500 else "")

    return _compress_schema_data(data)


def _compress_schema_data(data: Any) -> str:
    """Extract compact tool info from parsed SelectTools JSON."""
    tools: list[dict[str, Any]] = []
    query_id: str | None = None

    # Handle various response shapes
    if isinstance(data, dict):
        # Look for query_id at top level
        query_id = data.get("query_id") or data.get("queryId")

        # Look for tools list
        tool_list = data.get("tools") or data.get("results") or data.get("definitions")
        if isinstance(tool_list, list):
            tools = tool_list
        elif "function" in data or "name" in data:
            # Single tool
            tools = [data]
    elif isinstance(data, list):
        tools = data

    if not tools:
        # Fallback: just summarize keys/size
        summary = json.dumps(data, indent=None, ensure_ascii=False)
        return summary[:500] + ("..." if len(summary) > 500 else "")

    # Build compact summary
    lines = []
    if query_id:
        lines.append(f"query_id: {query_id}")
    lines.append(f"Selected {len(tools)} tools:")

    for tool_obj in tools:
        name = _extract_tool_name(tool_obj)
        desc = _extract_tool_description(tool_obj)

        if desc:
            # First sentence only
            first_sentence = desc.split(". ")[0].rstrip(".")
            lines.append(f"  - {name}: {first_sentence}")
        else:
            lines.append(f"  - {name}")

    lines.append("")
    lines.append(
        "Full schemas stored externally. Use Arcade_UseTool with the tool name and query_id above."
    )

    return "\n".join(lines)


def _strip_version_suffix(name: str) -> str:
    """Remove Arcade toolkit version suffix from tool names.

    Arcade may return versioned names like ``GoogleShopping_SearchProducts@3_1_2``.
    ``Arcade_UseTool`` expects the bare name without the ``@version`` suffix.
    """
    at_idx = name.find("@")
    return name[:at_idx] if at_idx > 0 else name


def _extract_tool_name(tool_obj: dict[str, Any]) -> str:
    """Extract tool name from various schema formats."""
    # OpenAI format: {"function": {"name": "..."}}
    func = tool_obj.get("function", {})
    if isinstance(func, dict) and func.get("name"):
        return _strip_version_suffix(func["name"])
    # Flat format: {"name": "..."}
    return _strip_version_suffix(tool_obj.get("name", "unknown"))


def _extract_tool_description(tool_obj: dict[str, Any]) -> str:
    """Extract tool description from various schema formats."""
    # OpenAI format
    func = tool_obj.get("function", {})
    if isinstance(func, dict) and func.get("description"):
        return func["description"]
    # Flat format
    return tool_obj.get("description", "")


# ---------------------------------------------------------------------------
# Retrieval strategy helpers
# ---------------------------------------------------------------------------


def search_content(
    content: str,
    query: str,
    context_lines: int = 2,
    max_matches: int = 20,
    max_result_chars: int = 30_000,
) -> dict[str, Any]:
    """Search within externalized content for a pattern.

    Supports regex patterns via :mod:`re`.  Falls back to literal substring
    match (case-insensitive) if the pattern is not valid regex.

    Returns matching lines with 1-indexed line numbers and surrounding
    context.
    """
    lines = content.splitlines()

    # Try compiling as regex; fall back to escaped literal
    try:
        pattern = re.compile(query, re.IGNORECASE)
    except re.error:
        pattern = re.compile(re.escape(query), re.IGNORECASE)

    # Find matching line indices
    match_indices: list[int] = []
    for i, line in enumerate(lines):
        if pattern.search(line):
            match_indices.append(i)
            if len(match_indices) >= max_matches:
                break

    if not match_indices:
        return {
            "matches": [],
            "total_matches": 0,
            "total_lines": len(lines),
            "truncated": False,
        }

    # Build context ranges and merge overlaps
    ranges: list[tuple[int, int]] = []
    for idx in match_indices:
        start = max(0, idx - context_lines)
        end = min(len(lines), idx + context_lines + 1)
        if ranges and start <= ranges[-1][1]:
            ranges[-1] = (ranges[-1][0], end)
        else:
            ranges.append((start, end))

    # Collect results, tracking output size
    matches: list[dict[str, Any]] = []
    total_chars = 0
    truncated = False

    for rng_start, rng_end in ranges:
        group: list[dict[str, Any]] = []
        for i in range(rng_start, rng_end):
            entry = {
                "line": i + 1,  # 1-indexed
                "text": lines[i],
                "is_match": i in match_indices,
            }
            total_chars += len(lines[i]) + 40  # overhead estimate
            if total_chars > max_result_chars:
                truncated = True
                break
            group.append(entry)
        if group:
            matches.append({"context": group})
        if truncated:
            break

    return {
        "matches": matches,
        "total_matches": len(match_indices),
        "total_lines": len(lines),
        "truncated": truncated,
    }


def get_content_lines(
    content: str,
    start_line: int,
    end_line: int,
    max_result_chars: int = 30_000,
) -> dict[str, Any]:
    """Return a specific line range from externalized content.

    Lines are 1-indexed inclusive.  Out-of-range bounds are clamped.
    """
    lines = content.splitlines()
    total = len(lines)

    # Clamp bounds
    start = max(1, min(start_line, total))
    end = max(start, min(end_line, total))

    selected = lines[start - 1 : end]
    text = "\n".join(selected)

    truncated = False
    if len(text) > max_result_chars:
        text = text[:max_result_chars]
        truncated = True

    return {
        "lines": text,
        "start_line": start,
        "end_line": end,
        "total_lines": total,
        "returned_chars": len(text),
        "truncated": truncated,
    }


def summarize_content(
    content: str,
    tool_name: str,
    status: str,
    max_result_chars: int = 8_000,
) -> dict[str, Any]:
    """Generate a structural summary of externalized content.

    JSON arrays:  per-item preview using ``_extract_preview_fields``.
    JSON objects: top-level key overview.
    Plain text:   paragraph-based chunking with first-line summaries.
    """
    total_chars = len(content)

    # Try JSON first
    try:
        data = json.loads(content)
        if isinstance(data, list):
            return _summarize_array(data, total_chars, max_result_chars)
        if isinstance(data, dict):
            return _summarize_object(data, total_chars, max_result_chars)
    except (json.JSONDecodeError, TypeError):
        pass

    # Plain text
    return _summarize_text(content, total_chars, max_result_chars)


def _summarize_array(data: list[Any], total_chars: int, max_chars: int) -> dict[str, Any]:
    """Summarize a JSON array."""
    parts: list[str] = [f"JSON array with {len(data)} items:"]
    chars = len(parts[0])

    for i, item in enumerate(data):
        if isinstance(item, dict):
            preview = _extract_preview_fields(item)
            if preview:
                line = f"  {i + 1}. {preview}"
            else:
                keys = list(item.keys())[:5]
                line = f"  {i + 1}. {{{', '.join(keys)}}}"
        else:
            line = f"  {i + 1}. {str(item)[:100]}"

        chars += len(line) + 1
        if chars > max_chars:
            parts.append(f"  ... and {len(data) - i} more items")
            break
        parts.append(line)

    return {
        "summary": "\n".join(parts),
        "content_type": "json_array",
        "total_chars": total_chars,
        "item_count": len(data),
    }


def _summarize_object(data: dict[str, Any], total_chars: int, max_chars: int) -> dict[str, Any]:
    """Summarize a JSON object."""
    parts: list[str] = [f"JSON object with {len(data)} keys:"]

    for key, val in data.items():
        if isinstance(val, list):
            desc = f"array[{len(val)}]"
        elif isinstance(val, dict):
            desc = f"object{{{', '.join(list(val.keys())[:3])}}}"
        elif isinstance(val, str):
            desc = f'"{val[:60]}{"..." if len(val) > 60 else ""}"'
        else:
            desc = str(val)[:60]
        parts.append(f"  {key}: {desc}")

    summary = "\n".join(parts)
    if len(summary) > max_chars:
        summary = summary[:max_chars] + "\n  ..."

    return {
        "summary": summary,
        "content_type": "json_object",
        "total_chars": total_chars,
    }


def _summarize_text(content: str, total_chars: int, max_chars: int) -> dict[str, Any]:
    """Summarize plain text by paragraphs."""
    # Split on double-newline or section separators
    chunks = re.split(r"\n\n+|\n---+\n", content)
    chunks = [c.strip() for c in chunks if c.strip()]

    parts: list[str] = [f"Text with {len(chunks)} sections, {total_chars:,} chars:"]
    chars = len(parts[0])

    for i, chunk in enumerate(chunks):
        chunk_lines = chunk.splitlines()
        first_line = chunk_lines[0][:120]
        line_info = f" ({len(chunk_lines)} lines)" if len(chunk_lines) > 1 else ""
        entry = f"  {i + 1}. {first_line}{line_info}"
        chars += len(entry) + 1
        if chars > max_chars:
            parts.append(f"  ... and {len(chunks) - i} more sections")
            break
        parts.append(entry)

    return {
        "summary": "\n".join(parts),
        "content_type": "text",
        "total_chars": total_chars,
        "section_count": len(chunks),
    }


# ---------------------------------------------------------------------------
# JSON navigation helpers
# ---------------------------------------------------------------------------


def parse_json_path(path: str) -> list[str | int]:
    """Parse a dot-bracket JSON path into segments.

    Examples::

        ""                 -> []
        "emails"           -> ["emails"]
        "emails[0]"        -> ["emails", 0]
        "emails[0].subject"-> ["emails", 0, "subject"]
        "data.users[2].name" -> ["data", "users", 2, "name"]
    """
    if not path or not path.strip():
        return []

    segments: list[str | int] = []
    current = ""
    i = 0
    while i < len(path):
        ch = path[i]
        if ch == ".":
            if current:
                segments.append(current)
                current = ""
        elif ch == "[":
            if current:
                segments.append(current)
                current = ""
            j = path.find("]", i + 1)
            if j == -1:
                raise ValueError(f"Unclosed bracket at position {i}")
            index_str = path[i + 1 : j]
            try:
                segments.append(int(index_str))
            except ValueError:
                segments.append(index_str.strip("'\""))
            i = j
        else:
            current += ch
        i += 1

    if current:
        segments.append(current)
    return segments


def _format_path(traversed: list[str]) -> str:
    """Join traversed segments into a display path."""
    return "".join(traversed) if traversed else "(root)"


def navigate_json_path(data: Any, segments: list[str | int]) -> tuple[Any, str | None]:
    """Navigate into a parsed JSON value following path segments.

    Returns ``(value, None)`` on success or ``(None, error_message)`` on failure.
    """
    current = data
    traversed: list[str] = []

    for seg in segments:
        if isinstance(seg, int):
            if not isinstance(current, list):
                return None, (
                    f"Expected array at '{_format_path(traversed)}', got {type(current).__name__}"
                )
            if seg < 0 or seg >= len(current):
                return None, (
                    f"Index {seg} out of range at '{_format_path(traversed)}' "
                    f"(array length: {len(current)})"
                )
            current = current[seg]
            traversed.append(f"[{seg}]")
        else:
            if not isinstance(current, dict):
                return None, (
                    f"Expected object at '{_format_path(traversed)}', got {type(current).__name__}"
                )
            if seg not in current:
                available = list(current.keys())[:10]
                suffix = f" (+{len(current) - 10} more)" if len(current) > 10 else ""
                return None, (
                    f"Key '{seg}' not found at '{_format_path(traversed)}'. "
                    f"Available keys: {available}{suffix}"
                )
            current = current[seg]
            traversed.append(f".{seg}" if traversed else seg)

    return current, None


def _describe_value(value: Any) -> Any:
    """Produce a short type annotation for a JSON value.

    Returns the original value unchanged for scalars and short strings.
    """
    if isinstance(value, dict):
        keys = list(value.keys())
        if len(keys) <= 5:
            return f"object{{{', '.join(keys)}}}"
        return f"object({len(keys)} keys)"
    if isinstance(value, list):
        return f"array[{len(value)}]"
    if isinstance(value, str):
        if len(value) <= 50:
            return value
        return f"string({len(value)} chars)"
    if isinstance(value, bool):
        return value
    if isinstance(value, int | float):
        return value
    if value is None:
        return None
    return str(value)[:60]


def truncate_at_depth(
    value: Any,
    depth: int | None,
    max_chars: int = 30_000,
) -> tuple[Any, bool]:
    """Truncate a JSON value at the specified depth.

    Args:
        value: Parsed JSON value.
        depth: How deep to render. 0 = annotation only, None = full.
        max_chars: Safety limit on serialized output size.

    Returns:
        ``(truncated_value, was_truncated)``
    """
    if depth is not None and depth < 0:
        depth = 0

    truncated = False

    def _truncate(val: Any, remaining: int | None) -> Any:
        nonlocal truncated

        if remaining is None:
            return val

        if remaining == 0:
            desc = _describe_value(val)
            if isinstance(val, dict | list) or (isinstance(val, str) and len(val) > 50):
                truncated = True
            return desc

        if isinstance(val, dict):
            return {k: _truncate(v, remaining - 1) for k, v in val.items()}

        if isinstance(val, list):
            return [_truncate(item, remaining - 1) for item in val]

        return val

    result = _truncate(value, depth)

    # Safety: if serialized result is too large, re-truncate shallower
    try:
        serialized = json.dumps(result, ensure_ascii=False, default=str)
        if len(serialized) > max_chars:
            truncated = True
            if depth is None or depth > 1:
                return truncate_at_depth(value, min(depth or 3, 2), max_chars)
    except (TypeError, ValueError):
        pass

    return result, truncated


def _json_type_name(value: Any) -> str:
    """Return the JSON type name for a Python value."""
    if isinstance(value, dict):
        return "object"
    if isinstance(value, list):
        return "array"
    if isinstance(value, str):
        return "string"
    if isinstance(value, bool):
        return "boolean"
    if isinstance(value, int | float):
        return "number"
    if value is None:
        return "null"
    return "unknown"


def _try_unwrap_mcp_content(data: Any) -> Any:
    """Unwrap the standard MCP ``[{"type":"text","text":"..."}]`` wrapper.

    If *data* matches the pattern, parse and return the inner JSON.
    Otherwise return *data* unchanged.
    """
    if (
        isinstance(data, list)
        and len(data) == 1
        and isinstance(data[0], dict)
        and data[0].get("type") == "text"
        and isinstance(data[0].get("text"), str)
    ):
        try:
            return json.loads(data[0]["text"])
        except (json.JSONDecodeError, TypeError):
            pass
    return data


def navigate_json(
    content: str,
    json_path: str = "",
    depth: int | None = None,
    max_chars: int = 30_000,
) -> dict[str, Any]:
    """Navigate and explore JSON content at a path and depth.

    Retrieval helper for ``mode='json'``.
    """
    try:
        data = json.loads(content)
    except (json.JSONDecodeError, TypeError) as e:
        return {
            "error": f"Content is not valid JSON: {e}",
            "hint": "Use mode='search' or mode='lines' for non-JSON content.",
        }

    # Auto-unwrap MCP [{"type":"text","text":"..."}] wrapper
    data = _try_unwrap_mcp_content(data)

    try:
        segments = parse_json_path(json_path)
    except ValueError as e:
        return {"error": f"Invalid json_path: {e}"}

    value, nav_error = navigate_json_path(data, segments)
    if nav_error is not None:
        return {"error": nav_error, "json_path": json_path}

    value_type = _json_type_name(value)
    result: dict[str, Any] = {
        "json_path": json_path or "(root)",
        "depth": depth,
        "value_type": value_type,
    }

    if isinstance(value, dict):
        result["child_count"] = len(value)
        result["available_keys"] = list(value.keys())
    elif isinstance(value, list):
        result["child_count"] = len(value)
        result["array_length"] = len(value)

    truncated_value, was_truncated = truncate_at_depth(value, depth, max_chars)
    result["truncated_at_depth"] = was_truncated

    if isinstance(truncated_value, dict | list):
        try:
            result["content"] = json.dumps(
                truncated_value, indent=2, ensure_ascii=False, default=str
            )
        except (TypeError, ValueError):
            result["content"] = str(truncated_value)
    else:
        result["content"] = truncated_value

    return result


# ---------------------------------------------------------------------------
# ToolResultStore
# ---------------------------------------------------------------------------


class ToolResultStore:
    """Append-only JSONL store for externalized tool results.

    Storage: ``~/.cadecoder/tool_results/tool_results.jsonl`` with a
    companion ``.lock`` file managed by :mod:`filelock`.  If disk writes
    fail the store falls back to in-memory-only mode for the current
    session.
    """

    def __init__(self, storage_dir: Path | None = None) -> None:
        if storage_dir is None:
            import os

            env_dir = os.environ.get("CADE_TOOL_RESULT_DIR")
            if env_dir:
                storage_dir = Path(env_dir)
            else:
                storage_dir = Path.home() / ".cadecoder" / "tool_results"
        self.storage_dir = storage_dir

        # In-memory session index: record_id -> ToolResultRecord
        self._index: dict[str, ToolResultRecord] = {}
        # Secondary index: tool_call_id -> record_id
        self._call_id_index: dict[str, str] = {}

        # Ensure directory exists
        try:
            self.storage_dir.mkdir(parents=True, exist_ok=True)
        except OSError as e:
            log.warning(f"Cannot create tool result storage dir: {e}")

        self._disk_ok = self.storage_dir.exists()

    # ------------------------------------------------------------------
    # Per-thread file helpers
    # ------------------------------------------------------------------

    def _thread_dir(self, thread_id: str) -> Path:
        """Return the per-thread storage directory, creating it if needed."""
        safe_id = thread_id.replace("/", "_").replace("..", "_") if thread_id else "_default"
        d = self.storage_dir / safe_id
        d.mkdir(parents=True, exist_ok=True)
        return d

    def _thread_jsonl(self, thread_id: str) -> Path:
        return self._thread_dir(thread_id) / "tool_results.jsonl"

    def _thread_lock(self, thread_id: str) -> Path:
        return self._thread_dir(thread_id) / "tool_results.lock"

    # ------------------------------------------------------------------
    # Core storage
    # ------------------------------------------------------------------

    def store_result(self, record: ToolResultRecord) -> ToolResultRef:
        """Append a single record to JSONL and return a compact reference."""
        # In-memory index
        self._index[record.record_id] = record
        self._call_id_index[record.tool_call_id] = record.record_id

        # Disk persistence — write to per-thread file
        if self._disk_ok:
            try:
                from filelock import FileLock

                jsonl_path = self._thread_jsonl(record.thread_id)
                lock_path = self._thread_lock(record.thread_id)
                lock = FileLock(str(lock_path), timeout=10)
                with lock:
                    with open(jsonl_path, "a", encoding="utf-8") as f:
                        f.write(record.to_json_line() + "\n")
            except Exception as e:
                log.warning(f"JSONL write failed, in-memory only: {e}")
                self._disk_ok = False

        return ToolResultRef(
            record_id=record.record_id,
            tool_name=record.tool_name,
            content_chars=record.content_chars,
            status=record.status,
            summary=record.summary,
        )

    def store_batch(
        self,
        tool_results: list[dict[str, Any]],
        tool_calls: list[dict[str, Any]],
        thread_id: str = "",
    ) -> list[ToolResultRef | None]:
        """Decide which results to externalize and store them.

        Returns a list parallel to *tool_results* where each entry is
        either a :class:`ToolResultRef` (externalized) or ``None`` (kept inline).
        """
        if not tool_results:
            return []

        refs: list[ToolResultRef | None] = [None] * len(tool_results)

        # Resolve tool names for each result
        tool_names: list[str] = []
        for i, result in enumerate(tool_results):
            tc = tool_calls[i] if i < len(tool_calls) else {}
            tool_names.append(result.get("name", tc.get("function", {}).get("name", "unknown")))

        # Compute per-result sizes via sys.getsizeof (O(1) for strings)
        import sys

        sizes = [sys.getsizeof(r.get("content", "")) for r in tool_results]
        total_bytes = sum(sizes)

        # Determine which results to externalize
        externalize_flags = [False] * len(tool_results)

        # Level 1: per-result threshold
        for i, size in enumerate(sizes):
            if size > PER_RESULT_BYTE_THRESHOLD:
                externalize_flags[i] = True

        # Level 2: aggregate threshold — externalize everything > TINY
        if total_bytes > AGGREGATE_BYTE_THRESHOLD:
            for i, size in enumerate(sizes):
                if size > TINY_RESULT_BYTES:
                    externalize_flags[i] = True

        # Never externalize retrieval tools — would cause infinite loops
        for i, name in enumerate(tool_names):
            if name in NEVER_EXTERNALIZE_TOOLS:
                externalize_flags[i] = False

        now = datetime.now(UTC).isoformat()
        externalized_count = 0

        for i, should_ext in enumerate(externalize_flags):
            if not should_ext:
                continue

            result = tool_results[i]
            tc = tool_calls[i] if i < len(tool_calls) else {}
            content = result.get("content") or ""
            tool_name = tool_names[i]
            status = result.get("status", "success")
            tool_call_id = tc.get("id", "")

            record = ToolResultRecord(
                record_id=str(ulid()).lower(),
                tool_call_id=tool_call_id,
                tool_name=tool_name,
                thread_id=thread_id,
                timestamp=now,
                content=content,
                content_chars=len(content),
                content_token_estimate=int(len(content) / 4),
                summary=_generate_summary(content, tool_name, status),
                status=status,
                turn_sequence=i,
            )

            refs[i] = self.store_result(record)
            externalized_count += 1

        if externalized_count:
            saved = sum(s for s, f in zip(sizes, externalize_flags) if f)
            log.info(
                f"Externalized {externalized_count}/{len(tool_results)} tool results "
                f"for turn (saved ~{saved:,} bytes)"
            )

        return refs

    # ------------------------------------------------------------------
    # Retrieval
    # ------------------------------------------------------------------

    def retrieve(self, record_id: str) -> ToolResultRecord | None:
        """Retrieve a record by ID. Checks in-memory index first, then disk."""
        if record_id in self._index:
            return self._index[record_id]
        return self._scan_disk_for(record_id=record_id)

    def retrieve_by_tool_call_id(self, tool_call_id: str) -> ToolResultRecord | None:
        """Retrieve a record by its tool_call_id."""
        rid = self._call_id_index.get(tool_call_id)
        if rid:
            return self._index.get(rid)
        return self._scan_disk_for(tool_call_id=tool_call_id)

    def list_by_thread(self, thread_id: str, limit: int = 50) -> list[ToolResultRecord]:
        """List records for a thread, most recent first."""
        # Check in-memory first
        results = [r for r in self._index.values() if r.thread_id == thread_id]

        # If we don't have enough, scan disk
        jsonl = self._thread_jsonl(thread_id)
        if len(results) < limit and self._disk_ok and jsonl.exists():
            seen_ids = {r.record_id for r in results}
            for record in self._iter_disk_records(thread_id):
                if record.record_id not in seen_ids:
                    results.append(record)
                    seen_ids.add(record.record_id)

        # Sort by record_id (ULIDs are time-ordered) descending
        results.sort(key=lambda r: r.record_id, reverse=True)
        return results[:limit]

    def list_records(
        self,
        thread_id: str | None = None,
        tool_name: str | None = None,
        limit: int = 20,
    ) -> list[ToolResultRecord]:
        """List records with optional filters."""
        all_records: dict[str, ToolResultRecord] = {}

        # In-memory
        for r in self._index.values():
            all_records[r.record_id] = r

        # Disk — scan specific thread or all thread dirs
        if self._disk_ok:
            thread_ids = [thread_id] if thread_id else self._list_thread_ids()
            for tid in thread_ids:
                for r in self._iter_disk_records(tid):
                    if r.record_id not in all_records:
                        all_records[r.record_id] = r

        # Apply filters
        results = list(all_records.values())
        if thread_id:
            results = [r for r in results if r.thread_id == thread_id]
        if tool_name:
            results = [r for r in results if r.tool_name == tool_name]

        results.sort(key=lambda r: r.record_id, reverse=True)
        return results[:limit]

    # ------------------------------------------------------------------
    # User context (pinned snippets)
    # ------------------------------------------------------------------

    def store_user_context(
        self,
        thread_id: str,
        name: str,
        content: str,
        source: str = "inline",
        source_path: str = "",
    ) -> ToolResultRecord:
        """Store user-pinned context as a record with tool_name='UserContext'."""
        source_info = f": {source_path}" if source_path else ""
        record = ToolResultRecord(
            record_id=str(ulid()).lower(),
            tool_call_id="",
            tool_name="UserContext",
            thread_id=thread_id,
            timestamp=datetime.now(UTC).isoformat(),
            content=content,
            content_chars=len(content),
            content_token_estimate=int(len(content) / 4),
            summary=f"{name} ({source}{source_info}, {len(content)} chars)",
            status="success",
            turn_sequence=0,
            name=name,
        )
        self.store_result(record)
        return record

    def list_user_context(self, thread_id: str) -> list[ToolResultRecord]:
        """List pinned context for a thread (records where tool_name='UserContext')."""
        all_records = self.list_by_thread(thread_id, limit=100)
        return [r for r in all_records if r.tool_name == "UserContext"]

    def get_user_context_by_name(self, thread_id: str, name: str) -> ToolResultRecord | None:
        """Find a pinned context record by name."""
        for r in self.list_user_context(thread_id):
            if r.name == name:
                return r
        return None

    def remove_user_context(self, thread_id: str, name: str) -> bool:
        """Remove a pinned context record by name."""
        record = self.get_user_context_by_name(thread_id, name)
        if not record:
            return False
        self._index.pop(record.record_id, None)
        if record.tool_call_id:
            self._call_id_index.pop(record.tool_call_id, None)
        self._rewrite_jsonl_without(thread_id, record.record_id)
        return True

    def _rewrite_jsonl_without(self, thread_id: str, exclude_record_id: str) -> None:
        """Rewrite JSONL file excluding a specific record. Protected by filelock."""
        if not self._disk_ok:
            return
        jsonl_path = self._thread_jsonl(thread_id)
        if not jsonl_path.exists():
            return
        try:
            from filelock import FileLock

            lock = FileLock(str(self._thread_lock(thread_id)), timeout=10)
            with lock:
                kept = []
                for record in self._iter_disk_records(thread_id):
                    if record.record_id != exclude_record_id:
                        kept.append(record.to_json_line())
                with open(jsonl_path, "w", encoding="utf-8") as f:
                    for line in kept:
                        f.write(line + "\n")
        except Exception as e:
            log.warning(f"JSONL rewrite failed for thread {thread_id}: {e}")

    # ------------------------------------------------------------------
    # Cleanup
    # ------------------------------------------------------------------

    def cleanup_old_records(self, max_age_hours: int = 72) -> int:
        """Rewrite JSONL files excluding records older than max_age_hours."""
        if not self._disk_ok:
            return 0

        cutoff = datetime.now(UTC) - timedelta(hours=max_age_hours)
        total_expired = 0

        for tid in self._list_thread_ids():
            jsonl_path = self._thread_jsonl(tid)
            if not jsonl_path.exists():
                continue

            kept: list[str] = []
            expired_ids: list[tuple[str, str]] = []

            for record in self._iter_disk_records(tid):
                try:
                    ts = datetime.fromisoformat(record.timestamp)
                    if ts.tzinfo is None:
                        ts = ts.replace(tzinfo=UTC)
                    if ts < cutoff:
                        expired_ids.append((record.record_id, record.tool_call_id))
                        continue
                except (ValueError, TypeError):
                    pass
                kept.append(record.to_json_line())

            if expired_ids:
                try:
                    from filelock import FileLock

                    lock = FileLock(str(self._thread_lock(tid)), timeout=10)
                    with lock:
                        with open(jsonl_path, "w", encoding="utf-8") as f:
                            for line in kept:
                                f.write(line + "\n")
                except Exception as e:
                    log.warning(f"Cleanup rewrite failed for thread {tid}: {e}")
                    continue

                for record_id, tool_call_id in expired_ids:
                    self._index.pop(record_id, None)
                    self._call_id_index.pop(tool_call_id, None)
                total_expired += len(expired_ids)

        if total_expired:
            log.info(f"Cleaned up {total_expired} old tool result records")
        return total_expired

    # ------------------------------------------------------------------
    # Internal disk helpers
    # ------------------------------------------------------------------

    def _scan_disk_for(
        self,
        record_id: str | None = None,
        tool_call_id: str | None = None,
    ) -> ToolResultRecord | None:
        """Scan JSONL files across all threads for a specific record."""
        if not self._disk_ok:
            return None

        for tid in self._list_thread_ids():
            for record in self._iter_disk_records(tid):
                if record_id and record.record_id == record_id:
                    self._index[record.record_id] = record
                    self._call_id_index[record.tool_call_id] = record.record_id
                    return record
                if tool_call_id and record.tool_call_id == tool_call_id:
                    self._index[record.record_id] = record
                    self._call_id_index[record.tool_call_id] = record.record_id
                    return record

        return None

    def _list_thread_ids(self) -> list[str]:
        """List thread IDs that have storage directories."""
        if not self.storage_dir.exists():
            return []
        return [
            d.name
            for d in self.storage_dir.iterdir()
            if d.is_dir() and (d / "tool_results.jsonl").exists()
        ]

    def _iter_disk_records(self, thread_id: str) -> list[ToolResultRecord]:
        """Read all valid records from a thread's JSONL, skipping corrupt lines."""
        records: list[ToolResultRecord] = []
        jsonl_path = self._thread_jsonl(thread_id)
        if not jsonl_path.exists():
            return records

        try:
            with open(jsonl_path, encoding="utf-8") as f:
                for line_no, line in enumerate(f, 1):
                    line = line.strip()
                    if not line:
                        continue
                    try:
                        records.append(ToolResultRecord.from_json_line(line))
                    except (json.JSONDecodeError, TypeError, KeyError) as e:
                        log.debug(f"Skipping corrupt JSONL line {line_no}: {e}")
        except OSError as e:
            log.warning(f"Failed to read JSONL for thread {thread_id}: {e}")

        return records
