"""
context_compactor.py

Auto-compacts conversation history when approaching context window limits.
Summarizes older messages via the same LLM, preserving key context.
"""

import json
import logging
import re
from typing import Any, Dict, List, Optional

from src.model_context import estimate_text_tokens, get_context_length, estimate_tokens
from src.llm_core import llm_call_async
from src.endpoint_resolver import resolve_endpoint
from src.settings import get_setting
from core.models import ChatMessage

logger = logging.getLogger(__name__)


def _content_as_text(content: Any) -> str:
    """Flatten a message's content to plain text.

    Handles the three shapes that flow through history: a plain string, a
    multimodal list of content blocks (vision/image attachments), and None
    (assistant turns that carried only native tool_calls persist content as
    None). Returns "" for anything without text so callers can safely slice
    the result.
    """
    if isinstance(content, str):
        return content
    if isinstance(content, list):
        return " ".join(
            b.get("text", "") for b in content
            if isinstance(b, dict) and b.get("text")
        )
    return ""


_MULTIMODAL_IMAGE_TYPES = {"image_url", "input_image", "image"}


def prune_multimodal_images(
    messages: List[Dict],
    *,
    max_images: int = 8,
) -> List[Dict]:
    """Keep uniformly sampled visual blocks across a multimodal history."""

    image_locations = []
    for message_index, message in enumerate(messages):
        content = message.get("content") if isinstance(message, dict) else None
        if not isinstance(content, list):
            continue
        for content_index, item in enumerate(content):
            if isinstance(item, dict) and item.get("type") in _MULTIMODAL_IMAGE_TYPES:
                image_locations.append((message_index, content_index))

    limit = max(0, int(max_images))
    if len(image_locations) <= limit:
        return list(messages)
    if limit == 0:
        selected = set()
    elif limit == 1:
        selected = {image_locations[-1]}
    else:
        last = len(image_locations) - 1
        selected = {
            image_locations[round(index * last / (limit - 1))]
            for index in range(limit)
        }

    pruned = []
    for message_index, message in enumerate(messages):
        if not isinstance(message, dict):
            pruned.append(message)
            continue
        cloned = dict(message)
        content = message.get("content")
        if isinstance(content, list):
            cloned["content"] = [
                item for content_index, item in enumerate(content)
                if not (
                    isinstance(item, dict)
                    and item.get("type") in _MULTIMODAL_IMAGE_TYPES
                    and (message_index, content_index) not in selected
                )
            ]
        pruned.append(cloned)
    return pruned


COMPACT_THRESHOLD = 0.85  # Trigger compaction at 85% of context window
SUMMARY_MAX_TOKENS = 1024
SMALL_CONTEXT_LIMIT = 8192  # Models with context <= this get aggressive trimming


def auto_compact_threshold_percent() -> int:
    """Configured auto-compaction threshold, clamped to a sane UI range."""
    try:
        value = int(get_setting("auto_compact_threshold_percent", int(COMPACT_THRESHOLD * 100)) or 85)
    except (TypeError, ValueError):
        value = int(COMPACT_THRESHOLD * 100)
    return max(50, min(95, value))

# Cursor-style self-summarization prompt — produces structured, dense summaries
SELF_SUMMARY_SYSTEM_PROMPT = """You are summarizing a conversation to preserve context after compaction. Produce a structured summary that lets the conversation continue seamlessly.

Use this format:

## Conversation Summary
**Turns summarized:** {count}  |  **Compactions so far:** {n}

### User Goal
One sentence describing what the user is trying to accomplish.

### What Was Done
- Bullet points of completed actions, decisions made, and key outputs
- Include specific file paths, function names, variable names, URLs, and config values
- Note any errors encountered and how they were resolved

### Current State
What is the system/code/task state right now? What was the last thing discussed?

### Pending / Next Steps
- What remains to be done
- Any open questions or blockers

### Key Context
- Important constraints, preferences, or decisions that must not be forgotten
- Specific values: model names, ports, paths, credentials references, versions

Keep the summary under 1000 tokens. Be dense — every token should carry information. Do not include pleasantries or meta-commentary."""


def normalize_compaction_summary(summary: str) -> str:
    """Remove redundant leading title text before adding our wrapper."""
    text = (summary or "").strip()
    text = re.sub(r"^(?:#{1,3}\s*)?Conversation Summary\s*", "", text, flags=re.IGNORECASE)
    text = re.sub(r"^\*\*Conversation Summary\*\*\s*", "", text, flags=re.IGNORECASE)
    return text.lstrip()


def _sanitize_tool_messages(msgs: List[Dict]) -> List[Dict]:
    """Drop orphaned `tool` messages and dangling assistant `tool_calls`.

    OpenAI's API requires every `role:"tool"` message to immediately
    follow an assistant message that carries `tool_calls` (or another
    tool message in the same batch). Front-trimming the history can cut
    the assistant `tool_calls` parent while keeping its tool responses,
    which triggers: "messages with role 'tool' must be a response to a
    preceding message with 'tool_calls'". This pass repairs that:
      - drops `tool` messages with no valid preceding tool_calls
      - drops assistant `tool_calls` messages whose tool responses were
        all trimmed away (some providers reject unanswered tool_calls)
    """
    # Pass 1: drop orphan tool messages.
    cleaned: List[Dict] = []
    in_batch = False  # are we right after an assistant tool_calls (or mid-batch)?
    for m in msgs:
        role = m.get("role")
        if role == "tool":
            if in_batch:
                cleaned.append(m)
            # else: orphan — drop
            continue
        if role == "assistant" and m.get("tool_calls"):
            in_batch = True
        else:
            in_batch = False
        cleaned.append(m)

    # Pass 2: drop assistant tool_calls messages that have NO following
    # tool response (dangling) — walk backwards so we know what follows.
    out: List[Dict] = []
    for i, m in enumerate(cleaned):
        if m.get("role") == "assistant" and m.get("tool_calls"):
            nxt = cleaned[i + 1] if i + 1 < len(cleaned) else None
            if not (nxt and nxt.get("role") == "tool"):
                # Dangling tool_calls — keep the message but strip the
                # tool_calls so it's a plain assistant turn (preserves any
                # text content the model produced alongside the calls).
                m = {k: v for k, v in m.items() if k != "tool_calls"}
                if not (m.get("content") or "").strip():
                    continue  # nothing left worth keeping
        out.append(m)
    return out


def _message_text_token_estimate(text: str) -> int:
    if not isinstance(text, str):
        return 4
    return estimate_text_tokens(text) + 4


def _truncate_text_to_token_budget(text: str, token_budget: int) -> str:
    """Trim a too-large current user message instead of dropping it entirely."""
    if token_budget <= 32:
        return "[Current user message omitted: it exceeded the model context window.]"

    if not isinstance(text, str):
        # This helper is typed/used as text downstream, so return an empty
        # string rather than the raw non-string (which would move the crash
        # into the caller that concatenates/measures the result).
        return ""
    if estimate_text_tokens(text) <= token_budget - 16:
        return text

    notice = (
        "\n\n[Notice: the pasted message was too large for this model's context "
        "window, so Odysseus kept the beginning and end.]"
    )
    # Binary replacement characters and dense scripts can approach one token
    # per character, while ASCII prose is closer to the historical 0.3 ratio.
    # Find the largest head/tail sample that fits the shared estimator instead
    # of guessing a character count from one language family.
    target = max(64, token_budget - estimate_text_tokens(notice) - 16)
    low, high = 1, len(text)
    best = 1
    while low <= high:
        keep_chars = (low + high) // 2
        head_len = max(1, int(keep_chars * 0.7))
        tail_len = max(0, keep_chars - head_len)
        sample = text[:head_len]
        if tail_len:
            sample += text[-tail_len:]
        if estimate_text_tokens(sample) <= target:
            best = keep_chars
            low = keep_chars + 1
        else:
            high = keep_chars - 1
    head_len = max(1, int(best * 0.7))
    tail_len = max(0, best - head_len)
    tail = text[-tail_len:].lstrip() if tail_len else ""
    return text[:head_len].rstrip() + notice + ("\n\n" + tail if tail else "")


def _truncate_tool_call_args(msg: Dict[str, Any], token_budget: int) -> Dict[str, Any]:
    """Shrink oversized assistant ``tool_calls`` arguments to fit ``token_budget``.

    A tool-only turn persists ``content=None`` with its whole payload in
    ``tool_calls[].function.arguments`` (e.g. a large create_document body), which
    the text-content truncation can't reach — so the message could stay over
    budget and the upstream call would 400. Replace each argument string that
    overflows its share of the budget with a small valid-JSON placeholder,
    preserving ``id``/``type``/``function.name`` so tool/result pairing and
    provider validation are unaffected. Returns msg unchanged when there is
    nothing oversized.
    """
    tool_calls = msg.get("tool_calls")
    if not isinstance(tool_calls, list) or not tool_calls:
        return msg
    # Budget left after whatever content survived (estimate_tokens counts tool
    # arguments too, so measure content alone here).
    content_tokens = estimate_tokens([{"role": msg.get("role", "assistant"), "content": msg.get("content")}])
    per_call = max(16, (max(0, token_budget - content_tokens)) // len(tool_calls))
    new_calls = []
    changed = False
    for tc in tool_calls:
        fn = tc.get("function") if isinstance(tc, dict) else None
        args = fn.get("arguments") if isinstance(fn, dict) else None
        if isinstance(args, str) and estimate_text_tokens(args) > per_call:
            new_fn = dict(fn)
            new_fn["arguments"] = json.dumps({"_truncated_for_context": len(args)})
            new_tc = dict(tc)
            new_tc["function"] = new_fn
            new_calls.append(new_tc)
            changed = True
        else:
            new_calls.append(tc)
    if not changed:
        return msg
    out = dict(msg)
    out["tool_calls"] = new_calls
    return out


def _truncate_message_to_token_budget(msg: Dict[str, Any], token_budget: int) -> Dict[str, Any]:
    """Return a copy of msg whose text content (and tool-call args) fit token_budget."""
    out = dict(msg)
    content = out.get("content", "")
    if isinstance(content, str):
        out["content"] = _truncate_text_to_token_budget(content, token_budget)
    elif isinstance(content, list):
        remaining = token_budget
        new_content = []
        for item in content:
            if not isinstance(item, dict) or item.get("type") != "text":
                new_content.append(item)
                continue
            text = item.get("text", "")
            truncated = _truncate_text_to_token_budget(text, remaining)
            cloned = dict(item)
            cloned["text"] = truncated
            new_content.append(cloned)
            remaining -= _message_text_token_estimate(truncated)
        out["content"] = new_content
    # A tool-only turn (content=None) carries its payload in tool_calls args,
    # which the branches above can't shrink — handle it so the message can fit.
    return _truncate_tool_call_args(out, token_budget)


def trim_for_context(messages: List[Dict], context_length: int, reserve_tokens: int = 512) -> List[Dict]:
    """Trim system messages to fit within context_length.

    For small-context models, progressively strips:
    1. RAG/memory system messages (keep preset system prompt)
    2. Older conversation turns
    Reserves space for the response.
    """
    budget = context_length - reserve_tokens
    used = estimate_tokens(messages)
    if used <= budget:
        return messages

    logger.info(f"Trimming messages: {used} tokens > {budget} budget (ctx={context_length})")

    # Separate system messages from conversation.
    # Messages marked _protected (e.g. active document) are never trimmed.
    system_msgs = []
    protected_msgs = []
    convo_msgs = []
    for msg in messages:
        if msg.get("_protected"):
            protected_msgs.append(msg)
        elif msg.get("role") == "system":
            system_msgs.append(msg)
        else:
            convo_msgs.append(msg)

    # Protected messages count toward budget but are never dropped
    protected_tokens = estimate_tokens(protected_msgs)
    budget -= protected_tokens

    # Priority: keep first system msg (preset prompt), drop others (memory, RAG, memo).
    # Exception: a research-spinoff primer (the seeded report that grounds a
    # "Discuss" chat) must never be dropped — it is the conversation's whole
    # knowledge base. Treat any system message carrying research_spinoff_from
    # metadata as essential alongside the leading system prompt.
    def _is_research_primer(m):
        return bool((m.get("metadata") or {}).get("research_spinoff_from"))

    def _is_compaction_summary(m):
        content = m.get("content")
        return isinstance(content, str) and content.lstrip().startswith(
            ("[Conversation summary", "Conversation Summary")
        )

    _primers = [m for m in system_msgs if _is_research_primer(m)]
    _summaries = [m for m in system_msgs if _is_compaction_summary(m)]
    _non_essential = [
        m for m in system_msgs
        if not _is_research_primer(m) and not _is_compaction_summary(m)
    ]
    # The base prompt and the conversation summary are the minimum state
    # needed to continue a task. Summaries used to be classified as ordinary
    # extra system context, so a large route prompt could trim them away right
    # after compaction and send the model the same lost-context request again.
    essential_system = (
        (_non_essential[:1] if _non_essential else [])
        + _primers
        + _summaries
    )
    extra_system = _non_essential[1:]

    # Try dropping extra system messages one by one (from the end)
    trimmed = essential_system + convo_msgs
    if estimate_tokens(trimmed) <= budget:
        # Dropping extras was enough — try adding back some
        result = list(essential_system)
        for msg in extra_system:
            candidate = result + [msg] + convo_msgs
            if estimate_tokens(candidate) <= budget:
                result.append(msg)
            else:
                break
        return _sanitize_tool_messages(result + protected_msgs + convo_msgs)

    # Still too big — truncate the first system message (but keep more than 500 chars)
    if essential_system:
        sys_text = essential_system[0].get("content", "")
        if len(sys_text) > 2000:
            truncated_system = dict(essential_system[0])
            truncated_system["content"] = sys_text[:2000] + "\n[System prompt truncated for context limits]"
            essential_system[0] = truncated_system
            trimmed = essential_system + convo_msgs
            if estimate_tokens(trimmed) <= budget:
                return _sanitize_tool_messages(essential_system + protected_msgs + convo_msgs)

    # Still too big — drop older conversation turns BUT always keep the latest
    # user turn. After a tool round the final message is usually an assistant
    # tool-call or a tool result, not the user's request; treating the last
    # message as "current" drops the real question and lets the model answer
    # stale context on the follow-up round.
    PROTECT_RECENT = 10
    def _is_direct_user_message(message: Dict) -> bool:
        # Qwen transports runtime corrections as user-role messages. Their
        # server-owned provenance must not make them replace the real request
        # as the protected tail's anchor. Never infer this from prompt wording.
        if message.get("role") != "user" or message.get("_harness_control"):
            return False
        metadata = message.get("metadata") or {}
        # Textual tool transports intentionally wrap external results as user
        # messages. They are context for the request, not a new request.
        return not (
            metadata.get("trusted") is False
            and bool(metadata.get("source"))
        )

    latest_user_idx = -1
    for idx in range(len(convo_msgs) - 1, -1, -1):
        if _is_direct_user_message(convo_msgs[idx]):
            latest_user_idx = idx
            break
    if latest_user_idx < 0:
        # Preserve the historical fallback for callers that only supply
        # synthetic context and no direct user turn.
        for idx in range(len(convo_msgs) - 1, -1, -1):
            if convo_msgs[idx].get("role") == "user":
                latest_user_idx = idx
                break
    if latest_user_idx >= 0:
        current_tail = convo_msgs[latest_user_idx:]
        prior_convo = convo_msgs[:latest_user_idx]
    else:
        current_tail = convo_msgs[-1:] if convo_msgs else []
        prior_convo = convo_msgs[:-1] if convo_msgs else []

    recent_prior_count = max(0, PROTECT_RECENT - len(current_tail))
    if len(prior_convo) > recent_prior_count:
        old_msgs = prior_convo[:-recent_prior_count] if recent_prior_count else prior_convo[:]
        recent_msgs = (prior_convo[-recent_prior_count:] if recent_prior_count else []) + current_tail
        while old_msgs and estimate_tokens(essential_system + old_msgs + recent_msgs) > budget:
            old_msgs.pop(0)
        convo_msgs = old_msgs + recent_msgs
    else:
        while prior_convo and estimate_tokens(essential_system + prior_convo + current_tail) > budget:
            prior_convo.pop(0)
        convo_msgs = prior_convo + current_tail

    # If the current request + tool tail is still too large, shrink that tail
    # instead of dropping the latest user turn. Native tool responses can be
    # huge (web search/fetch), and preserving a truncated source block is better
    # than sending the model a prompt with no active user request.
    if current_tail and estimate_tokens(essential_system + protected_msgs + convo_msgs) > budget:
        tail_start = len(convo_msgs) - len(current_tail)
        prefix = essential_system + protected_msgs + convo_msgs[:tail_start]
        available_for_tail = max(64 * len(current_tail), budget - estimate_tokens(prefix))
        per_tail_msg = max(64, available_for_tail // max(1, len(current_tail)))
        convo_msgs[tail_start:] = [
            _truncate_message_to_token_budget(msg, per_tail_msg)
            for msg in current_tail
        ]

    # Last ditch: if the tool tail still cannot fit, keep only the latest user
    # message (truncated if needed). Losing source output is bad; losing the user
    # request is worse and caused visibly crossed answers.
    if current_tail and estimate_tokens(essential_system + protected_msgs + convo_msgs) > budget:
        latest_user = next(
            (m for m in current_tail if _is_direct_user_message(m)),
            next((m for m in current_tail if m.get("role") == "user"), current_tail[0]),
        )
        prefix = essential_system + protected_msgs
        available_for_current = max(64, budget - estimate_tokens(prefix))
        convo_msgs = [_truncate_message_to_token_budget(latest_user, available_for_current)]

    result = _sanitize_tool_messages(essential_system + protected_msgs + convo_msgs)
    logger.info(f"Trimmed to {estimate_tokens(result)} tokens ({len(result)} messages)")
    return result


async def maybe_compact(
    session,
    endpoint_url: str,
    model: str,
    messages: List[Dict],
    headers: Optional[Dict] = None,
    owner: Optional[str] = None,
    *,
    persist: bool = True,
    compaction_state: Optional[Dict[str, Any]] = None,
    deterministic: bool = False,
) -> tuple:
    """Check context usage and compact if above threshold.

    Returns (messages, context_length, was_compacted).
    """
    context_length = get_context_length(endpoint_url, model)
    used = estimate_tokens(messages)
    pct = (used / context_length) * 100 if context_length else 0
    threshold = auto_compact_threshold_percent()

    if pct < threshold:
        return messages, context_length, False

    logger.info(
        f"Context at {pct:.1f}% ({used}/{context_length} tokens, threshold={threshold}%) — compacting"
    )

    if deterministic:
        # Unattended workers must not compete with policy generation for the
        # same saturated endpoint just to summarize their own transcript. Trim
        # to 75% so the next tool round has useful headroom. This path is not
        # persisted into interactive history and preserves the active request
        # through trim_for_context's protected-tail rules.
        reserve_tokens = max(512, int(context_length * 0.25))
        compacted = trim_for_context(
            messages,
            context_length,
            reserve_tokens=reserve_tokens,
        )
        changed = compacted != messages
        if changed:
            logger.info(
                "Deterministically compacted: %s -> %s tokens (%s -> %s messages)",
                used,
                estimate_tokens(compacted),
                len(messages),
                len(compacted),
            )
        return compacted, context_length, changed

    # Split into system preface and conversation
    system_msgs = []
    convo_msgs = []
    for msg in messages:
        if msg.get("role") == "system":
            system_msgs.append(msg)
        else:
            convo_msgs.append(msg)

    if len(convo_msgs) < 4:
        return messages, context_length, False

    # Split conversation: summarize older half, keep recent half
    split_point = len(convo_msgs) // 2
    older = convo_msgs[:split_point]
    recent = convo_msgs[split_point:]

    # Build the text to summarize
    convo_text = "\n".join(
        f"{msg.get('role', 'user').upper()}: {_content_as_text(msg.get('content'))[:2000]}"
        for msg in older
    )

    # Count prior compactions from existing summary messages
    compaction_count = sum(
        1 for m in system_msgs
        if "[Conversation summary" in m.get("content", "")
    )

    # Use utility model if configured, otherwise fall back to session model
    util_url, util_model, util_headers = resolve_endpoint("utility", owner=owner)
    compact_url = util_url or endpoint_url
    compact_model = util_model or model
    compact_headers = util_headers if util_url else headers

    prompt = SELF_SUMMARY_SYSTEM_PROMPT.replace(
        "{count}", str(len(older))
    ).replace(
        "{n}", str(compaction_count + 1)
    )
    summary_messages = [
        {"role": "system", "content": prompt},
        {"role": "user", "content": convo_text},
    ]

    try:
        summary = await llm_call_async(
            compact_url,
            compact_model,
            summary_messages,
            temperature=0.2,
            max_tokens=SUMMARY_MAX_TOKENS,
            headers=compact_headers,
            timeout=30,
        )
    except Exception as e:
        logger.error(f"Compaction summary failed: {e}")
        # Degrade gracefully: keep the conversation intact rather than
        # silently dropping the older half. was_compacted=False signals the
        # caller nothing was summarized; trim_for_context handles length.
        return messages, context_length, False
    summary = normalize_compaction_summary(summary)

    summary_msg = {
        "role": "system",
        "content": f"[Conversation summary — earlier messages were compacted]\n{summary}",
    }

    compacted = system_msgs + [summary_msg] + recent

    # Update session history to match. Pass len(system_msgs) so the
    # recent_history slice in _update_session_history uses the correct
    # offset — session.history INCLUDES the system messages, but
    # split_point is indexed against convo_msgs which does NOT. Without
    # this, the slice drops the leading system message(s).
    if compaction_state is not None:
        compaction_state.update({
            "split_point": split_point,
            "summary": summary,
            "system_msg_count": len(system_msgs),
            "applied": False,
        })
    if persist:
        _update_session_history(session, split_point, summary, system_msg_count=len(system_msgs))
        if compaction_state is not None:
            compaction_state["applied"] = True

    new_used = estimate_tokens(compacted)
    logger.info(
        f"Compacted: {used} -> {new_used} tokens "
        f"({len(older)} messages summarized, {len(recent)} kept)"
    )

    return compacted, context_length, True


def apply_compaction_state(session, compaction_state: Optional[Dict[str, Any]]) -> bool:
    """Persist a route-specific compaction after that route commits output.

    Candidate prompts may be compacted speculatively while an explicit
    foreground fallback chain is being tried.  Persisting at construction time
    would let an unavailable route rewrite history before another route answers,
    so callers hold this small plan and apply only the winning route's plan.
    """

    state = compaction_state if isinstance(compaction_state, dict) else None
    if not state or state.get("applied"):
        return False
    summary = state.get("summary")
    split_point = state.get("split_point")
    system_msg_count = state.get("system_msg_count", 0)
    if not isinstance(summary, str) or not isinstance(split_point, int):
        return False
    _update_session_history(
        session,
        split_point,
        summary,
        system_msg_count=system_msg_count if isinstance(system_msg_count, int) else 0,
    )
    state["applied"] = True
    return True


def apply_compaction_state_for_session(
    session_id: Optional[str],
    compaction_state: Optional[Dict[str, Any]],
) -> bool:
    """Resolve an in-memory session and apply a deferred compaction plan."""

    if not session_id:
        return False
    try:
        from core.models import get_session_manager_instance

        manager = get_session_manager_instance()
        session = manager.get_session(session_id) if manager else None
    except Exception:
        session = None
    return apply_compaction_state(session, compaction_state) if session else False


def _update_session_history(session, split_point: int, summary: str,
                            system_msg_count: int = 0):
    """Update the in-memory session history after compaction.

    `split_point` is the index in `convo_msgs` (system-stripped). The
    in-memory `session.history` includes leading system messages, so the
    actual recent-history slice starts at `system_msg_count + split_point`.
    Prepending `session.history[:system_msg_count]` to the new history
    preserves persona, preset, and RAG system messages that would
    otherwise be dropped.
    """
    if not session or not hasattr(session, "history"):
        return

    effective_split = system_msg_count + split_point
    if effective_split >= len(session.history):
        return

    # Keep the recent messages, prepend summary AND the leading system
    # messages so the system prompt survives compaction.
    system_prefix = list(session.history[:system_msg_count])
    recent_history = session.history[effective_split:]
    summary = normalize_compaction_summary(summary)
    summary_msg = ChatMessage(
        role="system",
        content=f"[Conversation summary]\n{summary}",
        metadata={"compacted": True, "summarized_count": split_point},
    )
    new_history = system_prefix + [summary_msg] + recent_history
    try:
        from core.models import get_session_manager_instance
        manager = get_session_manager_instance()
    except Exception:
        manager = None
    if manager and getattr(session, "id", None):
        if manager.replace_messages(session.id, new_history):
            return
    session.history = new_history