"""Context compaction plugin for automatic conversation summarization.

Monitors conversation length and automatically compresses old messages
into a summary when thresholds are reached. Keeps recent interactions
intact and swaps history seamlessly between LLM turns.
"""

import asyncio
import json
import logging
from datetime import datetime, timezone
from pathlib import Path
from typing import Any, Dict, List, Optional
from uuid import uuid4

from kollabor_events import EventType, Hook, HookPriority
from kollabor_events.data_models import ConversationMessage
from kollabor_plugins import BasePlugin

logger = logging.getLogger(__name__)

# Summarization system prompt template
SUMMARIZATION_SYSTEM_PROMPT = (
    "You are a conversation summarizer. Compress the following "
    "conversation into a concise briefing that preserves all critical "
    "context for the AI assistant to continue the conversation seamlessly.\n\n"
    "Preserve:\n"
    "- Key decisions and conclusions\n"
    "- File paths, function names, technical specifics\n"
    "- Unresolved questions and pending tasks\n"
    "- The user's current goals and working context\n"
    "- Any preferences or patterns established\n\n"
    "Do NOT include:\n"
    "- Greetings or meta-commentary\n"
    "- Tool call details (just mention what tools accomplished)\n"
    "- Duplicate information\n\n"
    "Format: structured bullet points, dense, under {max_tokens} tokens."
)

SUMMARY_INJECTION_PREFIX = (
    "[Previous Context Summary]\n"
    "The following is a summary of our earlier conversation:\n\n"
)
SUMMARY_INJECTION_SUFFIX = "\n\nContinue from where we left off."


class ContextCompactionPlugin(BasePlugin):
    """Automatically compacts conversation history by summarizing old messages."""

    @staticmethod
    def get_default_config() -> Dict[str, Any]:
        return {
            "plugins": {
                "context_compaction": {
                    "enabled": True,
                    "trigger_threshold": 16,
                    "keep_recent": 4,
                    "re_trigger_threshold": 12,
                    "count_mode": "interactions",
                    "summarization_profile": None,
                    "max_summary_tokens": 2000,
                    "log_compaction_events": True,
                }
            }
        }

    @staticmethod
    def get_config_widgets() -> Dict[str, Any]:
        return {
            "title": "Context Compaction",
            "widgets": [
                {
                    "type": "checkbox",
                    "label": "Enabled",
                    "config_path": "plugins.context_compaction.enabled",
                    "help": "Enable automatic context compaction",
                },
                {
                    "type": "slider",
                    "label": "Trigger Threshold",
                    "config_path": "plugins.context_compaction.trigger_threshold",
                    "min_value": 4,
                    "max_value": 50,
                    "step": 2,
                    "help": "Number of interactions before first compaction",
                },
                {
                    "type": "slider",
                    "label": "Keep Recent",
                    "config_path": "plugins.context_compaction.keep_recent",
                    "min_value": 1,
                    "max_value": 20,
                    "step": 1,
                    "help": "Number of recent interactions to preserve",
                },
                {
                    "type": "slider",
                    "label": "Re-trigger Threshold",
                    "config_path": "plugins.context_compaction.re_trigger_threshold",
                    "min_value": 4,
                    "max_value": 30,
                    "step": 2,
                    "help": "Interactions after compaction before re-triggering",
                },
                {
                    "type": "dropdown",
                    "label": "Count Mode",
                    "config_path": "plugins.context_compaction.count_mode",
                    "options": ["interactions", "messages"],
                    "help": "How to count conversation length",
                },
                {
                    "type": "slider",
                    "label": "Max Summary Tokens",
                    "config_path": "plugins.context_compaction.max_summary_tokens",
                    "min_value": 500,
                    "max_value": 8000,
                    "step": 500,
                    "help": "Maximum tokens for compaction summary",
                },
                {
                    "type": "checkbox",
                    "label": "Log Compaction Events",
                    "config_path": "plugins.context_compaction.log_compaction_events",
                    "help": "Log compaction events to conversation history",
                },
            ],
        }

    def __init__(self, name: str, event_bus, renderer, config) -> None:
        self.name = name
        self.event_bus = event_bus
        self.renderer = renderer
        self.config = config

        # State
        self._compaction_round: int = 0
        self._compaction_in_progress: bool = False
        self._pending_compaction: Optional[List[ConversationMessage]] = None
        self._pre_compaction_len: int = 0
        self._pending_session_id: Optional[str] = None
        self._consecutive_failures: int = 0
        self._disabled_for_session: bool = False
        self._compaction_task: Optional[asyncio.Task] = None

        # References set during initialize()
        self._llm_service = None
        self._conversation_logger = None
        self._conversation_manager = None
        self._profile_manager = None

        logger.info(f"ContextCompactionPlugin initialized: {name}")

    async def initialize(
        self,
        args=None,
        event_bus=None,
        config=None,
        command_registry=None,
        input_handler=None,
        renderer=None,
        llm_service=None,
        conversation_logger=None,
        conversation_manager=None,
        **kwargs,
    ) -> None:
        if event_bus:
            self.event_bus = event_bus
        if config:
            self.config = config
        if renderer:
            self.renderer = renderer

        self._llm_service = llm_service
        self._conversation_logger = conversation_logger
        self._conversation_manager = conversation_manager

        if llm_service and hasattr(llm_service, "profile_manager"):
            self._profile_manager = llm_service.profile_manager

        # Register status widget if widget_api available
        self._register_status_widget()

        logger.info("Context compaction plugin initialized")

    async def register_hooks(self) -> None:
        if not self.config.get("plugins.context_compaction.enabled", True):
            logger.info("Context compaction disabled, skipping hook registration")
            return

        post_hook = Hook(
            name="context_compaction_post",
            plugin_name=self.name,
            event_type=EventType.LLM_REQUEST_POST,
            priority=HookPriority.POSTPROCESSING.value,
            callback=self._on_llm_turn_complete,
        )
        await self.event_bus.register_hook(post_hook)

        pre_hook = Hook(
            name="context_compaction_pre",
            plugin_name=self.name,
            event_type=EventType.LLM_REQUEST_PRE,
            priority=HookPriority.PREPROCESSING.value,
            callback=self._apply_pending_compaction,
        )
        await self.event_bus.register_hook(pre_hook)

        logger.info("Context compaction hooks registered")

    async def shutdown(self) -> None:
        if self._compaction_task and not self._compaction_task.done():
            self._compaction_task.cancel()
            try:
                await self._compaction_task
            except asyncio.CancelledError:
                pass
        self._pending_compaction = None
        self._pending_session_id = None
        logger.info("Context compaction plugin shutdown")

    # ------------------------------------------------------------------
    # Status widget
    # ------------------------------------------------------------------

    def _register_status_widget(self) -> None:
        """Register a status widget via the widget_api service."""
        if not self.event_bus:
            return
        widget_api = self.event_bus.get_service("widget_api")
        if not widget_api:
            logger.debug("widget_api service not available, skipping widget")
            return
        try:
            widget_api.register_widget(
                id="context-compaction",
                name="Context",
                description="Context compaction status",
                render_fn=self._render_widget,
                default_width="auto",
                min_width=8,
            )
            logger.info("Registered context-compaction status widget")

            # Ensure widget is in the layout (row 1) if not already placed
            layout_mgr = self.event_bus.get_service("layout_manager")
            if layout_mgr:
                layout = layout_mgr.get_layout()
                already_placed = any(
                    w.id == "context-compaction"
                    for row in layout.rows
                    for w in row.widgets
                )
                if not already_placed:
                    layout_mgr.add_widget_to_row(1, "context-compaction")
                    logger.info("Added context-compaction widget to row 1")

        except Exception as e:
            logger.warning(f"Failed to register status widget: {e}")

    def _render_widget(self, width: int, context) -> str:
        """Render the context compaction status widget."""
        history = self._get_conversation_history()
        count_mode = self.config.get(
            "plugins.context_compaction.count_mode", "interactions"
        )
        interaction_count = 0
        if history:
            interaction_count = self._count_interactions(history, count_mode)

        if self._disabled_for_session:
            return "ctx: off"
        if self._compaction_in_progress:
            return "ctx: compacting..."
        if self._compaction_round > 0:
            return f"ctx: r{self._compaction_round}"

        return f"ctx: {interaction_count}"

    # ------------------------------------------------------------------
    # Hooks
    # ------------------------------------------------------------------

    async def _on_llm_turn_complete(
        self, data: Dict[str, Any], event
    ) -> Dict[str, Any]:
        """LLM_REQUEST_POST: check if compaction threshold reached."""
        if self._disabled_for_session or self._compaction_in_progress:
            return data

        if not self.config.get("plugins.context_compaction.enabled", True):
            return data

        history = self._get_conversation_history()
        if not history:
            return data

        count_mode = self.config.get(
            "plugins.context_compaction.count_mode", "interactions"
        )
        current_count = self._count_interactions(history, count_mode)

        threshold = self._get_current_threshold()

        if current_count >= threshold:
            logger.info(
                f"Compaction threshold reached: {current_count} >= {threshold} "
                f"(round {self._compaction_round})"
            )
            self._compaction_in_progress = True
            self._compaction_task = asyncio.ensure_future(self._run_compaction())

        return data

    async def _apply_pending_compaction(
        self, data: Dict[str, Any], event
    ) -> Dict[str, Any]:
        """LLM_REQUEST_PRE: apply staged compaction atomically."""
        if self._pending_compaction is None:
            return data

        history = self._get_conversation_history()
        if not history:
            self._pending_compaction = None
            self._pending_session_id = None
            self._pre_compaction_len = 0
            return data

        pending_session_id = self._pending_session_id
        current_session_id = self._get_current_session_id()
        if pending_session_id and current_session_id != pending_session_id:
            logger.info("Discarded staged compaction due to session change")
            self._pending_compaction = None
            self._pending_session_id = None
            self._pre_compaction_len = 0
            return data

        if len(history) < self._pre_compaction_len:
            logger.info(
                "Discarded staged compaction because conversation history was reset"
            )
            self._pending_compaction = None
            self._pending_session_id = None
            self._pre_compaction_len = 0
            return data

        # Capture any messages added during background compaction
        new_msgs = history[self._pre_compaction_len:]

        # Atomic swap via slice assignment
        pre_len = len(history)
        history[:] = self._pending_compaction + new_msgs
        post_len = len(history)

        logger.info(
            f"Applied compaction round {self._compaction_round}: "
            f"{pre_len} -> {post_len} messages"
        )

        # Sync conversation_manager if available
        if self._conversation_manager:
            self._conversation_manager._update_context_window()

        self._pending_compaction = None
        self._pending_session_id = None
        self._pre_compaction_len = 0

        return data

    # ------------------------------------------------------------------
    # Core compaction logic
    # ------------------------------------------------------------------

    def _get_conversation_history(self) -> Optional[List[ConversationMessage]]:
        """Get the live conversation_history list from llm_service."""
        if self._llm_service and hasattr(self._llm_service, "conversation_history"):
            return self._llm_service.conversation_history
        return None

    def _get_current_session_id(self) -> Optional[str]:
        """Return current session ID for stale-stage detection."""
        if self._conversation_logger and hasattr(self._conversation_logger, "session_id"):
            return self._conversation_logger.session_id
        return None

    def _get_current_threshold(self) -> int:
        if self._compaction_round == 0:
            return self.config.get(
                "plugins.context_compaction.trigger_threshold", 16
            )
        return self.config.get(
            "plugins.context_compaction.re_trigger_threshold", 12
        )

    def _count_interactions(
        self, history: List[ConversationMessage], mode: str = "interactions"
    ) -> int:
        """Count interactions or messages in history."""
        if mode == "messages":
            return len(history)
        if mode == "interactions":
            return sum(1 for msg in history if msg.role == "user")
        return len(history)

    def _find_split_point(
        self, history: List[ConversationMessage], keep_recent: int
    ) -> int:
        """Find the split point for compaction.

        Returns the index of the first message to keep (keep_recent count).
        Adjusts the split to never break an assistant+tool_calls / tool_result
        group, which would cause 400 errors from OpenAI and compatible APIs.
        """
        if len(history) <= keep_recent:
            return 0

        split = len(history) - keep_recent

        # If the split lands on a tool-role message, walk backward to include
        # the preceding assistant message that owns the tool_calls.
        while split > 0 and self._is_tool_result(history[split]):
            split -= 1

        # If we landed on an assistant message with tool_calls, the tool
        # results are in to_keep but the assistant isn't -- pull it in too.
        if split > 0 and self._has_tool_calls(history[split - 1]):
            # Actually we need to NOT split here -- the assistant at split-1
            # has tool_calls whose results start at split. Move split back
            # to include the assistant message.
            split -= 1
            # Walk back further in case this assistant is itself preceded
            # by tool results from an even earlier call (shouldn't happen
            # normally, but be safe).
            while split > 0 and self._is_tool_result(history[split]):
                split -= 1

        return max(split, 0)

    @staticmethod
    def _is_tool_result(msg: ConversationMessage) -> bool:
        """Check if a message is a tool result."""
        return msg.role == "tool" or msg.metadata.get("tool_call_id") is not None

    @staticmethod
    def _has_tool_calls(msg: ConversationMessage) -> bool:
        """Check if an assistant message contains tool_calls."""
        return (
            msg.role == "assistant"
            and bool(msg.metadata.get("tool_calls"))
        )

    def _format_messages_for_summary(
        self, messages: List[ConversationMessage]
    ) -> str:
        """Format messages into a text block for summarization.

        Skips tool-role messages and assistant tool_calls metadata since
        the summarizer only needs the high-level conversation flow.
        Tool outputs are captured by the assistant's follow-up response.
        """
        parts = []
        for msg in messages:
            # Skip tool result messages -- their output is echoed by the
            # assistant in its follow-up, so including them is redundant
            # and can confuse the summarizer with raw JSON.
            if self._is_tool_result(msg):
                continue

            role_label = msg.role.capitalize()
            content = msg.content
            if len(content) > 2000:
                content = content[:2000] + "... [truncated]"

            # For assistant messages with tool_calls, note which tools ran
            if self._has_tool_calls(msg):
                tool_names = [
                    tc.get("function", {}).get("name", tc.get("name", "?"))
                    for tc in msg.metadata["tool_calls"]
                ]
                content += f"\n[Used tools: {', '.join(tool_names)}]"

            parts.append(f"[{role_label}]: {content}")
        return "\n\n".join(parts)

    async def _run_compaction(self) -> None:
        """Background task: summarize old messages and stage compaction."""
        try:
            history = self._get_conversation_history()
            if not history:
                return
            history_snapshot = list(history)
            snapshot_len = len(history_snapshot)
            snapshot_session_id = self._get_current_session_id()

            keep_recent = self.config.get(
                "plugins.context_compaction.keep_recent", 4
            )
            max_summary_tokens = self.config.get(
                "plugins.context_compaction.max_summary_tokens", 2000
            )

            # Identify boundaries
            split_point = self._find_split_point(history_snapshot, keep_recent)
            if split_point <= 1:
                # Nothing meaningful to summarize
                logger.info("Not enough messages to summarize, skipping")
                return

            # Separate system message
            system_msg = None
            start_idx = 0
            if history_snapshot and history_snapshot[0].role == "system":
                system_msg = history_snapshot[0]
                start_idx = 1

            to_summarize = history_snapshot[start_idx:split_point]
            to_keep = history_snapshot[split_point:]

            if not to_summarize:
                logger.info("No messages to summarize after boundaries")
                return

            # Format and call LLM
            formatted = self._format_messages_for_summary(to_summarize)
            system_prompt = SUMMARIZATION_SYSTEM_PROMPT.format(
                max_tokens=max_summary_tokens
            )

            summary_text = await self._call_summarization_llm(
                formatted, system_prompt, max_summary_tokens
            )

            if not summary_text or len(summary_text) < 100:
                logger.warning(
                    f"Summary too short or empty ({len(summary_text) if summary_text else 0} chars), "
                    "treating as failure"
                )
                self._consecutive_failures += 1
                self._check_failure_threshold()
                return

            # Build compacted history
            compacted = self._build_compacted_history(
                system_msg, summary_text, to_keep
            )

            # Stage for safe application from the snapshot baseline.
            self._pre_compaction_len = snapshot_len
            self._pending_session_id = snapshot_session_id
            self._pending_compaction = compacted
            self._compaction_round += 1
            self._consecutive_failures = 0

            # Log compaction event
            await self._log_compaction_event(
                messages_summarized=len(to_summarize),
                messages_kept=len(to_keep),
                summary_length=len(summary_text),
                pre_count=snapshot_len,
                post_count=len(compacted),
            )

            logger.info(
                f"Compaction round {self._compaction_round} staged: "
                f"summarized {len(to_summarize)} msgs, keeping {len(to_keep)}"
            )

        except asyncio.CancelledError:
            logger.info("Compaction cancelled")
            raise
        except Exception as e:
            logger.error(f"Compaction failed: {e}", exc_info=True)
            self._consecutive_failures += 1
            self._check_failure_threshold()
        finally:
            self._compaction_in_progress = False

    async def _call_summarization_llm(
        self,
        formatted_messages: str,
        system_prompt: str,
        max_tokens: int,
    ) -> Optional[str]:
        """Call LLM to generate conversation summary."""
        from kollabor_ai.providers.registry import (
            ProviderRegistry,
            create_config_from_profile,
        )

        try:
            # Get profile for summarization
            profile_name = self.config.get(
                "plugins.context_compaction.summarization_profile", None
            )
            profile_data = None

            if profile_name and self._profile_manager:
                profile_obj = self._profile_manager.get_profile(profile_name)
                if profile_obj:
                    profile_data = profile_obj.to_dict()
                else:
                    logger.warning(
                        f"Summarization profile '{profile_name}' not found, "
                        "falling back to current"
                    )

            if not profile_data and self._profile_manager:
                profile_obj = self._profile_manager.get_active_profile()
                if profile_obj:
                    profile_data = profile_obj.to_dict()

            if not profile_data:
                logger.error("No profile available for summarization")
                return None

            # Create provider config and get provider
            provider_config = create_config_from_profile(profile_data)
            provider = await ProviderRegistry.get_provider(provider_config)

            # Build messages for the summarization call
            messages = [
                {"role": "system", "content": system_prompt},
                {
                    "role": "user",
                    "content": f"Summarize this conversation:\n\n{formatted_messages}",
                },
            ]

            # provider.call() returns UnifiedResponse
            # max_tokens is determined by the provider config from the profile
            response = await provider.call(messages=messages)

            # UnifiedResponse has get_text_content() helper
            if hasattr(response, "get_text_content"):
                return response.get_text_content()

            # Fallback for unexpected return types
            if isinstance(response, str):
                return response

            return str(response)

        except Exception as e:
            logger.error(f"Summarization LLM call failed: {e}", exc_info=True)
            return None

    def _build_compacted_history(
        self,
        system_msg: Optional[ConversationMessage],
        summary_text: str,
        to_keep: List[ConversationMessage],
    ) -> List[ConversationMessage]:
        """Build the new compacted history list."""
        compacted: List[ConversationMessage] = []

        if system_msg:
            compacted.append(system_msg)

        summary_content = (
            SUMMARY_INJECTION_PREFIX + summary_text + SUMMARY_INJECTION_SUFFIX
        )

        compacted.append(
            ConversationMessage(role="user", content=summary_content)
        )
        compacted.extend(to_keep)

        return compacted

    def _check_failure_threshold(self) -> None:
        """Disable compaction for session after 3 consecutive failures."""
        if self._consecutive_failures >= 3:
            logger.warning(
                "3 consecutive compaction failures, disabling for this session"
            )
            self._disabled_for_session = True

    # ------------------------------------------------------------------
    # JSONL logging
    # ------------------------------------------------------------------

    async def _log_compaction_event(
        self,
        messages_summarized: int,
        messages_kept: int,
        summary_length: int,
        pre_count: int,
        post_count: int,
    ) -> None:
        """Append compaction record to session JSONL."""
        if not self.config.get(
            "plugins.context_compaction.log_compaction_events", True
        ):
            return

        if not self._conversation_logger:
            return

        session_file = getattr(self._conversation_logger, "session_file", None)
        if not session_file or not isinstance(session_file, Path):
            return

        session_id = getattr(
            self._conversation_logger, "session_id", "unknown"
        )

        record = {
            "type": "context_compaction",
            "sessionId": session_id,
            "uuid": str(uuid4()),
            "timestamp": datetime.now(timezone.utc).isoformat(),
            "compaction_round": self._compaction_round,
            "messages_summarized": messages_summarized,
            "messages_kept": messages_kept,
            "summary_length": summary_length,
            "pre_message_count": pre_count,
            "post_message_count": post_count,
        }

        try:
            with open(session_file, "a") as f:
                f.write(json.dumps(record) + "\n")
            logger.debug(f"Logged compaction event for round {self._compaction_round}")
        except Exception as e:
            logger.warning(f"Failed to log compaction event: {e}")

    # ------------------------------------------------------------------
    # Status line
    # ------------------------------------------------------------------

    def get_status_lines(self) -> Dict[str, List[str]]:
        """Status line for area C showing compaction state."""
        if not self.config.get("plugins.context_compaction.enabled", True):
            return {"A": [], "B": [], "C": []}

        history = self._get_conversation_history()
        count_mode = self.config.get(
            "plugins.context_compaction.count_mode", "interactions"
        )
        interaction_count = 0
        if history:
            interaction_count = self._count_interactions(history, count_mode)

        if self._compaction_in_progress:
            status = "ctx: compacting..."
        elif self._compaction_round > 0:
            status = f"ctx: compacted (r{self._compaction_round})"
        else:
            status = f"ctx: {interaction_count}"

        if self._disabled_for_session:
            status += " [disabled]"

        return {"A": [], "B": [], "C": [status]}
