"""
cortex/hippo/compaction/leaf.py — Leaf pass functions (raw messages → depth-0 summaries).
"""

from __future__ import annotations
import json

import logging
from typing import Optional

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 cortex.hippo.summarize import SummarizeOptions

from .helpers import _extract_file_ids, estimate_tokens, format_timestamp
from .types import LeafChunkSelection, PassResult, SummarizeFn
from .utils import _get_item_tokens, _get_summary_by_context_item, _summarize_with_escalation

logger = logging.getLogger(__name__)


def _resolve_fresh_tail_ordinal(
    context_items: list[ContextItemRecord],
    config: HippoConfig,
) -> float:
    """Compute ordinal boundary for protected fresh messages.

    Returns the ordinal at which the fresh tail starts.
    Items with ordinal >= this value are protected.
    Returns 0 if all messages should be protected (fewer than fresh_tail_count).
    Returns float('inf') if no messages exist (nothing to protect).
    """
    if not context_items or config.fresh_tail_count <= 0:
        return float("inf")

    # Count messages from the end to find the fresh tail boundary
    message_count = 0
    for item in reversed(context_items):
        if item.item_type == "message":
            message_count += 1
            if message_count >= config.fresh_tail_count:
                return float(item.ordinal)

    # Fewer messages than fresh_tail_count — protect all of them
    if message_count > 0:
        return 0.0  # All ordinals >= 0, so everything is protected
    return float("inf")  # No messages at all


async def _select_oldest_leaf_chunk(
    conversation_id: int,
    conversation_store: ConversationStore,
    summary_store: SummaryStore,
    config: HippoConfig,
) -> LeafChunkSelection:
    """Select oldest contiguous raw-message chunk outside fresh tail."""
    context_items = await summary_store.get_context_items(conversation_id)
    fresh_tail_ordinal = _resolve_fresh_tail_ordinal(context_items, config)

    chunk: list[ContextItemRecord] = []
    chunk_tokens = 0
    started = False
    raw_tokens_outside_tail = 0

    for item in context_items:
        if item.ordinal >= fresh_tail_ordinal:
            break

        if item.item_type == "message":
            raw_tokens_outside_tail += await _get_item_tokens(
                item, conversation_store, summary_store
            )

            if not started:
                started = True

            if started:
                msg_tokens = await _get_item_tokens(
                    item, conversation_store, summary_store
                )
                if chunk and chunk_tokens + msg_tokens > config.leaf_chunk_tokens:
                    break
                chunk.append(item)
                chunk_tokens += msg_tokens
                if chunk_tokens >= config.leaf_chunk_tokens:
                    break
        elif item.item_type == "summary":
            if started:
                # Hit a summary — stop collecting contiguous messages
                break
            # Skip summaries before finding messages

    return LeafChunkSelection(
        items=chunk,
        raw_tokens_outside_tail=raw_tokens_outside_tail,
        threshold=config.leaf_chunk_tokens,
    )


async def _check_leaf_trigger(
    conversation_id: int,
    conversation_store: ConversationStore,
    summary_store: SummaryStore,
    config: HippoConfig,
) -> bool:
    """Check if raw message tokens outside fresh tail exceed leaf_chunk_tokens."""
    chunk = await _select_oldest_leaf_chunk(
        conversation_id, conversation_store, summary_store, config
    )
    return chunk.raw_tokens_outside_tail > config.leaf_chunk_tokens


async def _resolve_prior_leaf_summary_context(
    conversation_id: int,
    message_items: list[ContextItemRecord],
    summary_store: SummaryStore,
) -> Optional[str]:
    """Collect up to 2 most recent summary items before the chunk."""
    if not message_items:
        return None
    chunk_start_ordinal = min(item.ordinal for item in message_items)
    context_items = await summary_store.get_context_items(conversation_id)

    prior_summaries = []
    for item in context_items:
        if item.ordinal >= chunk_start_ordinal:
            break
        if item.item_type == "summary" and item.summary_id:
            summary = await _get_summary_by_context_item(item, summary_store)
            if summary:
                prior_summaries.append(summary.content)

    if not prior_summaries:
        return None

    # Take last 2
    recent = prior_summaries[-2:]
    return "\n\n".join(recent)


async def _leaf_pass(
    conversation_id: int,
    message_items: list[ContextItemRecord],
    summarize: SummarizeFn,
    previous_summary_content: Optional[str],
    conversation_store: ConversationStore,
    summary_store: SummaryStore,
    config: HippoConfig,
) -> PassResult:
    """Convert oldest raw message chunk to leaf summary (depth=0, kind='leaf')."""
    # Fetch full message records
    texts = []
    message_ids = []
    total_source_tokens = 0
    earliest_at = None
    latest_at = None

    for item in message_items:
        if item.message_id is None:
            continue
        msg = await conversation_store.get_message_by_id(item.message_id)
        if not msg:
            continue
        message_ids.append(msg.message_id)
        ts = format_timestamp(msg.created_at, config.timezone)
        texts.append(f"[{ts}]\n{msg.content}")
        total_source_tokens += msg.token_count

        if earliest_at is None or msg.created_at < earliest_at:
            earliest_at = msg.created_at
        if latest_at is None or msg.created_at > latest_at:
            latest_at = msg.created_at

    concatenated = "\n\n".join(texts)

    # Extract file references
    file_ids = _extract_file_ids(concatenated)

    # Summarize with escalation
    summary_result = await _summarize_with_escalation(
        concatenated, summarize,
        SummarizeOptions(
            previous_summary=previous_summary_content,
            is_condensed=False,
            depth=0,
        ),
    )

    summary_content = summary_result.content
    level = summary_result.mode_used  # "normal" or "aggressive"
    quality_score = summary_result.quality_score

    token_count = estimate_tokens(summary_content)

    # Get owner_id from conversation
    # Use first message's owner_id
    first_msg = await conversation_store.get_message_by_id(message_ids[0]) if message_ids else None
    owner_id = first_msg.owner_id if first_msg else "default"

    # Insert summary — use the returned record's ID directly (no post-query race condition)
    created_record = await summary_store.create_summary(
        conversation_id=conversation_id,
        owner_id=owner_id,
        kind="leaf",
        depth=0,
        content=summary_content,
        token_count=token_count,
        quality_score=quality_score,
        earliest_at=earliest_at,
        latest_at=latest_at,
        descendant_count=len(message_ids),
        file_ids=json.dumps(file_ids) if file_ids else None,
        descendant_token_count=0,
        source_message_token_count=total_source_tokens,
    )
    actual_summary_id = created_record.summary_id

    # Link to source messages
    await summary_store.link_summary_messages(actual_summary_id, message_ids)

    # Replace context range
    start_ordinal = min(item.ordinal for item in message_items)
    end_ordinal = max(item.ordinal for item in message_items)
    await summary_store.replace_context_range_with_summary(
        conversation_id, start_ordinal, end_ordinal, actual_summary_id
    )

    return PassResult(
        summary_id=actual_summary_id,
        level=level,
        content=summary_content,
        tokens_before=total_source_tokens,
        tokens_after=token_count,
        quality_score=quality_score,
    )
