"""Compact agent for conversation summarization.

This module provides a compact agent that can summarize conversation history
and return structured results including analysis and context for continuing
the conversation.
"""

from __future__ import annotations

from collections.abc import Awaitable, Callable, Sequence
from dataclasses import replace
from inspect import isawaitable
from pathlib import Path
from typing import cast
from uuid import uuid4

from pydantic import BaseModel, Field
from pydantic_ai import Agent, ModelSettings, UserContent
from pydantic_ai.messages import (
    BinaryContent,
    ImageUrl,
    ModelMessage,
    ModelRequest,
    ModelResponse,
    SystemPromptPart,
    TextPart,
    ToolReturnPart,
    UserPromptPart,
    VideoUrl,
)
from pydantic_ai.models import Model
from pydantic_ai.tools import RunContext

from ya_agent_sdk._config import AgentSettings
from ya_agent_sdk._logger import logger
from ya_agent_sdk.agents.models import infer_model
from ya_agent_sdk.context import AgentContext, ModelConfig
from ya_agent_sdk.events import CompactCompleteEvent, CompactFailedEvent, CompactStartEvent
from ya_agent_sdk.filters import (
    create_system_prompt_filter,
    fix_truncated_tool_args,
)
from ya_agent_sdk.usage import InternalUsage
from ya_agent_sdk.utils import get_latest_request_usage

# =============================================================================
# Constants
# =============================================================================

AGENT_NAME = "compact"

DEFAULT_COMPACT_INSTRUCTION = """Use `condense` to generate a summary and context of the conversation so far.
This summary covers important details of the historical conversation with the user which has been truncated.
It's crucial that you respond by ONLY asking the user what you should work on next.
You should NOT take any initiative or make any assumptions about continuing with work.
Keep this response CONCISE and wrap your analysis in `analysis` and `context` fields to organize your thoughts and ensure you've covered all necessary points.

IMPORTANT: If the message history contains any access to Skills (files in /skills/ directory, such as reading SKILL.md or using skill resources), you MUST include a reminder in the context to re-read the relevant skill documentation when resuming work."""

# Anthropic cache settings that should NOT be inherited by compact agent.
# The compact agent has different system prompts, tools, and message history;
# inheriting these would create separate cache entries that are rarely reused,
# wasting cache write tokens.
_ANTHROPIC_CACHE_KEYS = frozenset({
    "anthropic_cache_tool_definitions",
    "anthropic_cache_instructions",
    "anthropic_cache_messages",
})

# Maximum characters to keep in a single tool return content for compact
_MAX_TOOL_RETURN_CHARS = 500
# Characters to keep from the beginning and end when truncating
_TOOL_RETURN_KEEP_HEAD = 200
_TOOL_RETURN_KEEP_TAIL = 200


def _strip_anthropic_cache_settings(settings: ModelSettings) -> ModelSettings:
    """Strip Anthropic cache settings from inherited model settings.

    The compact agent has different system prompts, tools, and message history
    from the main agent. Inheriting cache settings would create separate cache
    entries that are rarely reused, wasting cache write tokens.

    Args:
        settings: Model settings potentially containing Anthropic cache keys.

    Returns:
        A copy with cache keys removed, or the original if no cache keys present.
    """
    if not any(k in settings for k in _ANTHROPIC_CACHE_KEYS):
        return settings
    return cast(ModelSettings, {k: v for k, v in settings.items() if k not in _ANTHROPIC_CACHE_KEYS})


# =============================================================================
# Pre-trimming for compact
# =============================================================================


def _truncate_str(content: str, max_chars: int = _MAX_TOOL_RETURN_CHARS) -> str:
    """Truncate a string, keeping head and tail portions.

    Args:
        content: The string to potentially truncate.
        max_chars: Maximum allowed length before truncation.

    Returns:
        Original string if within limit, otherwise head + marker + tail.
    """
    if len(content) <= max_chars:
        return content

    head = content[:_TOOL_RETURN_KEEP_HEAD]
    tail = content[-_TOOL_RETURN_KEEP_TAIL:]
    truncated_count = len(content) - _TOOL_RETURN_KEEP_HEAD - _TOOL_RETURN_KEEP_TAIL
    return f"{head}\n[... {truncated_count} chars truncated ...]\n{tail}"


def _is_media_content(item: object) -> bool:
    """Check if a UserContent item is image or video media."""
    if isinstance(item, (ImageUrl, VideoUrl)):
        return True
    return isinstance(item, BinaryContent) and (
        item.media_type.startswith("image/") or item.media_type.startswith("video/")
    )


def _media_to_placeholder(item: object) -> str:
    """Convert a media content item to a descriptive text placeholder."""
    if isinstance(item, ImageUrl):
        return f"[image: {item.url}]"
    if isinstance(item, VideoUrl):
        return f"[video: {item.url}]"
    if isinstance(item, BinaryContent):
        return f"[{item.media_type} binary content removed]"
    return "[media content removed]"


def _truncate_tool_return(part: ToolReturnPart) -> tuple[ToolReturnPart, bool]:
    """Truncate a ToolReturnPart if its content exceeds the limit.

    Returns:
        A tuple of (possibly replaced part, whether it was modified).
    """
    content_str = part.model_response_str()
    if len(content_str) > _MAX_TOOL_RETURN_CHARS:
        return replace(part, content=_truncate_str(content_str)), True
    return part, False


def _strip_media_from_user_prompt(part: UserPromptPart) -> tuple[UserPromptPart, bool]:
    """Replace image/video content in a UserPromptPart with text placeholders.

    Returns:
        A tuple of (possibly replaced part, whether it was modified).
    """
    content = part.content

    # Handle direct media content (e.g., content=ImageUrl(...))
    if _is_media_content(content):
        return replace(part, content=_media_to_placeholder(content)), True

    if not isinstance(content, Sequence) or isinstance(content, str):
        return part, False

    has_media = any(_is_media_content(item) for item in content)
    if not has_media:
        return part, False

    replaced = [_media_to_placeholder(item) if _is_media_content(item) else item for item in content]
    return replace(part, content=replaced), True


def _trim_history_for_compact(
    message_history: list[ModelMessage],
) -> list[ModelMessage]:
    """Pre-trim message history before sending to compact agent.

    Performs two operations to reduce token count:
    1. Truncates large ToolReturnPart content (keeps head + tail)
    2. Drops all image/video content from UserPromptPart

    This ensures the compact agent can process the history without
    exceeding its own context window, even when triggered at high
    token usage (e.g., 89-99% of context).

    Args:
        message_history: The full message history from the main agent.

    Returns:
        A trimmed copy of the message history suitable for the compact agent.
    """
    trimmed: list[ModelMessage] = []

    for message in message_history:
        if not isinstance(message, ModelRequest):
            # ModelResponse - keep as-is (text and tool calls are small)
            trimmed.append(message)
            continue

        new_parts = []
        parts_modified = False

        for part in message.parts:
            if isinstance(part, ToolReturnPart):
                part, modified = _truncate_tool_return(part)
                parts_modified = parts_modified or modified
            elif isinstance(part, UserPromptPart):
                part, modified = _strip_media_from_user_prompt(part)
                parts_modified = parts_modified or modified
            # SystemPromptPart, RetryPromptPart - keep as-is
            new_parts.append(part)

        if parts_modified:
            message = replace(message, parts=new_parts)
        trimmed.append(message)

    return trimmed


# =============================================================================
# Utilities
# =============================================================================


def _load_system_prompt() -> str:
    """Load system prompt from the prompts directory."""
    prompt_path = Path(__file__).parent / "prompts" / "compact.md"
    if prompt_path.exists():
        return prompt_path.read_text()
    return ""


# =============================================================================
# Models
# =============================================================================


class CondenseResult(BaseModel):
    analysis: str = Field(
        ...,
        description="""A summary of the conversation so far, capturing technical details, code patterns, and architectural decisions.""",
    )
    context: str = Field(
        ...,
        description="""The context to continue the conversation with. If applicable based on the current task, this should include:

1. Primary Request and Intent: Capture all of the user's explicit requests and intents in detail
2. Key Technical Concepts: List all important technical concepts, technologies, and frameworks discussed.
3. Files and Code Sections: Enumerate specific files and code sections examined, modified, or created. Pay special attention to the most recent messages and include full code snippets where applicable and include a summary of why this file read or edit is important.
4. Problem Solving: Document problems solved and any ongoing troubleshooting efforts.
5. Pending Tasks: Outline any pending tasks that you have explicitly been asked to work on.
6. Current Work: Describe in detail precisely what was being worked on immediately before this summary request, paying special attention to the most recent messages from both user and assistant. Include file names and code snippets where applicable.
7. Optional Next Step: List the next step that you will take that is related to the most recent work you were doing. IMPORTANT: ensure that this step is DIRECTLY in line with the user's explicit requests, and the task you were working on immediately before this summary request. If your last task was concluded, then only list next steps if they are explicitly in line with the users request. Do not start on tangential requests without confirming with the user first.
8. If there is a next step, include direct quotes from the most recent conversation showing exactly what task you were working on and where you left off. This should be verbatim to ensure there's no drift in task interpretation.
""",
    )
    original_prompt: str = Field(
        ...,
        description="The original prompt and key information from the user. "
        "Used as fallback when agent_ctx.user_prompts is not set.",
    )
    auto_load_files: list[str] = Field(
        default_factory=list,
        description="File paths to auto-load when resuming. "
        "Files will be read and injected into context on next request. "
        "Use for: files being actively edited, key references needed to continue work.",
    )


# =============================================================================
# Agent Factory
# =============================================================================


def get_compact_agent(
    model: str | Model | None = None,
    model_settings: ModelSettings | None = None,
    main_model: str | Model | None = None,
    main_model_settings: ModelSettings | None = None,
) -> Agent[AgentContext, CondenseResult]:
    """Create a compact agent.

    Args:
        model: Model string or Model instance. Highest priority.
        model_settings: Optional model settings dict.
        main_model: Fallback model inherited from main agent. Lowest priority.
        main_model_settings: Fallback model settings inherited from main agent.

    Model resolution priority:
        1. model parameter (explicit configuration)
        2. YA_AGENT_COMPACT_MODEL environment variable
        3. main_model parameter (inherited from main agent)

    Returns:
        Agent configured for compact with AgentContext as deps type.

    Raises:
        ValueError: If no model is available from any source.
    """
    effective_model: str | Model | None = model
    effective_settings: ModelSettings | None = model_settings

    # Priority: model > env var > main_model
    if effective_model is None:
        settings = AgentSettings()
        if settings.compact_model:
            effective_model = settings.compact_model
        elif main_model is not None:
            effective_model = main_model
        else:
            raise ValueError(
                "No model specified. Provide model parameter, set YA_AGENT_COMPACT_MODEL, "
                "or pass main_model for inheritance."
            )

    # model_settings: model_settings > main_model_settings
    if effective_settings is None and main_model_settings is not None:
        effective_settings = _strip_anthropic_cache_settings(main_model_settings)

    system_prompt = _load_system_prompt()
    return Agent[AgentContext, CondenseResult](
        model=infer_model(effective_model),
        model_settings=effective_settings,
        output_type=CondenseResult,
        deps_type=AgentContext,
        system_prompt=system_prompt,
        history_processors=[
            create_system_prompt_filter(system_prompt),  # Ensure system prompt is consistent
            fix_truncated_tool_args,
        ],
    )


# =============================================================================
# Utilities
# =============================================================================


def condense_result_to_markdown(result: CondenseResult) -> str:
    """Convert CondenseResult to markdown format.

    Args:
        result: The CondenseResult to convert.

    Returns:
        Markdown formatted string with analysis and context.
    """
    return f"""## Condensed conversation summary

### Analysis

{result.analysis}

### Context

{result.context}
"""


def _need_compact(ctx: AgentContext, message_history: list[ModelMessage]) -> bool:
    """Check if compaction is needed based on token usage threshold.

    Args:
        ctx: Agent context with model configuration.
        message_history: Current message history.

    Returns:
        True if compaction should be triggered.
    """
    if not message_history:
        return False

    model_cfg = ctx.model_cfg
    if model_cfg.context_window is None:
        logger.debug("Unknown context window, skipping compact check.")
        return False

    # Get current token usage from message history
    request_usage = get_latest_request_usage(message_history)
    if request_usage is None or request_usage.total_tokens is None:
        return False

    threshold_tokens = int(model_cfg.context_window * model_cfg.compact_threshold)
    current_tokens = request_usage.total_tokens

    logger.debug(f"Compact check: {current_tokens} tokens vs {threshold_tokens} threshold")
    return current_tokens >= threshold_tokens


def _build_compacted_messages(
    summary: str,
    original_prompt: str | Sequence[UserContent],
    steering_messages: list[str] | None = None,
) -> list[ModelMessage]:
    """Build compacted message history.

    Args:
        summary: The compacted summary content.
        original_prompt: The initial user prompt.
        steering_messages: Additional steering messages from user during execution.

    Returns:
        List of ModelMessage representing the compacted history.
    """
    request_parts: list[SystemPromptPart | UserPromptPart] = [
        SystemPromptPart(content="Placeholder system prompt"),
        UserPromptPart(
            content="You have exceeded the maximum token limit for this conversation. "
            "Please provide a summary of the conversation so far and what you should work on next "
            "and I'll resume the conversation."
        ),
    ]

    # Build final request parts with original prompt and steering messages
    final_parts: list[UserPromptPart] = [UserPromptPart(content=original_prompt)]

    # Append steering messages if any
    if steering_messages:
        for steering in steering_messages:
            final_parts.append(UserPromptPart(content=f"[User Steering] {steering}"))

    final_parts.append(UserPromptPart(content="<compact-complete>Context compacted. Resume task.</compact-complete>"))

    return [
        ModelRequest(parts=request_parts),
        ModelResponse(parts=[TextPart(content=summary)]),
        ModelRequest(parts=final_parts),
    ]


def create_compact_filter(
    model: str | Model | None = None,
    model_settings: ModelSettings | None = None,
    model_cfg: ModelConfig | None = None,
    main_model: str | Model | None = None,
    main_model_settings: ModelSettings | None = None,
) -> Callable[[RunContext[AgentContext], list[ModelMessage]], Awaitable[list[ModelMessage]]]:
    """Create a compact filter for automatic context compaction.

    The returned filter checks token usage and compacts the conversation history
    when usage exceeds the configured threshold (ModelConfig.compact_threshold).

    Args:
        model: Model string or Model instance for the compact agent. Highest priority.
        model_settings: Optional model settings for the compact agent.
        model_cfg: Model configuration for threshold checking.
        main_model: Fallback model inherited from main agent. Lowest priority.
        main_model_settings: Fallback model settings inherited from main agent.

    Model resolution priority:
        1. model parameter (explicit configuration)
        2. YA_AGENT_COMPACT_MODEL environment variable
        3. main_model parameter (inherited from main agent)

    Returns:
        An async filter function compatible with pydantic-ai history_processors.

    Example::

        compact_filter = await create_compact_filter(model="openai:gpt-4o-mini")
        agent = Agent(
            'openai:gpt-4',
            deps_type=AgentContext,
            history_processors=[compact_filter],
        )
    """
    agent = get_compact_agent(
        model=model,
        model_settings=model_settings,
        main_model=main_model,
        main_model_settings=main_model_settings,
    )

    async def compact_filter(
        ctx: RunContext[AgentContext],
        message_history: list[ModelMessage],
    ) -> list[ModelMessage]:
        """Filter that compacts message history when threshold is exceeded.

        Args:
            ctx: Runtime context containing AgentContext.
            message_history: Current message history to potentially compact.

        Returns:
            Original or compacted message history.
        """
        agent_ctx = ctx.deps

        if not _need_compact(agent_ctx, message_history):
            logger.debug("No need to compact history.")
            return message_history

        logger.info("Compacting conversation history...")

        # Generate event_id to correlate start/complete events
        event_id = uuid4().hex[:8]

        # Apply model wrapper if configured
        original_model = agent.model
        if agent_ctx.model_wrapper is not None:
            wrapper_metadata = agent_ctx.get_wrapper_metadata()
            wrapped = agent_ctx.model_wrapper(cast(Model, original_model), AGENT_NAME, wrapper_metadata)
            agent.model = await wrapped if isawaitable(wrapped) else wrapped

        try:
            # Emit start event
            await agent_ctx.emit_event(CompactStartEvent(event_id=event_id, message_count=len(message_history)))

            # Pre-trim history to reduce token count for compact agent
            trimmed_history = _trim_history_for_compact(message_history)

            # Run compact agent on trimmed message history with AgentContext as deps
            result = await agent.run(
                DEFAULT_COMPACT_INSTRUCTION,
                message_history=trimmed_history,
                deps=AgentContext(
                    model_cfg=model_cfg or ModelConfig(),
                ),
            )

            # Record usage in extra_usages

            model_id = cast(Model, agent.model).model_name
            agent_ctx.add_extra_usage(
                agent=AGENT_NAME,
                internal_usage=InternalUsage(model_id=model_id, usage=result.usage()),
                uuid=uuid4().hex,
            )

            condense_result: CondenseResult = result.output

            # Append auto_load_files for the auto_load_files filter to process
            # Use extend instead of assignment to preserve any files set by external callers
            agent_ctx.auto_load_files.extend(condense_result.auto_load_files)

            # Build summary with condense result and user prompts
            condense_markdown = condense_result_to_markdown(condense_result)

            # Build compacted messages
            # Priority: agent_ctx.user_prompts > condense_result.original_prompt
            # user_prompts is set by main agent from actual user input, while original_prompt
            # is extracted by LLM from conversation history and may be less accurate
            compacted = _build_compacted_messages(
                condense_markdown,
                agent_ctx.user_prompts or condense_result.original_prompt,
                agent_ctx.steering_messages or None,
            )

            # Emit complete event with summary
            await agent_ctx.emit_event(
                CompactCompleteEvent(
                    event_id=event_id,
                    summary_markdown=condense_markdown,
                    original_message_count=len(message_history),
                    compacted_message_count=len(compacted),
                )
            )

            # Clear steering_messages after successful compact (content is now in summary)
            agent_ctx.steering_messages.clear()

            # Force downstream filters to inject instructions after context reset
            agent_ctx.force_inject_instructions = True

            logger.info(f"Compacted history from {len(message_history)} messages to {len(compacted)} messages")
            return compacted

        except Exception as e:
            logger.error(f"Failed to compact history: {e}")
            # Emit failed event so consumers know compact did not succeed
            await agent_ctx.emit_event(
                CompactFailedEvent(event_id=event_id, error=str(e), message_count=len(message_history))
            )
            # On error, return original history
            return message_history

        finally:
            # Restore original model to avoid side effects on shared agent
            agent.model = original_model

    return compact_filter
