"""
Context Manager — Token-aware data compression for report_agent.

Adapted from hermes-agent-repo/agent/context_compressor.py. Manages the
token budget during multi-cycle report generation by pruning oversized
tool results and compressing intermediate LLM outputs.

Key patterns from context_compressor.py:
  - Token budgeting per message role
  - PRUNED placeholder for removed content
  - Priority-based content ranking

Usage:
    from agent.report_agent.context_manager import ReportContextManager

    mgr = ReportContextManager(max_tokens=8000)
    compressed = mgr.compress_tool_results(results, question="Tổng doanh thu")
    budget_info = mgr.get_budget_info()
"""

from __future__ import annotations

import logging
import re
from dataclasses import dataclass
from typing import Any, Dict, List, Optional, Tuple

logger = logging.getLogger(__name__)

# ─── Constants ──────────────────────────────────────────────────────────

# Approximate chars per token (conservative for Vietnamese/mixed content)
CHARS_PER_TOKEN = 3.5

# Placeholder for pruned content (mirrors Hermes pattern)
_PRUNED_TOOL_PLACEHOLDER = "[...kết quả quá dài — đã cắt bớt để tiết kiệm context...]"
_PRUNED_SUMMARY_PLACEHOLDER = "[...phân tích trước đã được tóm tắt...]"

# Default token limits per component
DEFAULT_TOOL_RESULT_BUDGET = 4000  # tokens per tool result
DEFAULT_TOTAL_CONTEXT_BUDGET = 12000  # total tokens for all context
DEFAULT_HISTORY_BUDGET = 3000  # tokens for conversation history


@dataclass
class BudgetInfo:
    """Current token budget status."""
    total_budget: int
    used_tokens: int
    remaining_tokens: int
    compressions_applied: int = 0
    items_pruned: int = 0


class ReportContextManager:
    """Token-aware context manager for report generation.

    Manages the token budget during multi-cycle report generation,
    ensuring tool results, history, and intermediate analysis fit
    within the LLM's context window.
    """

    def __init__(
        self,
        max_tokens: int = DEFAULT_TOTAL_CONTEXT_BUDGET,
        tool_budget: int = DEFAULT_TOOL_RESULT_BUDGET,
        history_budget: int = DEFAULT_HISTORY_BUDGET,
    ):
        self.max_tokens = max_tokens
        self.tool_budget = tool_budget
        self.history_budget = history_budget
        self._compressions = 0
        self._items_pruned = 0

    def _estimate_tokens(self, text: str) -> int:
        """Estimate token count from text length."""
        return max(1, int(len(text) / CHARS_PER_TOKEN))

    def _truncate_to_tokens(self, text: str, max_tokens: int) -> Tuple[str, bool]:
        """Truncate text to fit within token budget.

        Returns (truncated_text, was_truncated).
        """
        max_chars = int(max_tokens * CHARS_PER_TOKEN)
        if len(text) <= max_chars:
            return text, False
        return text[:max_chars] + f"\n{_PRUNED_TOOL_PLACEHOLDER}", True

    # ── Public API ──

    def compress_tool_results(
        self,
        results: List[Dict[str, Any]],
        question: str = "",
    ) -> List[Dict[str, Any]]:
        """Compress a list of tool results to fit within the token budget.

        Each result dict should have:
          - 'tool': tool name
          - 'result': result text/data
          - 'priority': optional priority (higher = keep more)

        Returns the compressed list with oversized results truncated.
        """
        if not results:
            return results

        compressed = []
        total_tokens_used = 0
        remaining_budget = self.max_tokens

        # Sort by priority (higher first) if available
        sorted_results = sorted(
            results,
            key=lambda r: r.get("priority", 0),
            reverse=True,
        )

        for item in sorted_results:
            tool = item.get("tool", "unknown")
            result_text = str(item.get("result", ""))

            tokens = self._estimate_tokens(result_text)
            per_item_budget = min(self.tool_budget, remaining_budget)

            if tokens > per_item_budget:
                # Truncate to fit
                truncated, was_cut = self._truncate_to_tokens(
                    result_text, per_item_budget
                )
                if was_cut:
                    self._compressions += 1
                    self._items_pruned += 1
                    logger.info(
                        "Compressed tool result '%s': %d → %d tokens",
                        tool, tokens, self._estimate_tokens(truncated),
                    )
                compressed.append({
                    **item,
                    "result": truncated,
                    "compressed": was_cut,
                })
                total_tokens_used += self._estimate_tokens(truncated)
            else:
                compressed.append(item)
                total_tokens_used += tokens

            remaining_budget = self.max_tokens - total_tokens_used
            if remaining_budget <= 0:
                # Drop remaining low-priority results
                dropped = len(sorted_results) - len(compressed)
                if dropped > 0:
                    self._items_pruned += dropped
                    logger.info(
                        "Dropped %d low-priority tool results (budget exhausted)",
                        dropped,
                    )
                break

        return compressed

    def compress_history(
        self,
        messages: List[Dict[str, str]],
    ) -> List[Dict[str, str]]:
        """Compress conversation history to fit within history budget.

        Keeps the most recent messages, summarizing older ones.
        Priority: system > last user > last assistant > older messages.
        """
        if not messages:
            return messages

        total_tokens = sum(
            self._estimate_tokens(m.get("content", "")) for m in messages
        )

        if total_tokens <= self.history_budget:
            return messages

        # Keep system message + last 2 user/assistant exchanges
        system_msgs = [m for m in messages if m.get("role") == "system"]
        non_system = [m for m in messages if m.get("role") != "system"]

        # Always keep the last 4 non-system messages (2 exchanges)
        keep_recent = non_system[-4:] if len(non_system) > 4 else non_system
        older = non_system[:-4] if len(non_system) > 4 else []

        if older:
            # Summarize older messages into a single context entry
            older_text = "\n".join(
                f"[{m.get('role', '?')}]: {m.get('content', '')[:200]}"
                for m in older
            )
            summary = {
                "role": "system",
                "content": f"{_PRUNED_SUMMARY_PLACEHOLDER}\n"
                           f"Tóm tắt {len(older)} tin nhắn trước:\n"
                           f"{older_text[:int(self.history_budget * CHARS_PER_TOKEN * 0.3)]}",
            }
            self._compressions += 1
            return system_msgs + [summary] + keep_recent

        return system_msgs + keep_recent

    def compress_sql_result(
        self,
        result_text: str,
        max_rows: int = 50,
    ) -> str:
        """Compress a SQL query result by limiting rows.

        Detects table-formatted results and truncates to max_rows.
        """
        if self._estimate_tokens(result_text) <= self.tool_budget:
            return result_text

        lines = result_text.split("\n")
        if len(lines) <= max_rows + 5:  # Header + separator + rows + footer
            return result_text

        # Keep header (first 3 lines: header, separator, first row pattern)
        header = lines[:3]
        data_lines = lines[3:]

        kept = data_lines[:max_rows]
        dropped = len(data_lines) - max_rows

        self._compressions += 1
        self._items_pruned += 1

        return "\n".join(
            header
            + kept
            + [f"\n... ({dropped} dòng nữa đã ẩn để tiết kiệm context)"]
        )

    def get_budget_info(self) -> BudgetInfo:
        """Get current budget status."""
        return BudgetInfo(
            total_budget=self.max_tokens,
            used_tokens=0,  # Reset per-call — caller tracks actual usage
            remaining_tokens=self.max_tokens,
            compressions_applied=self._compressions,
            items_pruned=self._items_pruned,
        )

    def reset_counters(self):
        """Reset compression counters for a new session."""
        self._compressions = 0
        self._items_pruned = 0
