"""
cortex/hippo/compaction/hardening.py — Edge-case hardening for the compaction pipeline.

Handles:
- Empty/blank message detection and graceful skip
- Very long message chunking (>100K tokens) before compaction
- Concurrent compaction guard (per-conversation DB flag + asyncio lock)
- Partial failure recovery (in-progress flag cleared on crash)
- Metrics counters for skipped/chunked/failed compactions

Issue: hippo-hardening (edge cases for empty/long/concurrent compaction)
"""

from __future__ import annotations

import asyncio
import logging
import time
from dataclasses import dataclass, field
from typing import Any, Optional

logger = logging.getLogger(__name__)

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

# Messages whose token count exceeds this threshold are chunked before
# being handed to the compaction pipeline.
LONG_MESSAGE_CHUNK_TOKEN_THRESHOLD = 100_000
# Size of each chunk when splitting a long message (in tokens ~ chars/4).
LONG_MESSAGE_CHUNK_SIZE = 50_000

# ── Metrics counters ──────────────────────────────────────────────────────────


@dataclass
class CompactionEdgeMetrics:
    """Counters for edge-case events in the compaction pipeline.

    These are in-process counters only — not persisted across restarts.
    """
    skipped_empty: int = 0          # messages skipped because content was empty
    chunked_long: int = 0           # messages split because they exceeded threshold
    failed_compactions: int = 0     # compaction runs that raised an exception
    concurrent_skips: int = 0       # compactions skipped because one was already in progress

    def as_dict(self) -> dict[str, int]:
        return {
            "skipped_empty": self.skipped_empty,
            "chunked_long": self.chunked_long,
            "failed_compactions": self.failed_compactions,
            "concurrent_skips": self.concurrent_skips,
        }


# Module-level singleton (reset in tests via reset_metrics())
_metrics = CompactionEdgeMetrics()


def get_metrics() -> CompactionEdgeMetrics:
    """Return the current module-level metrics snapshot."""
    return _metrics


def reset_metrics() -> None:
    """Reset all counters to zero (for testing)."""
    global _metrics
    _metrics = CompactionEdgeMetrics()


# ── Empty message detection ────────────────────────────────────────────────────


def is_empty_message(content: Any) -> bool:
    """Return True if a message carries no meaningful content.

    Handles plain strings, structured content arrays, and None.
    """
    if content is None:
        return True
    if isinstance(content, str):
        return not content.strip()
    if isinstance(content, list):
        # Only consider it non-empty if at least one text/tool block has content
        for block in content:
            if not isinstance(block, dict):
                # Non-dict items (strings, ints, etc.) count as content
                # if they are truthy; skip empty/falsy ones
                if block:
                    return False
                continue
            btype = block.get("type", "")
            # Non-text blocks (tool_use, tool_result, etc.) count as content
            if btype not in ("text", "thinking", "reasoning", ""):
                return False
            text = block.get("text") or block.get("thinking") or block.get("reasoning")
            if isinstance(text, str) and text.strip():
                return False
        # All blocks were empty text/thinking blocks
        return True
    # dict or other — treat as non-empty
    return False


def check_and_skip_empty(role: str, content: Any) -> bool:
    """Return True if the message should be skipped (empty content).

    Side-effect: increments _metrics.skipped_empty and logs a warning.
    """
    if is_empty_message(content):
        _metrics.skipped_empty += 1
        logger.warning(
            "compaction.hardening: skipping empty message (role=%s)", role
        )
        return True
    return False


# ── Long message chunking ──────────────────────────────────────────────────────


def estimate_content_tokens(content: Any) -> int:
    """Rough token estimate for any content shape (~4 chars per token)."""
    if content is None:
        return 0
    if isinstance(content, str):
        return max(1, len(content) // 4)
    if isinstance(content, list):
        import json
        try:
            serialized = json.dumps(content)
            return max(1, len(serialized) // 4)
        except (TypeError, ValueError):
            return 0
    return 0


def chunk_long_text(
    text: str,
    chunk_size_tokens: int = LONG_MESSAGE_CHUNK_SIZE,
) -> list[str]:
    """Split a long text string into chunks of approximately chunk_size_tokens.

    Splits on paragraph / sentence boundaries where possible.
    """
    char_limit = chunk_size_tokens * 4
    if len(text) <= char_limit:
        return [text]

    chunks: list[str] = []
    remaining = text
    while len(remaining) > char_limit:
        # Try to split at a double-newline (paragraph boundary) near the limit
        split_at = remaining.rfind("\n\n", 0, char_limit)
        if split_at == -1:
            # Try single newline
            split_at = remaining.rfind("\n", 0, char_limit)
        if split_at == -1 or split_at < char_limit // 2:
            # Try to split on whitespace to avoid mid-word breaks
            split_at = remaining.rfind(" ", 0, char_limit)
        if split_at == -1 or split_at < char_limit // 2:
            # Hard split at char boundary (no whitespace found)
            split_at = char_limit

        chunks.append(remaining[:split_at].rstrip())
        remaining = remaining[split_at:].lstrip()

    if remaining:
        chunks.append(remaining)
    return chunks


def maybe_chunk_long_content(
    role: str,
    content: Any,
    threshold: int = LONG_MESSAGE_CHUNK_TOKEN_THRESHOLD,
    chunk_size: int = LONG_MESSAGE_CHUNK_SIZE,
) -> tuple[bool, list[Any]]:
    """Check if content exceeds the long-message threshold.

    Returns ``(was_chunked, list_of_content_pieces)``.
    - If not long: returns ``(False, [content])``.
    - If long and plain string: splits into text chunks, returns ``(True, [str, ...])``.
    - If long and structured (list): JSON-serializes and splits, returns ``(True, [str, ...])``.

    Side-effect: increments _metrics.chunked_long if splitting occurred.
    """
    token_count = estimate_content_tokens(content)
    if token_count <= threshold:
        return False, [content]

    # Need to chunk
    if isinstance(content, str):
        text = content
    else:
        import json
        try:
            text = json.dumps(content)
        except (TypeError, ValueError):
            text = str(content)

    chunks = chunk_long_text(text, chunk_size)
    _metrics.chunked_long += 1
    logger.warning(
        "compaction.hardening: long message detected (role=%s, ~%d tokens) — "
        "split into %d chunks",
        role, token_count, len(chunks),
    )
    return True, chunks


# ── Concurrent compaction guard ────────────────────────────────────────────────


@dataclass
class _ConversationLockState:
    lock: asyncio.Lock = field(default_factory=asyncio.Lock)
    in_progress: bool = False


class CompactionConcurrencyGuard:
    """Per-conversation lock to prevent concurrent compaction runs.

    Usage::

        guard = CompactionConcurrencyGuard()

        async with guard.acquire(conv_id) as acquired:
            if not acquired:
                # Already running — skip
                return
            # ... do compaction ...

    The guard is purely in-process (asyncio-based). For multi-process
    deployments, combine with a DB-level ``compaction_in_progress`` flag
    (see ``mark_compaction_start`` / ``mark_compaction_end`` helpers below).
    """

    def __init__(self) -> None:
        self._states: dict[int, _ConversationLockState] = {}

    def _get_state(self, conversation_id: int) -> _ConversationLockState:
        if conversation_id not in self._states:
            self._states[conversation_id] = _ConversationLockState()
        return self._states[conversation_id]

    def is_in_progress(self, conversation_id: int) -> bool:
        """Return True if compaction is currently in progress for this conversation."""
        state = self._states.get(conversation_id)
        return state is not None and state.in_progress

    class _AcquireContext:
        def __init__(self, guard: "CompactionConcurrencyGuard", conversation_id: int) -> None:
            self._guard = guard
            self._conversation_id = conversation_id
            self._acquired = False

        async def __aenter__(self) -> bool:
            state = self._guard._get_state(self._conversation_id)
            if state.in_progress:
                _metrics.concurrent_skips += 1
                logger.debug(
                    "compaction.hardening: compaction already in progress "
                    "for conversation %d — skipping",
                    self._conversation_id,
                )
                self._acquired = False
                return False
            # Try to acquire without waiting
            if not state.lock.locked():
                await state.lock.acquire()
                state.in_progress = True
                self._acquired = True
                return True
            # Lock is held — skip rather than queue
            _metrics.concurrent_skips += 1
            logger.debug(
                "compaction.hardening: lock contention for conversation %d — skipping",
                self._conversation_id,
            )
            self._acquired = False
            return False

        async def __aexit__(self, exc_type, exc_val, exc_tb) -> None:
            if self._acquired:
                state = self._guard._get_state(self._conversation_id)
                state.in_progress = False
                if state.lock.locked():
                    state.lock.release()

    def acquire(self, conversation_id: int) -> "_AcquireContext":
        """Return an async context manager that yields True if acquired."""
        return self._AcquireContext(self, conversation_id)


# B-019: Module-level default guard for safe_compact when none is provided
_default_guard = CompactionConcurrencyGuard()


# ── Partial failure recovery ──────────────────────────────────────────────────

# DB flag helpers — these are thin wrappers that set/clear a per-conversation
# ``compaction_in_progress`` flag in the hippo_conversations table.
# On startup / next compaction attempt, if the flag is stale (crash recovery),
# it is cleared automatically so the next run proceeds normally.


async def mark_compaction_start(
    conn: Any, conversation_id: int
) -> None:
    """Set compaction_in_progress=1 for the conversation in the DB.

    Requires the hippo_conversations table to have a ``compaction_in_progress``
    column (added by migration; we guard with a try/except for backward compat).
    """
    try:
        await conn.execute(
            """
            UPDATE hippo_conversations
               SET compaction_in_progress = 1,
                   compaction_started_at  = datetime('now')
             WHERE conversation_id = ?
            """,
            (conversation_id,),
        )
        await conn.commit()
    except Exception as exc:
        # Column may not exist in older schemas — log and continue
        logger.debug(
            "mark_compaction_start: could not set flag (conv=%d): %s",
            conversation_id, exc,
        )


async def mark_compaction_end(
    conn: Any, conversation_id: int
) -> None:
    """Clear compaction_in_progress for the conversation in the DB.

    B-005 fix: Re-raises DB errors after logging to avoid permanently
    blocking conversations. The caller's finally block ensures this is
    always attempted, and _run_with_db_flag catches/logs any exception.
    """
    try:
        await conn.execute(
            """
            UPDATE hippo_conversations
               SET compaction_in_progress = 0,
                   compaction_started_at  = NULL
             WHERE conversation_id = ?
            """,
            (conversation_id,),
        )
        await conn.commit()
    except Exception as exc:
        logger.error(
            "mark_compaction_end: could not clear flag (conv=%d): %s",
            conversation_id, exc,
        )
        raise


# B-006: Default stale recovery timeout increased from 5min to 30min.
# Configurable via the stale_after_seconds parameter.
DEFAULT_STALE_COMPACTION_TIMEOUT_SECONDS = 1800  # 30 minutes


async def touch_compaction_heartbeat(
    conn: Any, conversation_id: int
) -> None:
    """Update compaction_started_at to 'now' as a heartbeat.

    B-006 fix: Long-running compactions call this periodically to signal
    they are still alive, preventing stale recovery from clearing them.
    """
    try:
        await conn.execute(
            """
            UPDATE hippo_conversations
               SET compaction_started_at = datetime('now')
             WHERE conversation_id = ?
               AND compaction_in_progress = 1
            """,
            (conversation_id,),
        )
        await conn.commit()
    except Exception as exc:
        logger.debug(
            "touch_compaction_heartbeat: failed (conv=%d): %s",
            conversation_id, exc,
        )


async def recover_stale_compaction_flags(
    conn: Any,
    stale_after_seconds: int = DEFAULT_STALE_COMPACTION_TIMEOUT_SECONDS,
) -> int:
    """Clear compaction_in_progress flags that have been set for too long.

    Returns the number of rows cleared (stale compactions recovered).
    Called on engine startup to clean up after crashes.
    """
    try:
        cursor = await conn.execute(
            """
            UPDATE hippo_conversations
               SET compaction_in_progress = 0,
                   compaction_started_at  = NULL
             WHERE compaction_in_progress = 1
               AND compaction_started_at IS NOT NULL
               AND (julianday('now') - julianday(compaction_started_at)) * 86400 > ?
            """,
            (stale_after_seconds,),
        )
        await conn.commit()
        cleared = cursor.rowcount if cursor.rowcount >= 0 else 0
        if cleared > 0:
            logger.warning(
                "compaction.hardening: recovered %d stale compaction flag(s) "
                "(stale_after=%ds)",
                cleared, stale_after_seconds,
            )
        return cleared
    except Exception as exc:
        logger.debug("recover_stale_compaction_flags: failed: %s", exc)
        return 0


# ── Safe compaction wrapper ───────────────────────────────────────────────────


async def safe_compact(
    conversation_id: int,
    compact_fn,
    guard: Optional[CompactionConcurrencyGuard] = None,
    conn: Optional[Any] = None,
    *args,
    **kwargs,
):
    """Run compact_fn with full hardening:

    1. Concurrent guard — skip if another compaction is already running.
    2. DB-level in-progress flag — survives process restarts.
    3. Failure counter — increments metrics on exception.
    4. Cleanup — always clears DB flag (partial failure recovery).

    Parameters
    ----------
    conversation_id
        Active conversation ID.
    compact_fn
        Async callable to invoke. Receives *args and **kwargs.
    guard
        Optional CompactionConcurrencyGuard instance.
    conn
        Optional aiosqlite connection for DB-flag updates.
    """
    # B-019 fix: If no guard provided, create a module-level default guard
    # to ensure concurrency protection is never bypassed.
    effective_guard = guard if guard is not None else _default_guard
    async with effective_guard.acquire(conversation_id) as acquired:
        if not acquired:
            return None
        return await _run_with_db_flag(
            conversation_id, compact_fn, conn, *args, **kwargs
        )


async def _run_with_db_flag(
    conversation_id: int,
    compact_fn,
    conn: Optional[Any],
    *args,
    **kwargs,
):
    """Run compact_fn with DB-level flag and failure tracking."""
    if conn is not None:
        await mark_compaction_start(conn, conversation_id)
    try:
        return await compact_fn(*args, **kwargs)
    except Exception as exc:
        _metrics.failed_compactions += 1
        logger.error(
            "compaction.hardening: compaction failed for conversation %d: %s",
            conversation_id, exc, exc_info=True,
        )
        raise
    finally:
        if conn is not None:
            await mark_compaction_end(conn, conversation_id)
