"""Context management: observation masking, tool result offloading, memory pressure.

After each agent turn, old tool results are replaced with one-line summaries
(observation masking).  Large results are offloaded to disk so the agent can
re-read them if needed.  When the context approaches the compaction threshold
a pressure warning is injected into the system prompt.

Based on JetBrains' Complexity Trap finding: simple observation masking
performs as well as LLM summarization, 7%+ cheaper.
"""

from __future__ import annotations

import hashlib
import json
import logging
from pathlib import Path
from typing import Any, Dict, List, Optional

logger = logging.getLogger(__name__)

# Keep this many recent user-turns with full tool results
KEEP_VERBATIM_TURNS = 10

# Tool results larger than this are saved to disk when masked
OFFLOAD_CHAR_THRESHOLD = 8_000

# Warn the agent when estimated tokens exceed this
PRESSURE_WARNING_THRESHOLD = 55_000

# Conservative chars-per-token estimate (matches cli_v2.py)
CHARS_PER_TOKEN = 4


# ---------------------------------------------------------------------------
# Observation masking
# ---------------------------------------------------------------------------

def mask_old_observations(
    history: list,
    project_path: str,
    keep_turns: int = KEEP_VERBATIM_TURNS,
) -> int:
    """Replace old ToolReturnPart content with one-line summaries.

    Keeps the last ``keep_turns`` user turns verbatim.  Older tool results
    are replaced in-place.  Large results are offloaded to disk first.

    Returns the number of tool results that were masked.
    """
    from pydantic_ai.messages import ModelRequest, ToolReturnPart

    turn_indices = _count_user_turns(history)
    if len(turn_indices) <= keep_turns:
        return 0

    # Everything before this index is eligible for masking
    keep_boundary = turn_indices[-keep_turns]
    masked_count = 0

    for i in range(keep_boundary):
        msg = history[i]
        if not isinstance(msg, ModelRequest):
            continue
        for part in msg.parts:
            if not isinstance(part, ToolReturnPart):
                continue
            content = part.content
            if not isinstance(content, str) or len(content) < 100:
                continue
            if content.startswith("[masked]"):
                continue

            # Find args from the matching ToolCallPart
            args = _find_tool_call_args(history, part.tool_call_id, max_lookback=i + 1)

            # Offload large results to disk
            offload_path = None
            if len(content) >= OFFLOAD_CHAR_THRESHOLD:
                offload_path = _offload_tool_result(content, part.tool_name, project_path)

            part.content = _summarize_tool_result(
                part.tool_name, content, args=args, offload_path=offload_path,
            )
            masked_count += 1

    return masked_count


# ---------------------------------------------------------------------------
# Tool result summaries (rule-based, no LLM call)
# ---------------------------------------------------------------------------

_SUMMARY_TEMPLATES: Dict[str, str] = {
    "read_file": "[masked] read_file({path}) -> {lines} lines, {chars} chars",
    "search_codebase": "[masked] search_codebase({query}) -> {chars} chars",
    "scan_codebase": "[masked] scan_codebase() -> {chars} chars of scan results",
    "run_tests": "[masked] run_tests({path}) -> {chars} chars",
    "detect_flaky_patterns": "[masked] detect_flaky_patterns({path}) -> {chars} chars",
    "heal_flaky_test": "[masked] heal_flaky_test({path}) -> {chars} chars",
    "fetch_production_errors": "[masked] fetch_production_errors() -> {chars} chars",
    "get_error_details": "[masked] get_error_details({error_id}) -> {chars} chars",
    "view_io_transcript": "[masked] view_io_transcript() -> {chars} chars",
    "review_files": "[masked] review_files({paths}) -> {chars} chars",
    "generate_fix_hypothesis": "[masked] generate_fix_hypothesis() -> {chars} chars",
}

_DEFAULT_TEMPLATE = "[masked] {tool_name}() -> {chars} chars"


def _summarize_tool_result(
    tool_name: str,
    content: str,
    args: Any = None,
    offload_path: Optional[str] = None,
) -> str:
    """Generate a one-line summary for a masked tool result."""
    chars = len(content)
    lines = content.count("\n") + 1 if content else 0

    # Extract hints from args
    hints: Dict[str, str] = {"tool_name": tool_name, "chars": f"{chars:,}", "lines": str(lines)}
    if isinstance(args, dict):
        hints["path"] = args.get("path", args.get("file_path", ""))
        hints["query"] = args.get("query", "")
        hints["error_id"] = args.get("error_id", "")
        hints["paths"] = str(args.get("paths", ""))[:60]
    elif isinstance(args, str):
        try:
            parsed = json.loads(args)
            if isinstance(parsed, dict):
                hints["path"] = parsed.get("path", parsed.get("file_path", ""))
                hints["query"] = parsed.get("query", "")
        except (json.JSONDecodeError, TypeError):
            pass

    template = _SUMMARY_TEMPLATES.get(tool_name, _DEFAULT_TEMPLATE)
    try:
        summary = template.format_map(hints)
    except (KeyError, IndexError):
        summary = f"[masked] {tool_name}() -> {chars:,} chars"

    if offload_path:
        summary += f"  Full: {offload_path}"

    return summary


# ---------------------------------------------------------------------------
# Tool result offloading
# ---------------------------------------------------------------------------

def _offload_tool_result(
    content: str,
    tool_name: str,
    project_path: str,
) -> Optional[str]:
    """Save a large tool result to ``.hikaflow/tmp/`` for later re-reading.

    Returns the relative path, or None if below threshold or on failure.
    """
    if len(content) < OFFLOAD_CHAR_THRESHOLD:
        return None

    tmp_dir = Path(project_path) / ".hikaflow" / "tmp"
    try:
        tmp_dir.mkdir(parents=True, exist_ok=True)
    except OSError:
        return None

    content_hash = hashlib.sha256(content.encode("utf-8", errors="replace")).hexdigest()[:12]
    filename = f"{tool_name}_{content_hash}.txt"
    filepath = tmp_dir / filename

    try:
        filepath.write_text(content, encoding="utf-8")
    except OSError:
        return None

    return f".hikaflow/tmp/{filename}"


# ---------------------------------------------------------------------------
# Memory pressure detection
# ---------------------------------------------------------------------------

def estimate_context_tokens(history: list) -> int:
    """Estimate token count from message history (chars / 4)."""
    total_chars = 0
    for msg in history:
        if hasattr(msg, "parts"):
            for part in msg.parts:
                if hasattr(part, "content"):
                    val = part.content
                    if isinstance(val, str):
                        total_chars += len(val)
                    elif isinstance(val, list):
                        for item in val:
                            if isinstance(item, str):
                                total_chars += len(item)
                if hasattr(part, "args"):
                    total_chars += len(str(part.args))
    return total_chars // CHARS_PER_TOKEN


def check_context_pressure(history: list) -> bool:
    """Return True if estimated tokens exceed the pressure warning threshold."""
    return estimate_context_tokens(history) >= PRESSURE_WARNING_THRESHOLD


def render_pressure_warning() -> str:
    """Static markdown warning injected via ``@agent.system_prompt``."""
    return (
        "\n## WARNING: Context Memory Pressure\n\n"
        "Your conversation context is approaching the compaction threshold. "
        "Important findings and file contents will be dropped soon.\n\n"
        "**Action required:** Use memory_insert or memory_rethink NOW to save "
        "any critical findings to your session_state or project_context blocks "
        "before they are lost to compaction.\n"
    )


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

def cleanup_offloaded_files(project_path: str) -> int:
    """Remove offloaded files from ``.hikaflow/tmp/``.  Returns count removed."""
    tmp_dir = Path(project_path) / ".hikaflow" / "tmp"
    if not tmp_dir.is_dir():
        return 0

    count = 0
    try:
        for f in tmp_dir.iterdir():
            if f.is_file() and f.suffix == ".txt":
                f.unlink(missing_ok=True)
                count += 1
        if not any(tmp_dir.iterdir()):
            tmp_dir.rmdir()
    except OSError:
        pass
    return count


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

def _count_user_turns(history: list) -> List[int]:
    """Return indices of ModelRequest messages that start user turns (chronological)."""
    from pydantic_ai.messages import ModelRequest, UserPromptPart

    indices: List[int] = []
    for i, msg in enumerate(history):
        if not isinstance(msg, ModelRequest):
            continue
        if any(isinstance(p, UserPromptPart) for p in msg.parts):
            indices.append(i)
    return indices


def _find_tool_call_args(
    history: list,
    tool_call_id: str,
    max_lookback: int,
) -> Any:
    """Find ToolCallPart args matching a tool_call_id."""
    from pydantic_ai.messages import ModelResponse, ToolCallPart

    for i in range(max_lookback - 1, -1, -1):
        msg = history[i]
        if not isinstance(msg, ModelResponse):
            continue
        for part in msg.parts:
            if isinstance(part, ToolCallPart) and part.tool_call_id == tool_call_id:
                return part.args
    return None
