"""Context management functions for the agent runner.

This module contains functions for estimating, compacting, and managing
conversation context during agent runs.
"""

from typing import Optional, TYPE_CHECKING

from ...utils.logger import log

if TYPE_CHECKING:
    from ..toolkit import AgentToolkit
    from ..events import AgentEventEmitter


def estimate_context_tokens(messages: list[dict], system_prompt: Optional[str] = None) -> int:
    """Estimate the current context window size in tokens.

    Args:
        messages: Conversation messages
        system_prompt: Optional system prompt to include in estimation

    Returns:
        Estimated token count for the context
    """
    total_chars = 0

    # Count characters in all messages
    for msg in messages:
        content = msg.get("content", "")
        if isinstance(content, str):
            total_chars += len(content)
        elif isinstance(content, list):
            # Handle multi-part messages (e.g., with images)
            for part in content:
                if isinstance(part, dict) and "text" in part:
                    total_chars += len(part["text"])

        # Add role overhead (~4 tokens per message for role/structure)
        total_chars += 16

    # Also count system prompt
    if system_prompt:
        total_chars += len(system_prompt)

    # Estimate: ~4 characters per token
    return total_chars // 4


def get_context_breakdown(
    messages: list[dict],
    system_prompt: Optional[str] = None,
) -> tuple[dict, list[dict]]:
    """Get breakdown of context usage by message type.

    Args:
        messages: Conversation messages
        system_prompt: Optional system prompt

    Returns:
        Tuple of (breakdown dict, list of largest messages)
    """
    breakdown = {
        "system_prompt": len(system_prompt) // 4 if system_prompt else 0,
        "user": 0,
        "assistant": 0,
        "tool_results": 0,
    }

    # Track individual message sizes for finding largest
    message_sizes = []

    for i, msg in enumerate(messages):
        role = msg.get("role", "unknown")
        content = msg.get("content", "")

        # Calculate content size
        if isinstance(content, str):
            size = len(content)
        elif isinstance(content, list):
            size = sum(len(p.get("text", "")) for p in content if isinstance(p, dict))
        else:
            size = 0

        tokens = size // 4

        # Categorize
        if role == "user":
            breakdown["user"] += tokens
        elif role == "assistant":
            breakdown["assistant"] += tokens
        elif role == "tool":
            breakdown["tool_results"] += tokens

        # Track for largest messages
        if tokens > 100:  # Only track substantial messages
            # Try to get a label for this message
            label = f"{role}[{i}]"
            if role == "tool":
                tool_call_id = msg.get("tool_call_id", "")
                # Try to find the tool name from previous assistant message
                for prev_msg in reversed(messages[:i]):
                    if prev_msg.get("role") == "assistant" and "tool_calls" in prev_msg:
                        for tc in prev_msg.get("tool_calls", []):
                            if tc.get("id") == tool_call_id:
                                label = tc.get("function", {}).get("name", "tool")
                                break
                        break

            message_sizes.append({
                "index": i,
                "role": role,
                "label": label,
                "tokens": tokens,
                "preview": content[:100] if isinstance(content, str) else str(content)[:100],
            })

    # Sort by size and get top 5
    message_sizes.sort(key=lambda x: x["tokens"], reverse=True)
    largest = message_sizes[:5]

    return breakdown, largest


def maybe_compact_context(
    messages: list[dict],
    provider: object,
    emitter: "AgentEventEmitter",
    system_prompt: Optional[str] = None,
    threshold: float = 0.8,
) -> list[dict]:
    """Proactively compact context if approaching limit.

    Args:
        messages: Current conversation messages
        provider: LLM provider instance
        emitter: Event emitter for notifications
        system_prompt: System prompt for token estimation
        threshold: Trigger compaction at this % of context limit (default 80%)

    Returns:
        Original or compacted messages
    """
    context_tokens = estimate_context_tokens(messages, system_prompt)
    context_limit = provider.get_context_limit()

    # Check if we need to compact
    if context_tokens < context_limit * threshold:
        return messages  # No compaction needed

    log.info(
        f"Context at {context_tokens:,}/{context_limit:,} tokens "
        f"({context_tokens/context_limit:.0%}), compacting..."
    )

    return compact_messages_with_llm(
        messages, emitter, target_tokens=int(context_limit * 0.5)
    )


def compact_messages_with_llm(
    messages: list[dict],
    emitter: "AgentEventEmitter",
    target_tokens: int,
) -> list[dict]:
    """Use fast LLM to summarize middle messages.

    Preserves:
    - First message (original user request)
    - Last 4 messages (recent context)
    - Summarizes everything in between

    Args:
        messages: Current conversation messages
        emitter: Event emitter for notifications
        target_tokens: Target token count after compaction

    Returns:
        Compacted messages list
    """
    from ..subagent import get_model_for_tier
    from ..providers import get_provider

    if len(messages) <= 5:
        return messages  # Too few to compact

    # Split messages
    first_msg = messages[0]
    recent_msgs = messages[-4:]
    middle_msgs = messages[1:-4]

    if not middle_msgs:
        return messages

    # Build summary prompt
    middle_content = format_messages_for_summary(middle_msgs)

    prompt = f"""Summarize this conversation history concisely.

PRESERVE (include verbatim if present):
- Code snippets and file paths
- Error messages
- Key decisions made
- Important tool results (file contents, search results)

CONDENSE:
- Repetitive searches
- Verbose tool outputs
- Intermediate reasoning

CONVERSATION HISTORY:
{middle_content}

OUTPUT FORMAT:
Provide a concise summary (max 2000 tokens) that captures the essential context needed to continue this task."""

    # Use fast model for summarization
    fast_model = get_model_for_tier("fast")
    fast_provider = get_provider(fast_model)

    try:
        emitter.emit_thinking("Compacting context with fast model...")

        response = fast_provider.chat(
            messages=[{"role": "user", "content": prompt}],
            system="You are a context summarizer. Be concise but preserve code and technical details.",
        )

        summary = response.content or ""

        log.info(
            f"Compacted {len(middle_msgs)} messages into summary "
            f"({len(summary)} chars)"
        )

        # Build compacted messages
        return [
            first_msg,
            {
                "role": "assistant",
                "content": f"[Context Summary]\n{summary}\n[End Summary]",
            },
            *recent_msgs,
        ]
    except Exception as e:
        log.warning(f"LLM compaction failed: {e}, falling back to truncation")
        return [first_msg] + recent_msgs


def format_messages_for_summary(messages: list[dict]) -> str:
    """Format messages for summarization prompt.

    Args:
        messages: Messages to format

    Returns:
        Formatted string for summarization
    """
    parts = []
    for msg in messages:
        role = msg.get("role", "unknown")
        content = msg.get("content", "")

        # Handle tool calls in assistant messages
        if role == "assistant" and "tool_calls" in msg:
            tool_calls = msg.get("tool_calls", [])
            tool_info = [
                f"Called: {tc.get('function', {}).get('name', 'unknown')}"
                for tc in tool_calls
            ]
            content = f"{content}\n[Tools: {', '.join(tool_info)}]" if content else f"[Tools: {', '.join(tool_info)}]"

        # Truncate very long content
        if len(content) > 4000:
            content = content[:4000] + "\n[...truncated...]"

        parts.append(f"[{role.upper()}]\n{content}")

    return "\n\n---\n\n".join(parts)


def get_reranked_context(
    toolkit: "AgentToolkit",
    current_query: str,
) -> dict:
    """Get reranked context items based on the current query.

    Args:
        toolkit: Agent toolkit instance
        current_query: Current query for relevance ranking

    Returns:
        Dict with item_count and items list
    """
    try:
        from ...context.service import ContextService
        from ...context.reranker import rerank_context_items

        # Get exploration steps for context extraction
        steps = toolkit.get_exploration_steps()
        if not steps:
            return {"item_count": 0, "items": []}

        # Use context service to extract context items from exploration
        service = ContextService(connection=toolkit.connection)
        terminal_id = service.get_terminal_id()

        # Update context with exploration steps
        service.update_context(
            terminal_id=terminal_id,
            exploration_steps=steps,
        )

        # Get context items
        items = service.get_context_items(terminal_id)
        if not items:
            return {"item_count": 0, "items": []}

        # Rerank by query relevance
        if current_query:
            items = rerank_context_items(
                items,
                current_query,
                top_k=20,
            )

        # Convert to serializable format
        result_items = []
        for item in items[:20]:  # Limit to 20 items
            result_items.append({
                "name": item.qualified_name,
                "type": item.entity_type,
                "file": item.file_path,
                "score": round(item.score, 3) if hasattr(item, 'score') else None,
            })

        return {
            "item_count": len(result_items),
            "items": result_items,
        }

    except Exception as e:
        log.debug(f"Failed to get reranked context: {e}")
        return {"item_count": 0, "items": []}


def emit_context_frame(
    toolkit: "AgentToolkit",
    emitter: "AgentEventEmitter",
    messages: list[dict],
    system_prompt: Optional[str],
    current_query: str,
    total_input_tokens: int,
    total_output_tokens: int,
) -> None:
    """Emit a context frame event with current exploration state.

    Args:
        toolkit: Agent toolkit instance
        emitter: Event emitter
        messages: Current conversation messages
        system_prompt: System prompt for estimation
        current_query: Current query for reranking
        total_input_tokens: Total input tokens used
        total_output_tokens: Total output tokens used
    """
    # Get exploration steps from toolkit session
    steps = toolkit.get_exploration_steps()

    # Estimate current context window tokens and get breakdown
    context_tokens = 0
    context_breakdown = {}
    largest_messages = []
    if messages:
        context_tokens = estimate_context_tokens(messages, system_prompt)
        context_breakdown, largest_messages = get_context_breakdown(messages, system_prompt)

    # Summarize exploration by tool
    tool_counts: dict[str, int] = {}
    entities_found = 0
    step_details: list[dict] = []

    for step in steps:
        tool_name = getattr(step, 'tool', 'unknown')
        tool_counts[tool_name] = tool_counts.get(tool_name, 0) + 1

        # Count entities from the step
        step_entities = getattr(step, 'entities_found', [])
        entities_found += len(step_entities)

        # Collect step details
        params = getattr(step, 'params', {})
        summary = getattr(step, 'result_summary', '')

        # Extract meaningful info based on tool type
        detail = {
            "tool": tool_name,
            "summary": summary,
        }

        # Add relevant params based on tool
        if tool_name == 'read_file' and 'file_path' in params:
            detail["file"] = params['file_path']
        elif tool_name == 'read_file' and 'path' in params:
            detail["file"] = params['path']
        elif tool_name in ('grep', 'semantic_search') and 'query' in params:
            detail["query"] = params['query']
        elif tool_name == 'glob' and 'pattern' in params:
            detail["pattern"] = params['pattern']
        elif tool_name == 'list_files' and 'path' in params:
            detail["path"] = params['path']

        # Add content preview if available
        content_preview = getattr(step, 'content_preview', None)
        if content_preview:
            detail["content_preview"] = content_preview

        # Add token count if available
        token_count = getattr(step, 'token_count', 0)
        if token_count > 0:
            detail["tokens"] = token_count

        # Add entities if any
        if step_entities:
            detail["entities"] = step_entities[:5]  # Limit to 5

        step_details.append(detail)

    exploration_steps = [
        {"tool": tool, "count": count}
        for tool, count in tool_counts.items()
    ]

    # Build context frame data
    adding = {
        "exploration_steps": exploration_steps,
        "entities_found": entities_found,
        "step_count": len(steps),
        "details": step_details[-20:],  # Last 20 steps
        "input_tokens": total_input_tokens,
        "output_tokens": total_output_tokens,
        "context_tokens": context_tokens,  # Current context window size
        "context_breakdown": context_breakdown,  # Tokens by message type
        "largest_messages": largest_messages,  # Top 5 biggest messages
    }

    # Get reranked context items
    reading = get_reranked_context(toolkit, current_query)

    # Emit the context frame
    emitter.emit_context_frame(adding=adding, reading=reading)
