"""Context pruning engine for tool result management.

Pure module — no DB access, no async. Takes config + content, returns pruned content.
"""

from __future__ import annotations

import re
from dataclasses import dataclass, field
from datetime import datetime, timezone
from typing import TYPE_CHECKING

if TYPE_CHECKING:
    from peon_mcp.projects.schemas import ContextPruningConfig


@dataclass
class ToolResult:
    """Represents a tool invocation result with metadata."""

    tool_name: str
    content: str
    params: dict = field(default_factory=dict)
    created_at: datetime = field(default_factory=lambda: datetime.now(timezone.utc))


def estimate_tokens(text: str) -> int:
    """Rough token estimate for English text (len / 4).

    Used for budgeting, not exact counting.
    """
    return len(text) // 4


def should_prune(tool_name: str, config: ContextPruningConfig) -> bool:
    """Return True if the tool's output should be subject to pruning.

    Returns False if pruning is disabled or the tool is in never_prune_tools.
    """
    if not config.enabled:
        return False
    if tool_name in config.never_prune_tools:
        return False
    return True


def summarize_large_output(content: str, max_chars: int) -> str:
    """Produce a structured summary for very large outputs (>2x max_chars).

    Extracts:
      - First 5 lines (headers / summary)
      - Lines containing 'error', 'fail', or 'warning' (case-insensitive)
      - Last 5 lines (summary / totals)
    """
    lines = content.splitlines()
    if not lines:
        return content

    first_lines = lines[:5]
    last_lines = lines[-5:] if len(lines) > 5 else []
    important_lines = [
        line
        for line in lines[5 : len(lines) - 5]
        if re.search(r"error|fail|warning", line, re.IGNORECASE)
    ]

    parts: list[str] = []
    parts.append("=== Summary (first 5 lines) ===")
    parts.extend(first_lines)

    if important_lines:
        parts.append(f"=== Notable lines ({len(important_lines)} found) ===")
        parts.extend(important_lines)

    if last_lines:
        parts.append("=== Last 5 lines ===")
        parts.extend(last_lines)

    summary = "\n".join(parts)
    if len(summary) > max_chars:
        summary = summary[:max_chars] + f"\n... [summary truncated to {max_chars} chars]"
    return summary


def prune_tool_result(
    content: str,
    tool_name: str,
    config: ContextPruningConfig,
) -> str:
    """Prune tool output to fit within the configured character limit.

    - Looks up per-tool limit (falls back to default_max_chars).
    - If content fits, returns it unchanged.
    - For very large outputs (>2x limit) delegates to summarize_large_output.
    - Otherwise applies head/tail preservation with a pruned-marker in the middle.
    - Lines containing always_preserve keywords are exempted from pruning.
    """
    limit = config.per_tool_limits.get(tool_name, config.default_max_chars)

    if len(content) <= limit:
        return content

    # For very large outputs use the structured summary approach
    if len(content) > 2 * limit:
        return summarize_large_output(content, limit)

    head_ratio, tail_ratio = config.head_tail_ratio[0], config.head_tail_ratio[1]
    head_chars = int(head_ratio * limit)
    tail_chars = int(tail_ratio * limit)

    # Collect lines that must be preserved
    preserve_keywords = [kw.lower() for kw in config.always_preserve]
    preserved_lines: list[tuple[int, str]] = []
    lines = content.splitlines(keepends=True)
    for idx, line in enumerate(lines):
        lower = line.lower()
        if any(kw in lower for kw in preserve_keywords):
            preserved_lines.append((idx, line))

    head = content[:head_chars]
    tail = content[-tail_chars:] if tail_chars > 0 else ""
    removed_chars = len(content) - head_chars - tail_chars

    # Build the pruned result
    pruned = head + f"\n... [pruned {removed_chars} chars] ...\n" + tail

    # Re-inject any preserved lines that were cut out of both head and tail
    head_end = head_chars
    tail_start = len(content) - tail_chars if tail_chars > 0 else len(content)

    # Determine char offsets for each line
    char_offset = 0
    line_char_offsets: list[int] = []
    for line in lines:
        line_char_offsets.append(char_offset)
        char_offset += len(line)

    missing_preserved: list[str] = []
    for idx, line in preserved_lines:
        start = line_char_offsets[idx]
        end = start + len(line)
        # If the line falls entirely in the pruned region, we need to inject it
        if start >= head_end and end <= tail_start:
            missing_preserved.append(line.rstrip("\n"))

    if missing_preserved:
        preserved_block = "\n".join(missing_preserved)
        pruned = (
            head
            + f"\n... [pruned {removed_chars} chars] ...\n"
            + f"[preserved lines]\n{preserved_block}\n[/preserved lines]\n"
            + tail
        )

    return pruned


def prune_stale_results(
    results: list[ToolResult],
    config: ContextPruningConfig,
) -> list[ToolResult]:
    """Replace the content of results that have exceeded their TTL.

    Stale results have their content replaced with an expiry notice.
    The tool name and params remain intact.
    """
    now = datetime.now(timezone.utc)
    ttl_minutes = config.ttl_minutes
    pruned: list[ToolResult] = []

    for result in results:
        created = result.created_at
        # Normalise naive datetimes to UTC for comparison
        if created.tzinfo is None:
            created = created.replace(tzinfo=timezone.utc)

        age_minutes = (now - created).total_seconds() / 60
        if age_minutes > ttl_minutes:
            expired = ToolResult(
                tool_name=result.tool_name,
                content=(
                    f"[Result expired after {ttl_minutes} minutes. "
                    "Re-run the tool if needed.]"
                ),
                params=result.params,
                created_at=result.created_at,
            )
            pruned.append(expired)
        else:
            pruned.append(result)

    return pruned
