"""Context window management — trims conversation history to fit token limits."""
from __future__ import annotations

from typing import Callable, Awaitable

from logos_adk.config import ContextWindowConfig
from logos_adk.state import AgentState
from logos_adk.token_counter import count_messages, estimate_tool_tokens
from logos_adk.types import Message

# Default summarization prompt
_DEFAULT_SUMMARY_PROMPT = (
    "Summarize the following conversation segment concisely. "
    "Preserve: key decisions, facts, user preferences, action items, names, numbers. "
    "Omit: pleasantries, repetition, verbose tool outputs. "
    "Return only the summary, no preamble."
)

# Type alias for the async summarize callback
SummarizeFn = Callable[[list[Message], str | None], Awaitable[str]]


def _split_system_and_conversation(
    messages: list[Message],
) -> tuple[list[Message], list[Message]]:
    """Separate leading system messages from the conversation."""
    system_msgs: list[Message] = []
    conv_start = 0
    for i, msg in enumerate(messages):
        if msg.role == "system":
            system_msgs.append(msg)
            conv_start = i + 1
        else:
            break
    return system_msgs, messages[conv_start:]


def _truncate(
    messages: list[Message],
    *,
    max_tokens: int,
    keep_recent: int,
    model: str,
) -> list[Message]:
    """Drop oldest non-system messages until under token limit."""
    system, conv = _split_system_and_conversation(messages)
    if not conv:
        return messages

    keep_recent = min(keep_recent, len(conv))
    recent = conv[-keep_recent:] if keep_recent > 0 else []
    middle = conv[: len(conv) - keep_recent] if keep_recent > 0 else list(conv)

    # Start with system + recent, add middle messages from newest to oldest
    result = system + recent
    if count_messages(result, model=model) >= max_tokens:
        return result  # Even minimum exceeds limit

    # Add middle messages from most recent backwards
    added: list[Message] = []
    for msg in reversed(middle):
        candidate = system + [msg] + added + recent
        if count_messages(candidate, model=model) <= max_tokens:
            added.insert(0, msg)
        else:
            break

    return system + added + recent


def trim_context(
    messages: list[Message],
    *,
    config: ContextWindowConfig,
    state: AgentState,
    model: str = "",
    tool_schemas: list[dict] | None = None,
    summarize_fn: SummarizeFn | None = None,
) -> list[Message]:
    """Trim conversation messages to fit within the configured token budget.

    This is a synchronous entry point for truncate strategy.
    For strategies that need LLM calls (sliding_window, chunked_summary),
    use ``async_trim_context`` instead.
    """
    if config.max_context_tokens <= 0:
        return messages

    tool_tokens = estimate_tool_tokens(tool_schemas, model=model)
    available = config.max_context_tokens - config.reserved_output_tokens - tool_tokens
    if available <= 0:
        available = config.max_context_tokens // 2

    current = count_messages(messages, model=model)
    if current <= available:
        return messages

    # For sync path, always use truncate regardless of strategy
    return _truncate(
        messages,
        max_tokens=available,
        keep_recent=config.keep_recent_messages,
        model=model,
    )


async def async_trim_context(
    messages: list[Message],
    *,
    config: ContextWindowConfig,
    state: AgentState,
    model: str = "",
    tool_schemas: list[dict] | None = None,
    summarize_fn: SummarizeFn | None = None,
) -> list[Message]:
    """Async version of trim_context — supports all three strategies."""
    if config.max_context_tokens <= 0:
        return messages

    tool_tokens = estimate_tool_tokens(tool_schemas, model=model)
    available = config.max_context_tokens - config.reserved_output_tokens - tool_tokens
    if available <= 0:
        available = config.max_context_tokens // 2

    current = count_messages(messages, model=model)
    if current <= available:
        return messages

    strategy = config.strategy

    if strategy == "truncate" or summarize_fn is None:
        return _truncate(
            messages, max_tokens=available, keep_recent=config.keep_recent_messages, model=model,
        )

    if strategy == "sliding_window":
        return await _sliding_window(
            messages,
            max_tokens=available,
            keep_recent=config.keep_recent_messages,
            model=model,
            summarize_fn=summarize_fn,
            summary_prompt=config.summary_system_prompt,
        )

    if strategy == "chunked_summary":
        return await _chunked_summary(
            messages,
            config=config,
            state=state,
            max_tokens=available,
            model=model,
            summarize_fn=summarize_fn,
        )

    # Unknown strategy — safe fallback
    return _truncate(
        messages, max_tokens=available, keep_recent=config.keep_recent_messages, model=model,
    )


async def _sliding_window(
    messages: list[Message],
    *,
    max_tokens: int,
    keep_recent: int,
    model: str,
    summarize_fn: SummarizeFn,
    summary_prompt: str | None,
) -> list[Message]:
    """Summarize all old messages into one block, keep recent as-is."""
    system, conv = _split_system_and_conversation(messages)
    if not conv:
        return messages

    keep_recent = min(keep_recent, len(conv))
    recent = conv[-keep_recent:] if keep_recent > 0 else []
    old = conv[: len(conv) - keep_recent] if keep_recent > 0 else list(conv)

    if not old:
        return _truncate(messages, max_tokens=max_tokens, keep_recent=keep_recent, model=model)

    prompt = summary_prompt or _DEFAULT_SUMMARY_PROMPT
    summary_text = await summarize_fn(old, prompt)
    summary_msg = Message(role="system", content=f"[Conversation summary]: {summary_text}")

    result = system + [summary_msg] + recent
    if count_messages(result, model=model) > max_tokens:
        return _truncate(result, max_tokens=max_tokens, keep_recent=keep_recent, model=model)
    return result


async def _chunked_summary(
    messages: list[Message],
    *,
    config: ContextWindowConfig,
    state: AgentState,
    max_tokens: int,
    model: str,
    summarize_fn: SummarizeFn,
) -> list[Message]:
    """Incremental chunked summarization with persistent summary chain.

    History zones: [system] [summary_chain] [buffer] [recent]
    - recent: last keep_recent_messages, always sent as-is
    - buffer: messages not yet summarized
    - summary_chain: list of summaries from prior chunks
    - When buffer reaches chunk_size, it's summarized and appended to chain
    - When chain exceeds max_summary_chain, two oldest merge into one
    """
    system, conv = _split_system_and_conversation(messages)
    if not conv:
        return messages

    keep_recent = min(config.keep_recent_messages, len(conv))
    recent = conv[-keep_recent:] if keep_recent > 0 else []
    older = conv[: len(conv) - keep_recent] if keep_recent > 0 else list(conv)

    # Load existing summary chain from state
    chain: list[str] = list(state.metadata.get("_context_summaries", []))

    prompt = config.summary_system_prompt or _DEFAULT_SUMMARY_PROMPT

    # Determine how many older messages are already covered by summaries.
    summarized_count: int = state.metadata.get("_context_summarized_count", 0)
    buffer = older[summarized_count:]

    # Chunk the buffer and summarize full chunks
    while len(buffer) >= config.chunk_size:
        chunk = buffer[: config.chunk_size]
        buffer = buffer[config.chunk_size :]
        summarized_count += config.chunk_size

        summary_text = await summarize_fn(chunk, prompt)
        chain.append(summary_text)

        # Merge oldest if chain is too long
        while len(chain) > config.max_summary_chain and len(chain) >= 2:
            merged = await summarize_fn(
                [
                    Message(role="system", content=chain[0]),
                    Message(role="system", content=chain[1]),
                ],
                "Merge these two conversation summaries into one concise summary. "
                "Preserve all key facts, decisions, and action items.",
            )
            chain = [merged] + chain[2:]

    # Persist updated chain
    state.metadata["_context_summaries"] = chain
    state.metadata["_context_summarized_count"] = summarized_count

    # Build output: system + chain summaries + remaining buffer + recent
    result = list(system)
    if chain:
        chain_text = "\n\n---\n\n".join(chain)
        result.append(Message(role="system", content=f"[Conversation history summaries]:\n{chain_text}"))
    result.extend(buffer)
    result.extend(recent)

    # Final check: if still over budget, truncate the buffer
    if count_messages(result, model=model) > max_tokens:
        return _truncate(result, max_tokens=max_tokens, keep_recent=keep_recent, model=model)

    return result
