#!/usr/bin/env python3
"""Adaptive compressor with structural tool-encoding merged in (COACH order).


Adaptive Token-Aware Compression Strategy

Based on research insights:
- Dynamically adjust compression ratio based on message type
- Preserve high-importance messages (errors, file changes)
- Aggressively compress low-importance messages (verbose output)
- Token-budget aware
"""

from __future__ import annotations
import re
from typing import Any


def estimate_tokens(text: str) -> int:
    """Rough token estimate: ~4 chars per token"""
    return len(text) // 4


def get_message_importance(msg: dict[str, Any]) -> float:
    """
    Score message importance (0-10).
    Higher score = keep more content.
    """
    content = str(msg.get('content', ''))

    # High importance signals
    score = 5.0  # baseline

    if 'error' in content.lower() or 'exception' in content.lower():
        score += 3.0
    if 'failed' in content.lower() or 'failure' in content.lower():
        score += 2.0
    if re.search(r'\btest\b.*\b(?:pass|fail)', content.lower()):
        score += 2.0
    if re.search(r'\bfile\b.*\b(?:changed|modified|created)', content.lower()):
        score += 2.0
    if '```' in content:  # code blocks
        score += 1.5

    # Low importance signals
    if re.search(r'^\s*(?:ok|done|success|completed)\s*$', content.lower()):
        score -= 2.0
    if len(content) > 2000 and content.count('\n') > 50:  # verbose output
        score -= 1.0

    return max(0.0, min(10.0, score))


def compress_by_importance(content: str, importance: float, max_tokens: int = 200) -> str:
    """
    Compress content based on importance score.
    Higher importance = keep more content.

    Splits are budgeted so head+tail is always a strict subset of the input,
    and the result is never longer than what was passed in.
    """
    current_tokens = estimate_tokens(content)

    # Scale max_tokens by importance (0.5x to 2x)
    importance_factor = 0.5 + (importance / 10.0) * 1.5
    target_tokens = int(max_tokens * importance_factor)

    if current_tokens <= target_tokens:
        return content

    lines = content.split('\n')

    # Fraction of lines to retain, by importance band. These are total
    # budgets (<1.0) so the elision marker always replaces real content.
    if importance >= 7.0:
        keep_fraction = 0.60
    elif importance >= 4.0:
        keep_fraction = 0.40
    else:
        keep_fraction = 0.20

    # Nothing to gain from eliding two lines or fewer.
    if len(lines) <= 3:
        return content

    budget = max(2, int(len(lines) * keep_fraction))
    if budget >= len(lines) - 1:
        budget = len(lines) - 2

    keep_start = max(1, (budget * 3) // 4)
    keep_end = max(1, budget - keep_start)

    head = lines[:keep_start]
    tail = lines[len(lines) - keep_end:]
    omitted = len(lines) - len(head) - len(tail)
    if omitted <= 0:
        return content

    result = '\n'.join(head + [f'[... {omitted} lines omitted ...]'] + tail)

    # Hard guarantee: never return more than we were given.
    return result if len(result) < len(content) else content


def compress_messages(
    messages: list[Any] | None = None,
    path: str | None = None,
    metadata: dict[str, Any] | None = None,
) -> list[Any]:
    """
    Adaptive compression based on message importance.
    """
    del path, metadata

    if not isinstance(messages, list):
        return []

    compressed = []

    for msg in messages:
        if not isinstance(msg, dict):
            compressed.append(msg)
            continue

        new_msg = msg.copy()

        if isinstance(new_msg.get('content'), str):
            importance = get_message_importance(new_msg)
            content = new_msg['content']

            # Apply adaptive compression
            compressed_content = compress_by_importance(content, importance)

            # Preserve the message even when compression yields nothing.
            # Six legacy compressors returned 0 messages for a 1-message input
            # with empty or whitespace-only content, and 1 message for a 2-message
            # input where the first was empty -- silently deleting a conversation
            # turn and changing the transcript structure the model sees. Emitting
            # the (empty) message keeps the message count invariant.
            new_msg['content'] = compressed_content
            compressed.append(new_msg)
        else:
            compressed.append(new_msg)

    return compressed


if __name__ == '__main__':
    # Test
    test_msg = {
        'role': 'assistant',
        'content': 'Error: Authentication failed\n' + 'x' * 1000
    }
    imp = get_message_importance(test_msg)
    print(f"Importance: {imp}/10")
    result = compress_messages([test_msg])
    print(f"Original: {len(test_msg['content'])} chars")
    print(f"Compressed: {len(result[0]['content'])} chars")


# --- structural tool-encoding merged in (COACH order) ---
import importlib.util as _u
from pathlib import Path as _P

_spec = _u.spec_from_file_location(
    "_struct_enc", _P("/var/lib/octave/sn114/external/SOMA-plugin")
    / "structural_cot_compressor.py")
_struct = _u.module_from_spec(_spec)
_spec.loader.exec_module(_struct)

_adaptive_only = compress_messages


def compress_messages(
    messages: list[dict[str, Any]] | None = None,
    path: str | None = None,
    metadata: dict[str, Any] | None = None,
):
    del path, metadata
    """Run adaptive importance selection, then structural tool-encoding.

    Order matters: adaptive prunes by importance, the structural pass then
    rewrites surviving tool scaffolding to the compact sigil encoding.
    """
    if not messages:
        return []
    out = _adaptive_only(messages)
    return _struct.compress_messages(out)