# context_summarizer.py — LLM-powered context summarization for long conversations.
# Compresses older messages into a single summary while preserving the recent tail,
# so the model retains awareness of earlier context without exceeding the context window.
from __future__ import annotations

import logging
import time
import re
from collections import Counter
from typing import Any
from utils.transcript_roles import is_tool_result_carrier_message, extract_inner_tool_response_text
from utils.llm_output_sanitize import sanitize_summary_text

_log = logging.getLogger(__name__)

_SUMMARY_SYSTEM_PROMPT = (
    "You are a conversation summarizer. Your job is to produce a concise, high-quality "
    "summary of the conversation history provided below. Focus on:\n"
    "- The user's original request and intent\n"
    "- Key findings, facts, data points, and conclusions reached so far\n"
    "- Important decisions made or actions taken (tool calls and their results)\n"
    "- A short **tool digest**: for each important tool invocation, note tool name, whether it "
    "succeeded (ok) or failed, and one line of outcome (path, count, or error prefix)\n"
    "- Any URLs, file paths, or specific values that were referenced\n"
    "- What the agent was working on most recently before this summary point\n\n"
    "CRITICAL — you MUST preserve these exactly (never omit or paraphrase away):\n"
    "- User refusals and constraints: any 'do not', 'never', 'stop doing', or similar directives\n"
    "- Deadlines, dates, and time-sensitive requirements\n"
    "- Exact file paths and URLs the user asked to track or work with\n"
    "- Tool failures or errors that changed the plan or safety assumptions\n\n"
    "Do NOT include filler, pleasantries, or meta-commentary about summarizing. "
    "Return ONLY the summary text — no preamble, no markdown headers."
)

_SUMMARY_USER_TEMPLATE = (
    "Summarize the following conversation history. Preserve all important context "
    "so the assistant can continue the task seamlessly.\n\n"
    "--- CONVERSATION TO SUMMARIZE ---\n\n{conversation}\n\n"
    "--- END ---\n\n"
    "Provide a concise summary (aim for 300-600 words)."
)

# Shorter retry prompt when the first completion loops on constraints / meta.
_SUMMARY_RETRY_SYSTEM_PROMPT = (
    "Summarize the chat below in plain prose only. Cover: what the user wanted, "
    "key facts, file paths and URLs, tool results and failures, and where things left off. "
    "Write the summary itself — do not restate rules about markdown, language, or formatting."
)

_MAX_CHARS_FOR_SUMMARY_INPUT = 40_000
_SUMMARY_MAX_TOKENS_PRIMARY = 768
_SUMMARY_MAX_TOKENS_RETRY = 512
# Regex: lines that echo summarization instructions instead of summarizing.
_INSTRUCTION_ECHO_RE = re.compile(
    r"(?i)(the summary should\b|should not include any markdown|"
    r"do not include any markdown|no preamble\b|return only the summary|"
    r"same language as the conversation|not include any markdown formatting)",
)


def _validate_summary_text(text: str | None) -> tuple[bool, str]:
    """Return (ok, reason). Rejects empty output, tight repetition loops, and instruction echo."""
    t = (text or "").strip()
    if not t:
        return False, "empty"
    lines = [ln.strip() for ln in t.splitlines() if ln.strip()]
    line_count = len(lines)
    top_line_freq = 0
    if line_count >= 1:
        top_line_freq = Counter(lines).most_common(1)[0][1]

    # Many identical lines (utility / small models often loop on one constraint sentence).
    if line_count >= 6 and top_line_freq / line_count >= 0.35:
        return False, "repetitive_lines"

    inst_hits = sum(1 for ln in lines if _INSTRUCTION_ECHO_RE.search(ln))
    if inst_hits >= 4:
        return False, "instruction_echo"
    if inst_hits >= 2 and line_count >= 6 and top_line_freq / line_count >= 0.25:
        return False, "instruction_echo_repetition"

    # Single blob: same sentence repeated with little variation.
    if line_count <= 3 and len(t) > 500:
        chunks = re.split(r"(?<=[.!?])\s+", t)
        chunks = [c.strip() for c in chunks if len(c.strip()) > 30]
        if len(chunks) >= 8:
            uniq = len(set(chunks))
            if uniq <= 2:
                return False, "repetitive_sentences"

    return True, ""


def _extractive_fallback_summary(conversation_text: str, *, max_body_chars: int = 8000) -> str:
    """Last resort: honest excerpt so context is not replaced by garbage."""
    note = (
        "[Note: automatic summarization failed, so the following is a trimmed excerpt of "
        "the older messages instead of an LLM summary.]\n\n"
    )
    body = conversation_text.strip()
    if len(body) > max_body_chars:
        body = body[:max_body_chars].rstrip() + "\n\n...[excerpt truncated]"
    return note + body


def _build_summary_message(summary_text: str) -> dict:
    return {
        "role": "user",
        "content": (
            "Summary of earlier conversation (older messages were compressed to fit context):\n\n"
            + summary_text
            + "\n\n---\n\n"
            "The conversation continues below. Use the summary above for context about what was discussed and accomplished earlier."
        ),
    }


def _format_messages_for_summary(messages: list[dict]) -> str:
    """Format a list of messages into a readable string for the summarization prompt."""
    parts = []
    for m in messages:
        role = m.get("role", "unknown").upper()
        content = m.get("content") or ""
        if isinstance(content, list):
            text_parts = [p.get("text", "") for p in content if isinstance(p, dict) and p.get("type") == "text"]
            content = " ".join(text_parts)
        if role == "ASSISTANT" and isinstance(content, str) and "<tool_call" in content.lower():
            try:
                from agent_loop import _parse_xml_tool_calls
                names = [tc.get("name", "") for tc in _parse_xml_tool_calls(content) if tc.get("name")]
                if names:
                    content = f"[Called tools: {', '.join(names)}]"
            except Exception:
                content = re.sub(r"<tool_call>[\s\S]*?</tool_call>", "[Called tool via XML]", content, flags=re.IGNORECASE)
        if role == "USER" and is_tool_result_carrier_message(m):
            inner = extract_inner_tool_response_text(content)
            content = inner or "[empty tool result]"
        if not content.strip():
            tool_calls = m.get("tool_calls")
            if tool_calls:
                tc_names = [tc.get("function", {}).get("name", "?") if isinstance(tc, dict) else "?" for tc in tool_calls]
                arg_hints: list[str] = []
                for tc in tool_calls:
                    if not isinstance(tc, dict):
                        arg_hints.append("")
                        continue
                    fn = tc.get("function") if isinstance(tc.get("function"), dict) else {}
                    raw_args = (fn.get("arguments") if isinstance(fn, dict) else None) or tc.get("arguments") or ""
                    sargs = str(raw_args).replace("\n", " ")[:120]
                    arg_hints.append(sargs)
                digest = ", ".join(f"{n}({h})" for n, h in zip(tc_names, arg_hints))
                content = f"[tool_digest: {digest}]"
            elif m.get("tool_call_id"):
                content = content or "[empty tool result]"
        rl = (m.get("role") or "").strip().lower()
        if rl == "tool":
            preview = (content or "").strip()
            pre_l = preview.lower()
            failed = any(
                pre_l.startswith(p)
                for p in ("tool error:", "error:", "unknown tool:", "operation blocked:", "blocked:", "refused:")
            )
            tc_id = m.get("tool_call_id") or ""
            tag = "FAIL" if failed else "ok"
            content = f"[tool_result id={tc_id} status={tag}] {preview[:1200]}"
        if len(content) > 1500:
            content = content[:1500] + "\n...[truncated for summary]"
        parts.append(f"{role}: {content}")
    return "\n\n".join(parts)


def _split_keeping_pairs(messages: list[dict], keep_recent: int) -> tuple[list[dict], list[dict]]:
    """Split messages into (to_summarize, to_keep), never breaking assistant+tool pairs.

    messages[0] is the system prompt and is excluded from both lists.
    The split ensures that if the boundary falls inside an assistant->tool sequence,
    we move the boundary earlier so the entire pair stays in to_keep.
    """
    if len(messages) <= keep_recent + 1:
        return [], messages

    body = messages[1:]
    split_idx = max(0, len(body) - keep_recent)

    # Walk backward from split_idx to avoid breaking an assistant+tool pair.
    # If messages[split_idx] is a tool result, include its preceding assistant message in to_keep.
    while split_idx > 0 and (
        body[split_idx].get("role") in ("tool",) or is_tool_result_carrier_message(body[split_idx])
    ):
        split_idx -= 1

    if split_idx <= 0:
        return [], messages

    to_summarize = body[:split_idx]
    to_keep = body[split_idx:]
    return to_summarize, [messages[0]] + to_keep


def _resolve_llm_params(
    base_url: str | None = None,
    api_key: str | None = None,
    model: str | None = None,
) -> tuple[str, str, str]:
    """Return (base_url, api_key, model) with config overrides applied."""
    if base_url is None or model is None:
        from agent_loop import get_backend_base_url, get_llmster_model
        if base_url is None:
            base_url = get_backend_base_url()
        if model is None:
            is_lmstudio = "1234" in (base_url or "")
            model = (get_llmster_model() or "local-model") if is_lmstudio else "local-model"
    if api_key is None:
        api_key = "lm-studio"
    try:
        from core.config import get_summarization_config
        cfg = get_summarization_config()
        cfg_model = (cfg.get("model") or "").strip()
        cfg_base_url = (cfg.get("base_url") or "").strip()
        cfg_api_key = (cfg.get("api_key") or "").strip()
        if cfg_model:
            model = cfg_model
        if cfg_base_url:
            base_url = cfg_base_url
        if cfg_api_key:
            api_key = cfg_api_key
    except Exception:
        pass
    return base_url, api_key, model


def _run_summary_llm(
    conversation_text: str,
    on_log=None,
    base_url: str | None = None,
    api_key: str | None = None,
    model: str | None = None,
) -> str:
    """Run the multi-stage LLM summarization pipeline.

    Returns the accepted summary text (never None — falls back to extractive excerpt).
    """
    base_url, api_key, model = _resolve_llm_params(base_url, api_key, model)

    user_primary = _SUMMARY_USER_TEMPLATE.format(conversation=conversation_text)
    user_retry = (
        "Below is a conversation transcript. Summarize the substance for an assistant continuing the work.\n\n"
        f"{conversation_text}"
    )

    def _accept_candidate(raw: str | None, source: str) -> str | None:
        sanitized = sanitize_summary_text(raw or "")
        ok, reason = _validate_summary_text(sanitized)
        if ok and sanitized:
            return sanitized
        if on_log and (raw or "").strip():
            on_log("WARN", f"Context summarization output rejected ({source}): {reason}")
        return None

    chosen: str | None = None

    try:
        from utility_model import utility_complete
        utility_messages = [
            {"role": "system", "content": _SUMMARY_SYSTEM_PROMPT},
            {"role": "user", "content": user_primary},
        ]
        t0 = time.monotonic()
        raw_u = utility_complete(
            utility_messages,
            max_tokens=_SUMMARY_MAX_TOKENS_PRIMARY,
            temperature=0.3,
            activity_operation="context_summarization",
        )
        chosen = _accept_candidate(raw_u, "utility_model")
        if chosen and on_log:
            elapsed = time.monotonic() - t0
            on_log(
                "AGENT",
                f"Context summarization complete via utility model ({elapsed:.1f}s, {len(chosen)} chars).",
            )
    except Exception:
        pass

    client = None
    try:
        from agent_loop import _get_openai_client
        client = _get_openai_client(base_url, api_key)
    except Exception as e:
        if on_log:
            on_log("WARN", f"Context summarization: could not create API client: {e}")

    if not chosen and client is not None:
        try:
            t0 = time.monotonic()
            response = client.chat.completions.create(
                model=model,
                messages=[
                    {"role": "system", "content": _SUMMARY_SYSTEM_PROMPT},
                    {"role": "user", "content": user_primary},
                ],
                max_tokens=_SUMMARY_MAX_TOKENS_PRIMARY,
                temperature=0.3,
                stream=False,
            )
            chosen = _accept_candidate(
                (response.choices[0].message.content or "").strip(),
                "chat_api",
            )
            if chosen and on_log:
                elapsed = time.monotonic() - t0
                on_log(
                    "AGENT",
                    f"Context summarization complete ({elapsed:.1f}s, {len(chosen)} chars).",
                )
        except Exception as e:
            if on_log:
                on_log("WARN", f"Context summarization primary API call failed: {e}")

    if not chosen and client is not None:
        try:
            t0 = time.monotonic()
            response = client.chat.completions.create(
                model=model,
                messages=[
                    {"role": "system", "content": _SUMMARY_RETRY_SYSTEM_PROMPT},
                    {"role": "user", "content": user_retry},
                ],
                max_tokens=_SUMMARY_MAX_TOKENS_RETRY,
                temperature=0.2,
                stream=False,
            )
            chosen = _accept_candidate(
                (response.choices[0].message.content or "").strip(),
                "chat_api_retry",
            )
            if chosen and on_log:
                elapsed = time.monotonic() - t0
                on_log(
                    "AGENT",
                    f"Context summarization complete after retry prompt ({elapsed:.1f}s, {len(chosen)} chars).",
                )
        except Exception as e:
            if on_log:
                on_log("WARN", f"Context summarization retry API call failed: {e}")

    if not chosen:
        chosen = _extractive_fallback_summary(conversation_text)
        if on_log:
            on_log(
                "WARN",
                "Context summarization fell back to excerpt (all LLM outputs were empty or unusable).",
            )

    return chosen


# ---------------------------------------------------------------------------
# Anchor-aware entry points
# ---------------------------------------------------------------------------


def compute_anchor_update(
    messages_with_row_ids: list[dict],
    keep_recent: int = 20,
    on_log=None,
    **llm_kw,
) -> dict[str, Any] | None:
    """Compute a context anchor from DB messages (with ``_row_id``).

    *messages_with_row_ids* should be the literal user/assistant history loaded
    via ``get_messages(thread_id, include_row_id=True)`` — **no** system prompt.
    The first message is always preserved literally (never folded into the
    summary), matching the ``_split_keeping_pairs`` invariant where
    ``messages[0]`` stays in ``to_keep``.

    Returns ``{"summary_text": str, "through_message_id": int}`` or ``None``
    when the conversation is too short to summarize.
    """
    if len(messages_with_row_ids) < 6:
        return None

    to_summarize, to_keep = _split_keeping_pairs(messages_with_row_ids, keep_recent)
    if len(to_summarize) < 4:
        return None

    summarized_with_ids = [m for m in to_summarize if "_row_id" in m]
    if not summarized_with_ids:
        return None
    through_message_id: int = summarized_with_ids[-1]["_row_id"]

    conversation_text = _format_messages_for_summary(to_summarize)
    if len(conversation_text) > _MAX_CHARS_FOR_SUMMARY_INPUT:
        conversation_text = (
            conversation_text[:_MAX_CHARS_FOR_SUMMARY_INPUT]
            + "\n\n...[older content truncated for summarization]"
        )

    if on_log:
        on_log(
            "AGENT",
            f"Context summarization: compressing {len(to_summarize)} older messages "
            f"({len(conversation_text)} chars) → keeping {len(to_keep) - 1} recent messages.",
        )

    summary_text = _run_summary_llm(conversation_text, on_log=on_log, **llm_kw)
    return {"summary_text": summary_text, "through_message_id": through_message_id}


def build_merged_history_for_main_agent(
    thread_id: str,
    literal_msgs_with_ids: list[dict],
) -> list[dict]:
    """Merge literal DB messages with the thread's context anchor for model input.

    Returns a message list suitable for the LLM — no ``_row_id`` keys,
    no system prompt (the worker adds that separately).

    Self-heals on orphan / corrupt anchors by deleting the row and falling back
    to the full literal history.
    """
    import chat_storage

    def _strip(m: dict) -> dict:
        return {k: v for k, v in m.items() if k != "_row_id"}

    if not literal_msgs_with_ids:
        anchor = chat_storage.get_anchor(thread_id)
        if anchor is not None:
            _log.warning("build_merged_history: empty messages but anchor exists for %s — deleting", thread_id)
            chat_storage.delete_anchor(thread_id)
        return []

    anchor = chat_storage.get_anchor(thread_id)
    if anchor is None:
        return [_strip(m) for m in literal_msgs_with_ids]

    through_id: int = anchor["through_message_id"]
    row_ids = [m["_row_id"] for m in literal_msgs_with_ids if "_row_id" in m]

    if not row_ids:
        _log.warning("build_merged_history: no _row_id in messages for %s — ignoring anchor", thread_id)
        chat_storage.delete_anchor(thread_id)
        return [_strip(m) for m in literal_msgs_with_ids]

    max_id = max(row_ids)
    row_id_set = set(row_ids)

    if through_id not in row_id_set or through_id >= max_id:
        _log.warning(
            "build_merged_history: stale/orphan anchor for %s (through_id=%d, max_id=%d) — deleting",
            thread_id, through_id, max_id,
        )
        chat_storage.delete_anchor(thread_id)
        return [_strip(m) for m in literal_msgs_with_ids]

    first_msg = literal_msgs_with_ids[0]
    summary_msg = _build_summary_message(anchor["summary"])
    tail = [m for m in literal_msgs_with_ids if m.get("_row_id", 0) > through_id]

    result: list[dict] = []
    if first_msg.get("_row_id", 0) <= through_id:
        result.append(_strip(first_msg))
    result.append(summary_msg)
    result.extend(_strip(m) for m in tail)
    return result


def summarize_older_messages(
    messages: list[dict],
    keep_recent: int = 20,
    on_log=None,
    base_url: str | None = None,
    api_key: str | None = None,
    model: str | None = None,
) -> list[dict] | None:
    """Summarize older messages via an LLM call and return the compressed message list.

    Returns a new list: [system_msg, summary_user_msg] + recent_tail.
    Returns None if summarization fails or there's nothing to summarize.

    This is the legacy entry point used by in-run context compression in
    ``agent_loop`` and by tests. For cross-turn persistence, prefer
    ``compute_anchor_update`` + ``chat_storage.upsert_anchor``.
    """
    to_summarize, to_keep = _split_keeping_pairs(messages, keep_recent)
    if len(to_summarize) < 4:
        return None

    conversation_text = _format_messages_for_summary(to_summarize)
    if len(conversation_text) > _MAX_CHARS_FOR_SUMMARY_INPUT:
        conversation_text = (
            conversation_text[:_MAX_CHARS_FOR_SUMMARY_INPUT]
            + "\n\n...[older content truncated for summarization]"
        )

    if on_log:
        on_log(
            "AGENT",
            f"Context summarization: compressing {len(to_summarize)} older messages "
            f"({len(conversation_text)} chars) → keeping {len(to_keep) - 1} recent messages.",
        )

    chosen = _run_summary_llm(
        conversation_text, on_log=on_log,
        base_url=base_url, api_key=api_key, model=model,
    )
    return [to_keep[0], _build_summary_message(chosen)] + to_keep[1:]
