"""Session lifecycle operations: consolidation, memory flush, index sync, /new.

Extracted from AgentLoop to allow independent development of session
management and memory archival logic.
"""

from __future__ import annotations

import asyncio
from typing import TYPE_CHECKING

from loguru import logger

from asedclaw.agent.memory import MemoryStore
from asedclaw.bus.events import OutboundMessage

if TYPE_CHECKING:
    from pathlib import Path

    from asedclaw.providers.base import LLMProvider
    from asedclaw.session.manager import Session


async def handle_new_session(
    msg,
    session: Session,
    *,
    workspace: Path,
    provider: LLMProvider,
    model: str,
    memory_window: int,
    consolidation_locks: dict,
    consolidating: set,
    memory_index=None,
) -> OutboundMessage:
    """Handle the ``/new`` slash command: archive → clear → respond."""
    lock = consolidation_locks.setdefault(session.key, asyncio.Lock())
    consolidating.add(session.key)
    try:
        async with lock:
            snapshot = session.messages[session.last_consolidated:]
            if snapshot:
                from asedclaw.session.manager import Session as SessionCls
                temp = SessionCls(key=session.key)
                temp.messages = list(snapshot)
                if not await consolidate_memory(
                    temp, workspace=workspace, provider=provider,
                    model=model, memory_window=memory_window,
                    archive_all=True, memory_index=memory_index,
                ):
                    return OutboundMessage(
                        channel=msg.channel, chat_id=msg.chat_id,
                        content="Memory archival failed, session not cleared. Please try again.",
                    )
    except Exception:
        logger.exception("/new archival failed for {}", session.key)
        return OutboundMessage(
            channel=msg.channel, chat_id=msg.chat_id,
            content="Memory archival failed, session not cleared. Please try again.",
        )
    finally:
        consolidating.discard(session.key)

    session.clear()
    return OutboundMessage(channel=msg.channel, chat_id=msg.chat_id, content="New session started.")


def maybe_consolidate(
    session: Session,
    *,
    memory_window: int,
    workspace: Path,
    provider: LLMProvider,
    model: str,
    consolidating: set,
    consolidation_locks: dict,
    consolidation_tasks: set,
    flush_done_for_cycle: dict,
    memory_index=None,
) -> None:
    """Trigger background consolidation if unconsolidated messages exceed window."""
    unconsolidated = len(session.messages) - session.last_consolidated
    if unconsolidated < memory_window or session.key in consolidating:
        return
    consolidating.add(session.key)
    lock = consolidation_locks.setdefault(session.key, asyncio.Lock())

    async def _do():
        try:
            async with lock:
                await consolidate_memory(
                    session, workspace=workspace, provider=provider,
                    model=model, memory_window=memory_window,
                    memory_index=memory_index,
                )
                flush_done_for_cycle.pop(session.key, None)
        finally:
            consolidating.discard(session.key)
            t = asyncio.current_task()
            if t is not None:
                consolidation_tasks.discard(t)

    task = asyncio.create_task(_do())
    consolidation_tasks.add(task)


async def maybe_memory_flush(
    session: Session,
    *,
    provider: LLMProvider,
    model: str,
    workspace: Path,
    memory_flush_config,
    flush_done_for_cycle: dict,
) -> None:
    """Silently flush memory if token or transcript thresholds are exceeded."""
    from asedclaw.agent.memory_flush import (
        run_memory_flush,
        should_flush,
        should_force_flush_by_size,
    )
    from asedclaw.agent.token_utils import estimate_messages_tokens

    cfg = memory_flush_config
    flush_done = flush_done_for_cycle.get(session.key, False)

    need_flush = False
    if not flush_done and cfg.enabled:
        total_tokens = estimate_messages_tokens(session.messages)
        if should_flush(total_tokens, cfg, flush_done):
            logger.info(
                "Memory flush triggered (tokens={}, threshold={})",
                total_tokens,
                cfg.context_window - cfg.reserve_tokens_floor - cfg.soft_threshold_tokens,
            )
            need_flush = True
        elif should_force_flush_by_size(session.transcript_bytes, cfg, flush_done):
            logger.info("Memory flush triggered (transcript_bytes={})", session.transcript_bytes)
            need_flush = True

    if need_flush:
        success = await run_memory_flush(session, provider, model, workspace, cfg)
        flush_done_for_cycle[session.key] = success


async def maybe_sync_index(
    session: Session,
    *,
    memory_index,
    index_synced_sessions: set,
) -> None:
    """Sync memory index on first access per session."""
    if memory_index is None:
        return
    if session.key in index_synced_sessions:
        return
    index_synced_sessions.add(session.key)
    try:
        await memory_index.sync()
    except Exception:
        logger.debug("Memory index sync failed", exc_info=True)


async def consolidate_memory(
    session: Session,
    *,
    workspace: Path,
    provider: LLMProvider,
    model: str,
    memory_window: int = 100,
    archive_all: bool = False,
    memory_index=None,
) -> bool:
    """Run memory consolidation and optionally trigger index sync."""
    result = await MemoryStore(workspace).consolidate(
        session, provider, model,
        archive_all=archive_all, memory_window=memory_window,
    )
    if result and memory_index is not None:
        asyncio.create_task(memory_index.sync())
    return result
