"""Context compression — tiered hot/warm/cold within a token budget.

Builds a prompt-ready context string from scored facts, events, clusters,
and meta-observations. Three tiers with configurable budget ratios:

- **Hot (Tier 1)**: Top facts in full format (core memory).
- **Warm (Tier 2)**: Remaining facts grouped by entity, compact format.
- **Cold (Tier 3)**: Cluster summaries, meta-observations, events.
"""

from __future__ import annotations

import logging
from collections import defaultdict
from dataclasses import dataclass
from datetime import UTC, datetime
from typing import Any

from arandu.config import MemoryConfig
from arandu.read._helpers import is_fact_stale

logger = logging.getLogger(__name__)

__all__ = ["CompressedContext", "compress_broad_context", "compress_context"]


# ---------------------------------------------------------------------------
# Token estimation
# ---------------------------------------------------------------------------


def _est_tokens(text: str) -> int:
    """Heuristic: ~4 chars per token."""
    return max(0, len(text) // 4)


# ---------------------------------------------------------------------------
# Result dataclass
# ---------------------------------------------------------------------------


@dataclass
class CompressedContext:
    """Result of context compression.

    Attributes:
        context_text: Final prompt-ready context string.
        hot_count: Number of facts in hot tier (Tier 1).
        warm_count: Number of facts in warm tier (Tier 2).
        cold_count: Number of items in cold tier (Tier 3).
        total_tokens: Estimated token count of context_text.
    """

    context_text: str = ""
    hot_count: int = 0
    warm_count: int = 0
    cold_count: int = 0
    total_tokens: int = 0


# ---------------------------------------------------------------------------
# Fact formatting helpers
# ---------------------------------------------------------------------------


def _fact_line_full(entry: dict[str, Any], stale: bool) -> str:
    """Format a fact entry in full (Tier 1) format."""
    entity = entry.get("entity", "")
    attribute = entry.get("attribute", "")
    value = entry.get("value", "")
    date = entry.get("date", "")
    prefix = "[STALE] " if stale else ""
    return f"  • {prefix}[{entity}] {attribute}: {value} (since {date})"


def _fact_line_compact(entry: dict[str, Any]) -> str:
    """Format a fact entry in compact (Tier 2) format."""
    attribute = entry.get("attribute", "")
    value = entry.get("value", "")
    return f"{attribute}: {value}"


# ---------------------------------------------------------------------------
# Main compression
# ---------------------------------------------------------------------------


async def compress_context(
    facts: list[dict[str, Any]],
    events: list[dict[str, Any]],
    config: MemoryConfig,
    *,
    clusters: list[Any] | None = None,
    meta_observations: list[Any] | None = None,
    stale_keys: set[str] | None = None,
    stale_threshold_days: int = 90,
    now: datetime | None = None,
) -> CompressedContext:
    """Build tiered context text within token budget.

    Tier 1 (hot): Top 5 facts in full format.
    Tier 2 (warm): Remaining facts grouped by entity, compact.
    Tier 3 (cold): Cluster summaries, meta-observations, events.

    Args:
        facts: Scored fact dicts (must have ``score``, ``fact``, ``entity``,
            ``attribute``, ``value``, ``date`` keys).
        events: Event dicts with ``date`` and ``text`` keys.
        config: Memory configuration with token budget and tier ratios.
        clusters: Optional cluster objects with ``label`` and ``summary_text``.
        meta_observations: Optional meta-observation objects.
        stale_keys: Attribute keys that are considered always-stale.
        stale_threshold_days: Days after which a fact is stale.
        now: Current timestamp (defaults to UTC now).

    Returns:
        CompressedContext with tiered context text.
    """
    if now is None:
        now = datetime.now(UTC)

    token_budget = config.context_max_tokens
    budget_hot = int(token_budget * config.hot_tier_ratio)
    budget_warm = int(token_budget * config.warm_tier_ratio)
    budget_cold = token_budget - budget_hot - budget_warm

    # Sort by score descending
    all_facts = sorted(facts, key=lambda x: x.get("score", 0.0), reverse=True)

    lines: list[str] = []
    tokens_used = 0
    hot_count = 0
    warm_count = 0
    cold_count = 0

    # -- Tier 1: Hot (top 5, full format) --
    section1: list[str] = []
    has_stale = False
    for entry in all_facts[:5]:
        fact_obj = entry.get("fact")
        stale = False
        if fact_obj is not None:
            stale = is_fact_stale(
                fact_obj,
                now,
                stale_keys=stale_keys,
                stale_threshold_days=stale_threshold_days,
            )
            if stale:
                has_stale = True
        line = _fact_line_full(entry, stale)
        cost = _est_tokens(line)
        if tokens_used + cost > budget_hot and section1:
            break
        section1.append(line)
        tokens_used += cost
        hot_count += 1

    if section1:
        lines.append("═══ CORE MEMORY (assert confidently) ═══")
        lines.extend(section1)
        if has_stale:
            lines.append("  ⚠ Facts marked [STALE] should be verified before asserting.")
        lines.append("")

    # -- Tier 2: Warm (remaining facts, grouped by entity) --
    warm_facts = all_facts[5:]
    by_entity: dict[str, list[dict[str, Any]]] = defaultdict(list)
    for entry in warm_facts:
        entity = entry.get("entity", "unknown")
        by_entity[entity].append(entry)

    section2: list[str] = []
    for entity_key, entries in list(by_entity.items())[:15]:
        parts = [_fact_line_compact(e) for e in entries[:8]]
        if parts:
            line = f"  • [{entity_key}] " + " | ".join(parts)
            cost = _est_tokens(line)
            if tokens_used + cost > budget_hot + budget_warm:
                break
            section2.append(line)
            tokens_used += cost
            warm_count += len(entries[:8])

    if section2:
        lines.append("═══ EXTENDED CONTEXT (lower certainty) ═══")
        lines.extend(section2)
        lines.append("")

    # -- Tier 3: Cold (clusters + meta_observations + events) --
    section3: list[str] = []
    total_budget = budget_hot + budget_warm + budget_cold

    for c in (clusters or [])[:3]:
        summary = getattr(c, "summary_text", None) or ""
        label = getattr(c, "label", "?")
        if summary:
            short = summary[:200] + ("..." if len(summary) > 200 else "")
            line = f"  • [cluster:{label}] {short}"
            cost = _est_tokens(line)
            if tokens_used + cost > total_budget:
                break
            section3.append(line)
            tokens_used += cost
            cold_count += 1

    for o in (meta_observations or [])[:3]:
        otype = getattr(o, "observation_type", "pattern")
        title = getattr(o, "title", "")
        conf = getattr(o, "confidence", 0.7) or 0.7
        times = getattr(o, "times_reinforced", 1) or 1
        line = f"  • [{otype}] {title} — detected {times}x (confidence: {conf:.2f})"
        cost = _est_tokens(line)
        if tokens_used + cost > total_budget:
            break
        section3.append(line)
        tokens_used += cost
        cold_count += 1

    if section3:
        lines.append("═══ PATTERNS & INSIGHTS ═══")
        lines.extend(section3)
        lines.append("")

    # Events (max 5, 80 char snippet)
    event_lines: list[str] = []
    for e in (events or [])[:5]:
        date = e.get("date", "")
        text = e.get("text", "")
        if not text and e.get("event"):
            ev = e["event"]
            text = getattr(ev, "text", "") or ""
            if not date and getattr(ev, "occurred_at", None):
                occ = ev.occurred_at
                date = occ.strftime("%Y-%m-%d") if hasattr(occ, "strftime") else str(occ)[:10]
        text = (text or "")[:80]
        ellipsis = "..." if len(text) >= 80 else ""
        line = f"  • ({date}) {text}{ellipsis}"
        cost = _est_tokens(line)
        if tokens_used + cost > total_budget + 200:
            break
        event_lines.append(line)
        tokens_used += cost
        cold_count += 1

    if event_lines:
        lines.append("═══ RELEVANT_EVENTS (history — do NOT assert as current truth) ═══")
        lines.extend(event_lines)

    context_text = "\n".join(lines)
    total_tokens = _est_tokens(context_text)

    return CompressedContext(
        context_text=context_text,
        hot_count=hot_count,
        warm_count=warm_count,
        cold_count=cold_count,
        total_tokens=total_tokens,
    )


# ---------------------------------------------------------------------------
# Broad cluster-first compression
# ---------------------------------------------------------------------------


async def compress_broad_context(
    cluster_facts: dict[str, list[dict[str, Any]]],
    clusters: list[Any],
    config: MemoryConfig,
    *,
    meta_observations: list[Any] | None = None,
    events: list[dict[str, Any]] | None = None,
) -> CompressedContext:
    """Build context for broad queries using clusters as primary unit.

    Tier 1 (50%): Cluster summaries + representative facts inline.
    Tier 2 (25%): Meta-observations (top 5).
    Tier 3 (25%): Recent events (top 5).

    Args:
        cluster_facts: Mapping of cluster_label → fact dicts.
        clusters: Cluster objects with ``label``, ``summary_text``, ``fact_count``.
        config: Memory configuration.
        meta_observations: Optional meta-observation objects.
        events: Optional event dicts.

    Returns:
        CompressedContext with cluster-first context text.
    """
    token_budget = config.context_max_tokens
    budget_t1 = int(token_budget * 0.50)
    budget_t2 = int(token_budget * 0.25)
    budget_t3 = int(token_budget * 0.25)

    lines: list[str] = []
    tokens_used = 0
    hot_count = 0
    warm_count = 0
    cold_count = 0

    # -- Tier 1: Cluster summaries + inline facts --
    section1: list[str] = []
    sorted_clusters = sorted(
        clusters,
        key=lambda c: getattr(c, "fact_count", 0) or 0,
        reverse=True,
    )

    for cluster in sorted_clusters:
        label = getattr(cluster, "label", "?")
        summary = getattr(cluster, "summary_text", "") or ""
        summary_short = summary[:200] + ("..." if len(summary) > 200 else "")

        facts_for_cluster = cluster_facts.get(label, [])
        fact_parts: list[str] = []
        for fe in facts_for_cluster:
            attr = fe.get("attribute", "")
            val = fe.get("value", "")
            if attr and val:
                fact_parts.append(f"{attr}: {val}")

        inline = " | ".join(fact_parts) if fact_parts else ""
        line = f"  • [{label}] {summary_short}" + (f" ({inline})" if inline else "")
        cost = _est_tokens(line)
        if tokens_used + cost > budget_t1 and section1:
            break
        section1.append(line)
        tokens_used += cost
        hot_count += 1

    if section1:
        lines.append("═══ BRAIN OVERVIEW ═══")
        lines.extend(section1)
        lines.append("")

    # -- Tier 2: Meta-observations --
    section2: list[str] = []
    for o in (meta_observations or [])[:5]:
        otype = getattr(o, "observation_type", "pattern")
        title = getattr(o, "title", "")
        conf = getattr(o, "confidence", 0.7) or 0.7
        times = getattr(o, "times_reinforced", 1) or 1
        line = f"  • [{otype}] {title} — detected {times}x (confidence: {conf:.2f})"
        cost = _est_tokens(line)
        if tokens_used + cost > budget_t1 + budget_t2:
            break
        section2.append(line)
        tokens_used += cost
        warm_count += 1

    if section2:
        lines.append("═══ PATTERNS & INSIGHTS ═══")
        lines.extend(section2)
        lines.append("")

    # -- Tier 3: Events --
    event_lines: list[str] = []
    for e in (events or [])[:5]:
        date = e.get("date", "")
        text = e.get("text", "")
        if not text and e.get("event"):
            ev = e["event"]
            text = getattr(ev, "text", "") or ""
            if not date and getattr(ev, "occurred_at", None):
                occ = ev.occurred_at
                date = occ.strftime("%Y-%m-%d") if hasattr(occ, "strftime") else str(occ)[:10]
        text = (text or "")[:80]
        ellipsis = "..." if len(text) >= 80 else ""
        line = f"  • ({date}) {text}{ellipsis}"
        cost = _est_tokens(line)
        if tokens_used + cost > budget_t1 + budget_t2 + budget_t3:
            break
        event_lines.append(line)
        tokens_used += cost
        cold_count += 1

    if event_lines:
        lines.append("═══ RELEVANT_EVENTS (history — do NOT assert as current truth) ═══")
        lines.extend(event_lines)

    context_text = "\n".join(lines)
    total_tokens = _est_tokens(context_text)

    return CompressedContext(
        context_text=context_text,
        hot_count=hot_count,
        warm_count=warm_count,
        cold_count=cold_count,
        total_tokens=total_tokens,
    )
