"""
cortex/hippo/compaction/engine.py — CompactionEngine class (orchestrator).

Mirrors lossless-claw CompactionEngine with async aiosqlite ops.
"""

from __future__ import annotations

import logging
import math
from typing import Awaitable, Callable, Optional

from cortex.hippo.budget import CompactionBudget
from cortex.hippo.config import HippoConfig
from cortex.hippo.models import ContextItemRecord
from cortex.hippo.store.conversation_store import ConversationStore
from cortex.hippo.store.summary_store import SummaryStore

from .condensed import (
    _condensed_pass,
    _select_oldest_chunk_at_depth,
    _select_shallowest_condensation_candidate,
)
from .events import _persist_compaction_events
from .leaf import (
    _check_leaf_trigger,
    _leaf_pass,
    _resolve_fresh_tail_ordinal,
    _resolve_prior_leaf_summary_context,
    _select_oldest_leaf_chunk,
)
from .types import CompactionDecision, CompactionMetrics, CompactionResult, PassResult, SummarizeFn
from .utils import (
    _get_item_tokens,
    _get_summary_by_context_item,
    _resolve_owner_id,
    _summarize_with_escalation,
)

logger = logging.getLogger(__name__)


class CompactionEngine:
    """Two-phase compaction engine: leaf pass + condensation.

    Mirrors lossless-claw CompactionEngine with async aiosqlite ops.
    """

    def __init__(
        self,
        conversation_store: ConversationStore,
        summary_store: SummaryStore,
        config: HippoConfig,
        budget: Optional[CompactionBudget] = None,
        on_compaction_metrics: Optional[Callable] = None,
        heartbeat_fn: Optional[Callable[[int], Awaitable[None]]] = None,
    ) -> None:
        self.conversation_store = conversation_store
        self.summary_store = summary_store
        self.config = config
        self.budget = budget
        self._on_compaction_metrics = on_compaction_metrics
        self._last_compaction_metrics: Optional[CompactionMetrics] = None
        self._heartbeat_fn = heartbeat_fn

    async def _touch_heartbeat(self, conversation_id: int) -> None:
        """Call the heartbeat callback if configured (best-effort)."""
        if self._heartbeat_fn is not None:
            try:
                await self._heartbeat_fn(conversation_id)
            except Exception:
                logger.debug(
                    "heartbeat callback failed for conv=%d", conversation_id,
                    exc_info=True,
                )

    # ── Public: evaluate ──────────────────────────────────────────────

    async def evaluate(
        self,
        conversation_id: int,
        token_budget: int,
        observed_token_count: Optional[int] = None,
    ) -> CompactionDecision:
        """Check if compaction threshold is exceeded."""
        stored_tokens = await self.summary_store.get_context_token_count(conversation_id)
        current_tokens = max(stored_tokens, observed_token_count or 0)
        threshold = math.floor(self.config.context_threshold * token_budget)

        should = current_tokens > threshold
        return CompactionDecision(
            should_compact=should,
            reason=(
                f"tokens {current_tokens} > threshold {threshold}"
                if should
                else f"tokens {current_tokens} <= threshold {threshold}"
            ),
            current_tokens=current_tokens,
            threshold=threshold,
        )

    # ── Private: metrics ──────────────────────────────────────────────

    def _emit_metrics(self, metrics: CompactionMetrics) -> None:
        import json as _json
        self._last_compaction_metrics = metrics
        logger.info("compaction_metrics: %s", _json.dumps(metrics.to_dict()))
        if self._on_compaction_metrics:
            try:
                self._on_compaction_metrics(metrics)
            except Exception:
                logger.warning("metrics callback failed", exc_info=True)

    # ── Public: compact ───────────────────────────────────────────────

    async def compact_leaf(
        self,
        conversation_id: int,
        token_budget: int,
        summarize: SummarizeFn,
        force: bool = False,
        previous_summary_content: Optional[str] = None,
    ) -> CompactionResult:
        """Run one incremental leaf compaction pass."""
        import time as _time
        _t0 = _time.monotonic()
        tokens_before = await self.summary_store.get_context_token_count(conversation_id)

        if not force:
            decision = await self.evaluate(conversation_id, token_budget)
            if not decision.should_compact:
                return CompactionResult(
                    action_taken=False,
                    tokens_before=tokens_before,
                    tokens_after=tokens_before,
                )

        chunk = await _select_oldest_leaf_chunk(
            conversation_id, self.conversation_store, self.summary_store, self.config
        )
        if not chunk.items:
            return CompactionResult(
                action_taken=False,
                tokens_before=tokens_before,
                tokens_after=tokens_before,
            )

        prior_context = previous_summary_content
        if prior_context is None:
            prior_context = await _resolve_prior_leaf_summary_context(
                conversation_id, chunk.items, self.summary_store
            )

        pass_result = await _leaf_pass(
            conversation_id, chunk.items, summarize, prior_context,
            self.conversation_store, self.summary_store, self.config,
        )

        # B-006: heartbeat after leaf pass to signal liveness
        await self._touch_heartbeat(conversation_id)

        tokens_after_leaf = await self.summary_store.get_context_token_count(conversation_id)

        # Persist compaction event for leaf pass
        owner_id = await _resolve_owner_id(conversation_id, self.conversation_store)
        await _persist_compaction_events(
            conversation_id=conversation_id,
            owner_id=owner_id,
            tokens_before=tokens_before,
            tokens_after_leaf=tokens_after_leaf,
            tokens_after_final=tokens_after_leaf,
            conversation_store=self.conversation_store,
            leaf_result=pass_result,
        )

        created_ids = [pass_result.summary_id]
        messages_compacted = len(chunk.items)
        tokens_after = tokens_after_leaf
        did_condensed = False
        rounds = 1

        # ── Incremental condensed passes (best-effort) ────────────
        # After a leaf pass, check if enough same-depth summaries exist
        # to trigger condensation. This piggybacks on leaf compaction
        # rather than waiting for a full sweep.
        incremental_max_depth = self.config.incremental_condensed_max_depth
        if incremental_max_depth != 0:
            max_depth = (
                incremental_max_depth
                if incremental_max_depth > 0
                else self.config.max_rounds  # unlimited (-1) → cap at max_rounds
            )
            try:
                for target_depth in range(max_depth):
                    chunk = await _select_oldest_chunk_at_depth(
                        conversation_id, target_depth,
                        self.summary_store, self.config,
                    )
                    fanout = (
                        self.config.leaf_min_fanout
                        if target_depth == 0
                        else self.config.condensed_min_fanout
                    )
                    if (
                        len(chunk.items) < fanout
                        or chunk.summary_tokens < self.config.condensed_min_chunk_tokens
                    ):
                        break

                    t_before = await self.summary_store.get_context_token_count(
                        conversation_id
                    )
                    condense_result = await _condensed_pass(
                        conversation_id,
                        chunk.items,
                        target_depth,
                        summarize,
                        self.summary_store,
                    )
                    t_after = await self.summary_store.get_context_token_count(
                        conversation_id
                    )

                    # Persist condensed event
                    await _persist_compaction_events(
                        conversation_id=conversation_id,
                        owner_id=owner_id,
                        tokens_before=t_before,
                        tokens_after_leaf=t_before,
                        tokens_after_final=t_after,
                        conversation_store=self.conversation_store,
                        condense_result=condense_result,
                    )

                    created_ids.append(condense_result.summary_id)
                    tokens_after = t_after
                    did_condensed = True
                    rounds += 1

                    # B-006: heartbeat after each condensed pass
                    await self._touch_heartbeat(conversation_id)

                    if t_after >= t_before:
                        break  # No progress
            except Exception:
                # Condensed pass failure must NOT fail the leaf pass
                logger.debug(
                    "Incremental condensed pass failed (conv=%s), continuing",
                    conversation_id,
                    exc_info=True,
                )

        result = CompactionResult(
            action_taken=True,
            tokens_before=tokens_before,
            tokens_after=tokens_after,
            created_summary_ids=created_ids,
            messages_compacted=messages_compacted,
            condensed=did_condensed,
            rounds=rounds,
        )
        self._emit_metrics(CompactionMetrics(
            messages_processed=messages_compacted,
            summaries_created=len(created_ids),
            tokens_saved=max(0, tokens_before - tokens_after),
            rounds_executed=rounds,
            duration_ms=(_time.monotonic() - _t0) * 1000,
            mode="LEAF",
        ))
        return result

    async def compact_full_sweep(
        self,
        conversation_id: int,
        token_budget: int,
        summarize: SummarizeFn,
        force: bool = False,
        hard_trigger: bool = False,
    ) -> CompactionResult:
        """Full sweep: Phase 1 (all leaf chunks) + Phase 2 (condensation)."""
        import time as _time
        _t0 = _time.monotonic()
        tokens_before = await self.summary_store.get_context_token_count(conversation_id)

        if not force:
            decision = await self.evaluate(conversation_id, token_budget)
            leaf_trigger = await _check_leaf_trigger(
                conversation_id, self.conversation_store, self.summary_store, self.config
            )
            if not decision.should_compact and not leaf_trigger:
                return CompactionResult(
                    action_taken=False,
                    tokens_before=tokens_before,
                    tokens_after=tokens_before,
                )

        created_ids: list[str] = []
        rounds = 0
        did_condensed = False
        last_leaf_result: Optional[PassResult] = None
        last_condense_result: Optional[PassResult] = None

        # Phase 1: Leaf passes (convert raw messages to leaf summaries)
        for _ in range(self.config.max_rounds):
            # Budget gate
            if self.budget is not None:
                bd = self.budget.check_budget(str(conversation_id))
                if not bd.allowed:
                    logger.warning(
                        "compact_full_sweep: budget check failed in leaf phase (%s)",
                        bd.reason,
                    )
                    break

            chunk = await _select_oldest_leaf_chunk(
                conversation_id, self.conversation_store, self.summary_store, self.config
            )
            if not chunk.items:
                break

            prior_context = await _resolve_prior_leaf_summary_context(
                conversation_id, chunk.items, self.summary_store
            )

            t_before = await self.summary_store.get_context_token_count(conversation_id)
            pass_result = await _leaf_pass(
                conversation_id, chunk.items, summarize, prior_context,
                self.conversation_store, self.summary_store, self.config,
            )
            t_after = await self.summary_store.get_context_token_count(conversation_id)

            # Record token usage from summarization
            if self.budget is not None:
                self.budget.record_summarization(
                    str(conversation_id),
                    pass_result.tokens_before,
                    pass_result.tokens_after,
                )

            created_ids.append(pass_result.summary_id)
            rounds += 1
            last_leaf_result = pass_result

            # B-006: heartbeat after each leaf pass
            await self._touch_heartbeat(conversation_id)

            if t_after >= t_before:
                # No progress — stop leaf passes
                break

        tokens_after_leaf = await self.summary_store.get_context_token_count(conversation_id)

        # Phase 2: Condensed passes (merge same-depth summaries)
        for _ in range(self.config.max_rounds):
            # Budget gate
            if self.budget is not None:
                bd = self.budget.check_budget(str(conversation_id))
                if not bd.allowed:
                    logger.warning(
                        "compact_full_sweep: budget check failed in condensed phase (%s)",
                        bd.reason,
                    )
                    break

            candidate = await _select_shallowest_condensation_candidate(
                conversation_id, hard_trigger, self.summary_store, self.config
            )
            if candidate is None:
                break

            t_before = await self.summary_store.get_context_token_count(conversation_id)
            pass_result = await _condensed_pass(
                conversation_id,
                candidate.chunk.items,
                candidate.target_depth,
                summarize,
                self.summary_store,
            )
            t_after = await self.summary_store.get_context_token_count(conversation_id)

            # Record token usage from summarization
            if self.budget is not None:
                self.budget.record_summarization(
                    str(conversation_id),
                    pass_result.tokens_before,
                    pass_result.tokens_after,
                )

            created_ids.append(pass_result.summary_id)
            rounds += 1
            did_condensed = True
            last_condense_result = pass_result

            # B-006: heartbeat after each condensed pass
            await self._touch_heartbeat(conversation_id)

            if t_after >= t_before:
                break

        tokens_after = await self.summary_store.get_context_token_count(conversation_id)

        # Persist compaction events
        if last_leaf_result or last_condense_result:
            owner_id = await _resolve_owner_id(conversation_id, self.conversation_store)
            await _persist_compaction_events(
                conversation_id=conversation_id,
                owner_id=owner_id,
                tokens_before=tokens_before,
                tokens_after_leaf=tokens_after_leaf,
                tokens_after_final=tokens_after,
                conversation_store=self.conversation_store,
                leaf_result=last_leaf_result,
                condense_result=last_condense_result,
            )

        full_result = CompactionResult(
            action_taken=rounds > 0,
            tokens_before=tokens_before,
            tokens_after=tokens_after,
            created_summary_ids=created_ids,
            condensed=did_condensed,
            rounds=rounds,
        )
        if full_result.action_taken:
            self._emit_metrics(CompactionMetrics(
                summaries_created=len(created_ids),
                tokens_saved=max(0, tokens_before - tokens_after),
                rounds_executed=rounds,
                duration_ms=(_time.monotonic() - _t0) * 1000,
                mode="FULL",
            ))
        return full_result

    async def compact_until_under(
        self,
        conversation_id: int,
        token_budget: int,
        target_tokens: Optional[int] = None,
        current_tokens: Optional[int] = None,
        summarize: SummarizeFn = None,
    ) -> CompactionResult:
        """Compact until under target, up to max_rounds."""
        target = target_tokens or math.floor(self.config.context_threshold * token_budget)
        total_rounds = 0
        did_condensed = False
        all_ids: list[str] = []
        tokens_start = current_tokens or await self.summary_store.get_context_token_count(
            conversation_id
        )

        for _ in range(self.config.max_rounds):
            t = await self.summary_store.get_context_token_count(conversation_id)
            if t <= target:
                break

            # Budget gate
            if self.budget is not None:
                session_id = str(conversation_id)
                decision = self.budget.check_budget(session_id)
                if not decision.allowed:
                    logger.warning(
                        "compact_until_under: budget check failed (%s), stopping",
                        decision.reason,
                    )
                    break

            result = await self.compact_full_sweep(
                conversation_id, token_budget, summarize, force=True, hard_trigger=True
            )
            total_rounds += result.rounds
            all_ids.extend(result.created_summary_ids)
            if result.condensed:
                did_condensed = True

            # Record reduction ratio for circuit breaker
            if self.budget is not None and result.tokens_before > 0:
                ratio = (
                    (result.tokens_before - result.tokens_after) / result.tokens_before
                    if result.tokens_before > 0
                    else 0.0
                )
                self.budget.record_compaction_result(str(conversation_id), ratio)

            # B-006: heartbeat after each sweep round
            await self._touch_heartbeat(conversation_id)

            if not result.action_taken:
                break

        tokens_final = await self.summary_store.get_context_token_count(conversation_id)
        return CompactionResult(
            action_taken=total_rounds > 0,
            tokens_before=tokens_start,
            tokens_after=tokens_final,
            created_summary_ids=all_ids,
            condensed=did_condensed,
            rounds=total_rounds,
        )

    # ── Backward-compat private methods (delegate to module functions) ─

    async def _select_oldest_leaf_chunk(self, conversation_id: int):
        return await _select_oldest_leaf_chunk(
            conversation_id, self.conversation_store, self.summary_store, self.config
        )

    async def _select_oldest_chunk_at_depth(self, conversation_id, target_depth, fresh_tail_ordinal=None):
        return await _select_oldest_chunk_at_depth(
            conversation_id, target_depth, self.summary_store, self.config, fresh_tail_ordinal
        )

    async def _select_shallowest_condensation_candidate(self, conversation_id, hard_trigger):
        return await _select_shallowest_condensation_candidate(
            conversation_id, hard_trigger, self.summary_store, self.config
        )

    async def _leaf_pass(self, conversation_id, message_items, summarize, previous_summary_content=None):
        return await _leaf_pass(
            conversation_id, message_items, summarize, previous_summary_content,
            self.conversation_store, self.summary_store, self.config,
        )

    async def _condensed_pass(self, conversation_id, summary_items, target_depth, summarize):
        return await _condensed_pass(
            conversation_id, summary_items, target_depth, summarize, self.summary_store,
        )

    async def _summarize_with_escalation(self, source_text, summarize, options=None):
        return await _summarize_with_escalation(source_text, summarize, options)

    def _resolve_fresh_tail_ordinal(self, context_items):
        return _resolve_fresh_tail_ordinal(context_items, self.config)

    async def _check_leaf_trigger(self, conversation_id):
        return await _check_leaf_trigger(
            conversation_id, self.conversation_store, self.summary_store, self.config
        )

    async def _get_item_tokens(self, item):
        return await _get_item_tokens(item, self.conversation_store, self.summary_store)

    async def _get_summary_by_context_item(self, item):
        return await _get_summary_by_context_item(item, self.summary_store)

    async def _persist_compaction_events(self, **kwargs):
        return await _persist_compaction_events(
            conversation_store=self.conversation_store, **kwargs
        )

    async def _persist_compaction_event(self, **kwargs):
        from .events import _persist_compaction_event
        return await _persist_compaction_event(
            conversation_store=self.conversation_store, **kwargs
        )

    async def _resolve_owner_id(self, conversation_id):
        return await _resolve_owner_id(conversation_id, self.conversation_store)

    async def _resolve_prior_leaf_summary_context(self, conversation_id, message_items):
        return await _resolve_prior_leaf_summary_context(
            conversation_id, message_items, self.summary_store
        )
