"""Context compaction — auto-summarize long conversations to save tokens."""

import hashlib
import json
import os
import re
import time
from datetime import datetime
from typing import Callable, Dict, List, Optional, Tuple

from salmalm.constants import (
    COMPACTION_THRESHOLD,
    COMPLEX_INDICATORS,
    DATA_DIR,
    KST,
    MODEL_TIERS,
    SIMPLE_QUERY_MAX_CHARS,
    TOOL_HINT_KEYWORDS,
)
from salmalm.security.crypto import log


def _persist_compaction_summary(session_id: str, summary: str) -> None:
    """Save compaction summary to DB for cross-session restoration.

    OpenClaw-style: when a session is restored after restart, the last
    compaction summary is injected so the AI retains prior context.
    """
    if not summary or not session_id:
        return
    try:
        from salmalm.core.core import _get_db

        conn = _get_db()
        conn.execute("""CREATE TABLE IF NOT EXISTS compaction_summaries (
            session_id TEXT PRIMARY KEY,
            summary TEXT NOT NULL,
            created_at TEXT NOT NULL,
            message_count INTEGER DEFAULT 0
        )""")
        conn.execute(
            "INSERT OR REPLACE INTO compaction_summaries (session_id, summary, created_at) VALUES (?,?,?)",
            (session_id, summary[:10000], datetime.now(KST).isoformat()),
        )
        conn.commit()
    except Exception as e:
        log.warning(f"[PKG] Summary persist error: {e}")


def _restore_compaction_summary(session_id: str) -> Optional[str]:
    """Restore last compaction summary for a session (cross-session continuity)."""
    try:
        from salmalm.core.core import _get_db

        conn = _get_db()
        row = conn.execute(
            "SELECT summary FROM compaction_summaries WHERE session_id=?",
            (session_id,),
        ).fetchone()
        return row[0] if row else None
    except Exception as e:  # noqa: broad-except
        return None


def _msg_content_str(msg: dict) -> str:
    """Extract text content from a message (handles list content blocks)."""
    c = msg.get("content", "")
    if isinstance(c, str):
        return c
    if isinstance(c, list):
        return " ".join(b.get("text", "") for b in c if isinstance(b, dict) and b.get("type") == "text")
    return str(c)


def _stage1_strip_images(messages: list) -> list:
    """Stage 1: Strip binary/image data from old messages."""
    trimmed = []
    for m in messages:
        if m["role"] == "user" and isinstance(m.get("content"), list):
            new_content = []
            for block in m["content"]:
                if isinstance(block, dict) and block.get("type") in ("image", "image_url"):
                    new_content.append({"type": "text", "text": "[Image attached]"})
                else:
                    new_content.append(block)
            trimmed.append({**m, "content": new_content})
        else:
            trimmed.append(m)
    return trimmed


def _stage2_trim_tool_results(trimmed: list) -> list:
    """Stage 2: Trim long tool results to 500 chars."""
    result = []
    for m in trimmed:
        if m["role"] == "tool" and len(_msg_content_str(m)) > 500:
            result.append({**m, "content": _msg_content_str(m)[:500] + "\n... [truncated]"})
        elif m["role"] == "user" and isinstance(m.get("content"), list):
            new_blocks = []
            for block in m["content"]:
                if isinstance(block, dict) and block.get("type") == "tool_result":
                    content = block.get("content", "")
                    if isinstance(content, str) and len(content) > 500:
                        new_blocks.append({**block, "content": content[:500] + "\n... [truncated]"})
                    else:
                        new_blocks.append(block)
                else:
                    new_blocks.append(block)
            result.append({**m, "content": new_blocks})
        else:
            result.append(m)
    return result


def _stage3_split_messages(trimmed: list) -> tuple:
    """Stage 3: Split into system, old important (user/assistant), and recent."""
    system_msgs = [m for m in trimmed if m["role"] == "system"]
    non_system = [m for m in trimmed if m["role"] != "system"]
    recent = non_system[-10:]
    old = non_system[:-10]
    old_important = [m for m in old if m["role"] in ("user", "assistant")]
    return system_msgs, old_important, recent


def _stage4_truncate_old(old_important: list) -> list:
    """Stage 4: Truncate verbose old messages."""
    result = []
    for m in old_important:
        txt = _msg_content_str(m)
        if m["role"] == "assistant" and len(txt) > 800:
            result.append({**m, "content": txt[:800] + "\n... [compacted]"})
        elif m["role"] == "user" and len(txt) > 500:
            result.append({**m, "content": txt[:500] + "\n... [compacted]"})
        else:
            result.append(m)
    return result


def _stage5_llm_summarize(
    to_summarize: list, system_msgs: list, recent: list, stage4: list, messages: list, total_chars: int, session
) -> list:
    """Stage 5: LLM summarization of old context."""
    if not to_summarize:
        return system_msgs + recent

    summary_parts = []
    for m in to_summarize[-30:]:
        role = m["role"]
        txt = _msg_content_str(m)[:400]
        if txt.strip():
            summary_parts.append(f"[{role}]: {txt}")

    from salmalm.core.llm import call_llm
    from salmalm.core.core import router

    summary_model = router._pick_available(1)
    _summ_msgs = [
        {
            "role": "system",
            "content": (
                "Summarize the following conversation concisely but thoroughly. "
                "You MUST preserve:\n1. Key decisions and conclusions\n"
                "2. Task progress and what was accomplished\n"
                "3. Important facts, numbers, file paths, code context\n"
                "4. User preferences and constraints mentioned\n"
                "5. Any pending/blocked items\n"
                "Write in the same language as the conversation. "
                "Use 5-15 sentences. Do NOT start with 'The conversation...' — "
                "write as a factual summary that can serve as context for continued work."
            ),
        },
        {"role": "user", "content": "\n".join(summary_parts)},
    ]
    try:
        summary_result = call_llm(_summ_msgs, model=summary_model, max_tokens=1200)
    except Exception as e:
        log.error(f"[PKG] Stage 5 LLM failed: {e} — using stage 4 result")
        return stage4

    summary_content = summary_result.get("content", "")
    if not summary_content or len(summary_content) < 20:
        log.warning("[PKG] Stage 5 produced empty/short summary — using stage 4 result")
        return stage4

    original_chars = sum(len(_msg_content_str(m)) for m in to_summarize)
    if len(summary_content) > original_chars:
        log.warning(f"[PKG] Summary ({len(summary_content)}) > original ({original_chars}) — using stage 4")
        return stage4

    compacted = (
        system_msgs
        + [
            {
                "role": "system",
                "content": f"[Previous conversation summary — {len(to_summarize)} messages compacted]\n{summary_content}",
            }
        ]
        + recent
    )

    if session:
        try:
            _persist_compaction_summary(getattr(session, "id", ""), summary_content)
        except Exception as e:
            log.warning(f"[PKG] Summary persistence error: {e}")

    log.info(
        f"[PKG] Stage 5 compacted: {len(messages)} -> {len(compacted)} messages, {total_chars} → {sum(len(_msg_content_str(m)) for m in compacted)} chars"
    )

    try:
        result = memory_manager.auto_curate(days_back=3)
        if "No new" not in result:
            log.info(f"[MEM] Post-compaction auto-curate: {result}")
    except Exception as e:
        log.warning(f"[MEM] Auto-curate error: {e}")

    return compacted


def compact_messages(
    messages: list,
    model: Optional[str] = None,
    session: Optional["Session"] = None,
    on_status: Optional[Callable] = None,
) -> list:
    """Multi-stage compaction: trim tool results → drop old tools → summarize.

    OpenClaw-quality compaction with 5 stages:
      1. Strip binary/image data from old messages
      2. Trim long tool results (keep first 500 chars)
      3. Drop old tool messages, keep user/assistant
      4. Truncate verbose old assistant messages
      5. LLM summarization of remaining old context

    Hard limits: max 100 messages, max 500K chars (≈125K tokens).
    Pre-compaction: flush memory to persistent store.
    Post-compaction: inject summary as system message, not user message.
    """
    # Empty conversation guard
    if not messages:
        return messages

    MAX_MESSAGES = 100
    MAX_CHARS = 500_000

    # OpenClaw-style: pre-compaction memory flush
    total_chars_check = sum(len(_msg_content_str(m)) for m in messages)
    if session and total_chars_check > COMPACTION_THRESHOLD * 0.8:  # noqa: F405
        try:
            memory_manager.flush_before_compaction(session)
        except Exception as e:
            log.warning(f"[MEM] Memory flush error: {e}")

    # Hard message count limit
    if len(messages) > MAX_MESSAGES:
        system_msgs = [m for m in messages if m["role"] == "system"][:1]
        recent = [m for m in messages if m["role"] != "system"][-40:]
        messages = system_msgs + recent
        log.warning(f"[CUT] Hard msg limit: truncated to {len(messages)} messages")

    total_chars = sum(len(_msg_content_str(m)) for m in messages)

    # Hard char limit — emergency truncation
    if total_chars > MAX_CHARS:
        system_msgs = [m for m in messages if m["role"] == "system"][:1]
        recent = [m for m in messages if m["role"] != "system"][-20:]
        messages = system_msgs + recent
        total_chars = sum(len(_msg_content_str(m)) for m in messages)
        log.warning(f"[CUT] Hard char limit: truncated to {len(messages)} msgs ({total_chars} chars)")

    if total_chars < COMPACTION_THRESHOLD:  # noqa: F405
        return messages

    # Notify UI that compaction is in progress
    if on_status:
        try:
            on_status("compacting", "✨ Compacting context...")
        except Exception as e:
            log.debug(f"Suppressed: {e}")

    log.info(f"[PKG] Compacting {len(messages)} messages ({total_chars} chars)")

    trimmed = _stage1_strip_images(messages)
    trimmed = _stage2_trim_tool_results(trimmed)
    total_after_trim = sum(len(_msg_content_str(m)) for m in trimmed)
    if total_after_trim < COMPACTION_THRESHOLD:  # noqa: F405
        log.info(f"[PKG] Stage 2 sufficient: {total_chars} -> {total_after_trim} chars")
        return trimmed

    system_msgs, old_important, recent = _stage3_split_messages(trimmed)
    stage3 = system_msgs + old_important + recent
    total_after_drop = sum(len(_msg_content_str(m)) for m in stage3)
    if total_after_drop < COMPACTION_THRESHOLD:  # noqa: F405
        log.info(f"[PKG] Stage 3 sufficient: {total_chars} -> {total_after_drop} chars")
        return stage3

    stage4_old = _stage4_truncate_old(old_important)
    stage4 = system_msgs + stage4_old + recent
    total_after_trunc = sum(len(_msg_content_str(m)) for m in stage4)
    if total_after_trunc < COMPACTION_THRESHOLD:  # noqa: F405
        log.info(f"[PKG] Stage 4 sufficient: {total_chars} -> {total_after_trunc} chars")
        return stage4

    return _stage5_llm_summarize(stage4_old, system_msgs, recent, stage4, messages, total_chars, session)


# ============================================================
import math  # noqa: F811


# TFIDFSearch extracted to salmalm/core/search.py
from salmalm.core.search import TFIDFSearch  # noqa: E402

_tfidf = TFIDFSearch()

# ============================================================
# LLM CRON MANAGER — Scheduled tasks with LLM execution
# ============================================================
# LLMCronManager extracted to salmalm/core/llm_cron.py
from salmalm.core.llm_cron import LLMCronManager  # noqa: E402


def _estimate_tokens(messages: list) -> int:
    """Estimate token count using chars/4 approximation (stdlib only)."""
    total_chars = sum(len(_msg_content_str(m)) for m in messages)
    return total_chars // 4


def estimate_tokens(text_or_messages) -> int:
    """Estimate tokens — accepts string or message list."""
    if isinstance(text_or_messages, str):
        return max(1, len(text_or_messages) // 4)
    return _estimate_tokens(text_or_messages)


# ── Enhanced compaction helpers (from original compaction.py) ──
from salmalm.core.cost import estimate_tokens  # noqa: F401


def _msg_text(msg: dict) -> str:
    """Extract text content from a message dict."""
    c = msg.get("content", "")
    if isinstance(c, str):
        return c
    if isinstance(c, list):
        return " ".join(p.get("text", "") for p in c if isinstance(p, dict))
    return str(c)


def _importance_score(msg: dict) -> float:
    """Score message importance for compaction prioritization (0-10 scale)."""
    role = msg.get("role", "")
    text = _msg_text(msg).lower()
    if role == "system":
        return 10.0
    if role == "tool":
        if "error" in text[:200]:
            return 3.0
        return 0.5
    # User preference/decision messages
    if any(
        kw in text
        for kw in (
            "always",
            "never",
            "remember",
            "decide",
            "conclusion",  # noqa: E127
            "important",
            "approved",
            "항상",
            "결정",
            "결론",
            "기억",
        )
    ):  # noqa: E127
        return 5.0 if role == "user" else 4.0
    if role == "user":
        return 3.0
    return 2.0


def _split_by_importance(messages: List[dict], keep_recent: int = 2) -> tuple:
    """Split messages into (system, old, recent) for compaction.

    Returns (system_msgs, old_msgs_to_summarize, recent_msgs_to_keep).
    """
    system = [m for m in messages if m.get("role") == "system"]
    non_system = [m for m in messages if m.get("role") != "system"]
    if len(non_system) <= keep_recent:
        return system, [], non_system
    recent = non_system[-keep_recent:]
    old = non_system[:-keep_recent]
    return system, old, recent


def _extract_key_facts(messages: List[dict]) -> List[str]:
    """Extract key facts from messages for summary."""
    facts = []
    for m in messages:
        text = _msg_text(m)
        if not text.strip() or len(text) < 10:
            continue
        # Preference/decision indicators
        indicators = (
            "always",
            "never",
            "remember",
            "must",
            "decided",
            "conclusion",
            "result",
            "error",
            "fixed",
            "created",
            "deleted",
            "updated",
            "use ",
            "prefer",
            "항상",
            "결정",
            "결론",
            "수정",
            "생성",
            "삭제",
        )
        lower = text.lower()
        if any(kw in lower for kw in indicators):
            # Use first 200 chars as the fact
            facts.append(text[:200].strip())
    return facts[:20]


def enhanced_compact(messages: List[dict], target_tokens: int = 4000, max_tokens: int = 0) -> List[dict]:
    """Enhanced compaction using importance scoring."""
    from salmalm.core.cost import estimate_tokens as _et

    limit = target_tokens or max_tokens or 4000
    total = sum(_et(_msg_text(m)) for m in messages)
    if total <= limit:
        return messages
    system, old, recent = _split_by_importance(messages, keep_recent=max(2, len(messages) // 3))
    facts = _extract_key_facts(old)
    summary_msg = {
        "role": "system",
        "content": f"[Compacted {len(old)} messages]\nKey facts: " + ("; ".join(facts) if facts else "None extracted"),
    }
    return system + [summary_msg] + recent
