"""
Core compression logic for token-diet.
Handles both rule-based (offline) and API-based (Claude) compression.
"""

from __future__ import annotations

import re
from dataclasses import dataclass, field
from enum import Enum
from typing import List, Optional, Tuple


class CompressionLevel(str, Enum):
    GENTLE = "gentle"
    BALANCED = "balanced"
    AGGRESSIVE = "aggressive"


@dataclass
class CompressionResult:
    original_text: str
    compressed_text: str
    original_tokens: int
    compressed_tokens: int
    techniques_applied: List[str] = field(default_factory=list)
    used_api: bool = False

    @property
    def tokens_saved(self) -> int:
        return self.original_tokens - self.compressed_tokens

    @property
    def percent_saved(self) -> float:
        if self.original_tokens == 0:
            return 0.0
        return (self.tokens_saved / self.original_tokens) * 100.0

    @property
    def original_chars(self) -> int:
        return len(self.original_text)

    @property
    def compressed_chars(self) -> int:
        return len(self.compressed_text)

    @property
    def cost_saved_gpt4o(self) -> float:
        """Estimated cost savings per call at GPT-4o input pricing (~$5/1M tokens)."""
        return (self.tokens_saved / 1_000_000) * 5.0


def estimate_tokens(text: str) -> int:
    """Estimate token count using the 4-chars-per-token heuristic."""
    return max(1, len(text) // 4)


def count_tokens(text: str) -> int:
    """Count tokens, using tiktoken if available, else the estimator."""
    try:
        import tiktoken  # type: ignore[import]

        enc = tiktoken.get_encoding("cl100k_base")
        return len(enc.encode(text))
    except ImportError:
        return estimate_tokens(text)


# ---------------------------------------------------------------------------
# Rule-based substitutions
# ---------------------------------------------------------------------------

# Each entry: (pattern, replacement, description)
# Patterns are compiled with re.IGNORECASE where the flag is True.
_FILLER_PATTERNS: List[Tuple[str, str, str, bool]] = [
    # Opening politeness phrases
    (r"\bPlease\s+", "", "Removed filler phrases", True),
    (r"\bCould you please\s+", "", "Removed filler phrases", True),
    (r"\bCould you\s+", "", "Removed filler phrases", True),
    (r"\bI would like you to\s+", "", "Removed filler phrases", True),
    (r"\bI'd like you to\s+", "", "Removed filler phrases", True),
    (r"\bI want you to\s+", "", "Removed filler phrases", True),
    (r"\bI need you to\s+", "", "Removed filler phrases", True),
    (r"\bMake sure to\s+", "", "Removed filler phrases", True),
    (r"\bMake sure that you\s+", "", "Removed filler phrases", True),
    (r"\bIt is important that you\s+", "", "Removed filler phrases", True),
    (r"\bIt is important that\s+", "", "Removed filler phrases", True),
    (r"\bPlease note that\s+", "", "Removed filler phrases", True),
    (r"\bNote that\s+", "", "Removed filler phrases", True),
    (r"\bKindly\s+", "", "Removed filler phrases", True),
    (r"\bFeel free to\s+", "", "Removed filler phrases", True),
    # Verbose constructions
    (r"\bin order to\b", "to", "Shortened verbose constructions", True),
    (r"\bat this point in time\b", "now", "Shortened verbose constructions", True),
    (r"\bat the present time\b", "now", "Shortened verbose constructions", True),
    (r"\bdue to the fact that\b", "because", "Shortened verbose constructions", True),
    (r"\bin the event that\b", "if", "Shortened verbose constructions", True),
    (r"\bfor the purpose of\b", "for", "Shortened verbose constructions", True),
    (r"\bwith the exception of\b", "except", "Shortened verbose constructions", True),
    (r"\bprior to\b", "before", "Shortened verbose constructions", True),
    (r"\bsubsequent to\b", "after", "Shortened verbose constructions", True),
    (r"\bin spite of the fact that\b", "although", "Shortened verbose constructions", True),
    (r"\bregardless of the fact that\b", "although", "Shortened verbose constructions", True),
    (r"\bhas the ability to\b", "can", "Shortened verbose constructions", True),
    (r"\bis able to\b", "can", "Shortened verbose constructions", True),
    (r"\bwas able to\b", "could", "Shortened verbose constructions", True),
    (r"\bin the process of\b", "while", "Shortened verbose constructions", True),
    (r"\ba large number of\b", "many", "Shortened verbose constructions", True),
    (r"\ba number of\b", "some", "Shortened verbose constructions", True),
    (r"\bthe majority of\b", "most", "Shortened verbose constructions", True),
    (r"\bon a regular basis\b", "regularly", "Shortened verbose constructions", True),
    (r"\bat all times\b", "always", "Shortened verbose constructions", True),
    # Redundancies
    (
        r"\bsummarize and (create|write|produce) a summary\b",
        "summarize",
        "Collapsed redundant instructions",
        True,
    ),
    (
        r"\bsummariz(?:e|ing) (?:the )?(?:content|information|material|text) (?:and|that is) (?:found|contained|provided) (?:in|within)\b",
        "summarize",
        "Collapsed redundant instructions",
        True,
    ),
    (
        r"\bplease (?:note|be aware|keep in mind) that\b",
        "",
        "Collapsed redundant instructions",
        True,
    ),
    (
        r"\bcomprehensive(ly)?\b",
        r"thorough\1",
        "Shortened verbose constructions",
        True,
    ),
]

_BALANCED_PATTERNS: List[Tuple[str, str, str, bool]] = [
    (r"\bthat which\b", "what", "Shortened verbose constructions", True),
    (r"\bthe reason why\b", "why", "Shortened verbose constructions", True),
    (r"\bwith regards to\b", "about", "Shortened verbose constructions", True),
    (r"\bwith respect to\b", "about", "Shortened verbose constructions", True),
    (r"\bwith reference to\b", "about", "Shortened verbose constructions", True),
    (r"\bin relation to\b", "about", "Shortened verbose constructions", True),
    (r"\bvery important\b", "critical", "Shortened verbose constructions", True),
    (r"\bextremely important\b", "critical", "Shortened verbose constructions", True),
    (r"\bvery\b", "", "Removed filler phrases", True),
    (r"\bquite\b", "", "Removed filler phrases", True),
    (r"\brather\b", "", "Removed filler phrases", True),
    (r"\bsomewhat\b", "", "Removed filler phrases", True),
    (r"\bbasically\b", "", "Removed filler phrases", True),
    (r"\bactually\b", "", "Removed filler phrases", True),
    (r"\bessentially\b", "", "Removed filler phrases", True),
    (r"\bfundamentally\b", "", "Removed filler phrases", True),
    (r"\bgenerally speaking,?\s*", "", "Removed filler phrases", True),
    (r"\bof course,?\s*", "", "Removed filler phrases", True),
    (r"\bnaturally,?\s*", "", "Removed filler phrases", True),
    (r"\bobviously,?\s*", "", "Removed filler phrases", True),
    (r"\bclearly,?\s*", "", "Removed filler phrases", True),
    (r"\bin fact,?\s*", "", "Removed filler phrases", True),
    (r"\bas a matter of fact,?\s*", "", "Removed filler phrases", True),
    (r"\bto be honest(?:,| with you)?\s*", "", "Removed filler phrases", True),
    (r"\btruly\b", "", "Removed filler phrases", True),
    (r"\bcertainly\b", "", "Removed filler phrases", True),
    (r"\babsolutely\b", "", "Removed filler phrases", True),
    (r"\bdefinitely\b", "", "Removed filler phrases", True),
]

_AGGRESSIVE_PATTERNS: List[Tuple[str, str, str, bool]] = [
    # Articles in some contexts
    (r"\bin the following\b", "below", "Shortened verbose constructions", True),
    (r"\bas follows\b", ":", "Shortened verbose constructions", True),
    (r"\bthe following\b", "this", "Shortened verbose constructions", True),
    (r"\bas mentioned (?:above|before|previously)\b", "(above)", "Shortened verbose constructions", True),
    (r"\bmore specifically\b", "specifically", "Shortened verbose constructions", True),
    (r"\bmore importantly\b", "importantly", "Shortened verbose constructions", True),
    (r"\bin summary\b", "summary:", "Shortened verbose constructions", True),
    (r"\bin conclusion\b", "conclusion:", "Shortened verbose constructions", True),
    (r"\bto summarize\b", "summary:", "Shortened verbose constructions", True),
    (r"\bto conclude\b", "conclusion:", "Shortened verbose constructions", True),
    # Wordy question openings
    (r"\bCan you (?:please )?(?:help me )?\b", "", "Removed filler phrases", True),
    (r"\bI (?:am|'m) (?:looking|trying) to\b", "", "Removed filler phrases", True),
    (r"\bI (?:am|'m) (?:interested|hoping) (?:to|in)\b", "", "Removed filler phrases", True),
    (r"\bI (?:am|'m) (?:writing|reaching out) to\b", "", "Removed filler phrases", True),
    # Remove double spaces that accumulate
    (r"  +", " ", "Cleaned up whitespace", False),
    (r" ,", ",", "Cleaned up whitespace", False),
    (r" \.", ".", "Cleaned up whitespace", False),
]


def _apply_patterns(
    text: str,
    patterns: List[Tuple[str, str, str, bool]],
    technique_counts: dict,
) -> str:
    """Apply a list of (pattern, replacement, technique, ignore_case) tuples."""
    for pattern, replacement, technique, ignore_case in patterns:
        flags = re.IGNORECASE if ignore_case else 0
        new_text, n = re.subn(pattern, replacement, text, flags=flags)
        if n > 0:
            technique_counts[technique] = technique_counts.get(technique, 0) + n
            text = new_text
    return text


def _normalise_whitespace(text: str) -> str:
    """Collapse multiple spaces / blank lines left by removals."""
    # Multiple spaces → one
    text = re.sub(r"[ \t]+", " ", text)
    # More than two consecutive newlines → two
    text = re.sub(r"\n{3,}", "\n\n", text)
    # Trim trailing spaces on each line
    text = "\n".join(line.rstrip() for line in text.splitlines())
    return text.strip()


def rule_based_compress(
    text: str, level: CompressionLevel
) -> Tuple[str, List[str]]:
    """
    Apply deterministic rule-based compression.
    Returns (compressed_text, list_of_technique_descriptions).
    """
    technique_counts: dict = {}

    # Gentle: only the safest, most unambiguous substitutions
    text = _apply_patterns(text, _FILLER_PATTERNS, technique_counts)

    if level in (CompressionLevel.BALANCED, CompressionLevel.AGGRESSIVE):
        text = _apply_patterns(text, _BALANCED_PATTERNS, technique_counts)

    if level == CompressionLevel.AGGRESSIVE:
        text = _apply_patterns(text, _AGGRESSIVE_PATTERNS, technique_counts)

    text = _normalise_whitespace(text)

    techniques = [
        f"{desc} ({count})" for desc, count in technique_counts.items() if count > 0
    ]
    return text, techniques


# ---------------------------------------------------------------------------
# API-based compression
# ---------------------------------------------------------------------------

_META_PROMPT = """\
Rewrite this prompt to use fewer tokens while preserving its exact meaning and intent.
Remove filler words, redundancy, and verbose constructions.
Return ONLY the compressed prompt, no explanation.
Compression level: {level}

PROMPT TO COMPRESS:
{user_prompt}"""


def api_compress(
    text: str,
    level: CompressionLevel,
    api_key: str,
) -> str:
    """Call Claude claude-haiku-4-5-20251001 to compress the prompt. Returns compressed text."""
    try:
        import anthropic
    except ImportError as exc:
        raise ImportError(
            "The 'anthropic' package is required for API compression. "
            "Install it with: pip install anthropic"
        ) from exc

    client = anthropic.Anthropic(api_key=api_key)
    meta_prompt = _META_PROMPT.format(level=level.value, user_prompt=text)

    message = client.messages.create(
        model="claude-haiku-4-5-20251001",
        max_tokens=4096,
        messages=[{"role": "user", "content": meta_prompt}],
    )
    return message.content[0].text.strip()


# ---------------------------------------------------------------------------
# Public interface
# ---------------------------------------------------------------------------

MIN_TOKENS_TO_COMPRESS = 50


class Compressor:
    """
    Main compressor.  If ``api_key`` is provided, uses Claude API for
    intelligent compression.  Otherwise falls back to rule-based compression.
    """

    def __init__(
        self,
        api_key: Optional[str] = None,
        level: CompressionLevel = CompressionLevel.BALANCED,
    ) -> None:
        self.api_key = api_key
        self.level = level

    def compress(self, text: str) -> CompressionResult:
        original_tokens = count_tokens(text)

        if original_tokens < MIN_TOKENS_TO_COMPRESS:
            return CompressionResult(
                original_text=text,
                compressed_text=text,
                original_tokens=original_tokens,
                compressed_tokens=original_tokens,
                techniques_applied=["Already under 50 tokens — nothing to compress"],
                used_api=False,
            )

        used_api = False
        techniques: List[str] = []

        if self.api_key:
            try:
                compressed = api_compress(text, self.level, self.api_key)
                used_api = True
                # After API compression, still run cheap whitespace normalisation
                compressed = _normalise_whitespace(compressed)
                # Count techniques loosely
                techniques = ["AI-powered compression (Claude claude-haiku-4-5-20251001)"]
            except Exception as exc:  # pragma: no cover
                # Gracefully degrade to rule-based on any API error
                compressed, techniques = rule_based_compress(text, self.level)
                techniques.insert(
                    0, f"API unavailable ({exc.__class__.__name__}), used rule-based fallback"
                )
        else:
            compressed, techniques = rule_based_compress(text, self.level)

        compressed_tokens = count_tokens(compressed)

        return CompressionResult(
            original_text=text,
            compressed_text=compressed,
            original_tokens=original_tokens,
            compressed_tokens=compressed_tokens,
            techniques_applied=techniques,
            used_api=used_api,
        )
