"""Conversation context builder and compactor."""

from __future__ import annotations

from typing import Any, Iterable

from .modelrelay import (
    chat_completion,
    get_model_context_window,
    get_model_max_output_tokens,
    resolve_model,
)
from .prompts import CONVERSATION_SUMMARY_SYSTEM_PROMPT, CONVERSATION_SUMMARY_USER_TEMPLATE

CHARS_PER_TOKEN = 4
SAFETY_MARGIN = 0.9
SUMMARY_FRACTION = 0.25
MIN_SUMMARY_CHARS = 400
MAX_SUMMARY_CHARS = 2000
DEFAULT_CONTEXT_WINDOW = 32000
DEFAULT_MAX_OUTPUT_TOKENS = 4096
SUMMARY_PROMPT_OVERHEAD = 400


def _format_messages(messages: Iterable[dict[str, Any]]) -> list[str]:
    formatted: list[str] = []
    for message in messages:
        role = str(message.get("role", "")).lower()
        text = str(message.get("text", "")).strip()
        if not text:
            continue
        label = "User" if role == "user" else "Assistant" if role == "assistant" else "System"
        formatted.append(f"{label}: {text}")
    return formatted


def _context_char_budget(model: str) -> int:
    context_window = get_model_context_window(model) or DEFAULT_CONTEXT_WINDOW
    max_output_tokens = get_model_max_output_tokens(model) or DEFAULT_MAX_OUTPUT_TOKENS
    available_tokens = max(1, context_window - max_output_tokens)
    return int(available_tokens * CHARS_PER_TOKEN * SAFETY_MARGIN)


def _summary_char_budget(max_context_chars: int) -> int:
    summary_budget = int(max_context_chars * SUMMARY_FRACTION)
    summary_budget = max(MIN_SUMMARY_CHARS, min(summary_budget, MAX_SUMMARY_CHARS))
    return min(summary_budget, max_context_chars // 2)


def _assemble_context(summary: str | None, recent: list[str]) -> str:
    segments: list[str] = []
    if summary:
        segments.append("Summary so far:")
        segments.append(summary.strip())
    if recent:
        segments.append("Recent conversation:")
        segments.extend(recent)
    return "\n".join(segments).strip()


def build_conversation_context(
    messages: list[dict[str, Any]],
    summary: str | None,
    model: str | None,
) -> tuple[str | None, str | None]:
    resolved_model = resolve_model(model)
    formatted = _format_messages(messages)
    if not formatted and not summary:
        return None, summary
    if not formatted and summary:
        return _assemble_context(summary, []), summary

    max_context_chars = _context_char_budget(resolved_model)
    context_without_summary = "\n".join(formatted)
    if len(context_without_summary) <= max_context_chars:
        return context_without_summary, summary

    summary_budget = _summary_char_budget(max_context_chars)
    available_for_recent = max_context_chars - summary_budget

    recent: list[str] = []
    running = 0
    for entry in reversed(formatted):
        entry_len = len(entry) + 1
        if running + entry_len > available_for_recent and recent:
            break
        running += entry_len
        recent.append(entry)
    recent.reverse()

    split_index = len(formatted) - len(recent)
    older = formatted[:split_index]

    updated_summary = summary
    if older:
        updated_summary = _summarize_messages(
            summary=summary,
            messages=older,
            model=resolved_model,
            max_summary_chars=summary_budget,
            max_context_chars=max_context_chars,
        )

    context = _assemble_context(updated_summary, recent)
    if len(context) > max_context_chars:
        updated_summary = _truncate_summary(updated_summary, max_context_chars - len("\n".join(recent)))
        context = _assemble_context(updated_summary, recent)
        while len(context) > max_context_chars and len(recent) > 1:
            recent.pop(0)
            context = _assemble_context(updated_summary, recent)

    return context, updated_summary


def _truncate_summary(summary: str | None, max_chars: int) -> str | None:
    if summary is None:
        return None
    trimmed = summary.strip()
    if max_chars <= 0:
        return ""
    if len(trimmed) <= max_chars:
        return trimmed
    trimmed = trimmed[:max_chars].rstrip()
    if trimmed.endswith("."):
        return trimmed
    return trimmed + "..."


def _summarize_messages(
    summary: str | None,
    messages: list[str],
    model: str,
    max_summary_chars: int,
    max_context_chars: int,
) -> str:
    existing_summary = summary.strip() if summary else ""
    remaining = list(messages)
    current_summary = existing_summary

    while remaining:
        chunk, remaining = _split_messages_to_fit(
            remaining,
            current_summary,
            max_context_chars,
        )
        current_summary = _summarize_chunk(
            current_summary,
            chunk,
            model,
            max_summary_chars,
        )

    return _truncate_summary(current_summary, max_summary_chars) or ""


def _split_messages_to_fit(
    messages: list[str],
    existing_summary: str,
    max_context_chars: int,
) -> tuple[list[str], list[str]]:
    available = max_context_chars - SUMMARY_PROMPT_OVERHEAD - len(existing_summary)
    if available <= 0:
        return messages[:1], messages[1:]

    chunk: list[str] = []
    running = 0
    for entry in messages:
        entry_len = len(entry) + 1
        if running + entry_len > available and chunk:
            break
        running += entry_len
        chunk.append(entry)

    remaining = messages[len(chunk):]
    return chunk, remaining


def _summarize_chunk(
    existing_summary: str,
    messages: list[str],
    model: str,
    max_summary_chars: int,
) -> str:
    formatted_messages = "\n".join(messages)
    user_prompt = CONVERSATION_SUMMARY_USER_TEMPLATE.format(
        existing_summary=existing_summary or "None",
        messages=formatted_messages,
        max_chars=max_summary_chars,
    )
    summary_text = chat_completion(
        [
            {"role": "system", "content": CONVERSATION_SUMMARY_SYSTEM_PROMPT},
            {"role": "user", "content": user_prompt},
        ],
        model=model,
        temperature=0.0,
        max_tokens=max(128, min(1024, max_summary_chars // CHARS_PER_TOKEN)),
    )
    return summary_text.strip()
