import sys
from typing import TYPE_CHECKING, Any, Awaitable, Callable

from zrb.config.config import CFG
from zrb.context.any_context import zrb_print
from zrb.llm.agent.summarizer import create_summarizer_agent
from zrb.llm.config.limiter import LLMLimiter
from zrb.llm.config.limiter import llm_limiter as default_llm_limiter
from zrb.util.cli.style import stylize_error, stylize_yellow
from zrb.util.markdown import make_markdown_section

if TYPE_CHECKING:
    from pydantic_ai.messages import ModelMessage
else:
    ModelMessage = Any


def is_turn_start(msg: Any) -> bool:
    """Identify start of a new user interaction (User Prompt without Tool Return)."""
    from pydantic_ai.messages import ModelRequest, ToolReturnPart, UserPromptPart

    if not isinstance(msg, ModelRequest):
        return False
    # In pydantic_ai, ModelRequest parts can be list of various parts
    has_user = any(isinstance(p, UserPromptPart) for p in msg.parts)
    has_return = any(isinstance(p, ToolReturnPart) for p in msg.parts)
    return has_user and not has_return


def message_to_text(msg: Any) -> str:
    """Convert a pydantic_ai message to a readable text representation for summarization."""
    from pydantic_ai.messages import (
        ModelRequest,
        ModelResponse,
        TextPart,
        ToolCallPart,
        ToolReturnPart,
        UserPromptPart,
    )

    if isinstance(msg, ModelRequest):
        parts = []
        for p in msg.parts:
            if isinstance(p, UserPromptPart):
                parts.append(f"User: {p.content}")
            elif isinstance(p, ToolReturnPart):
                parts.append(f"Tool Result ({p.tool_name}): {p.content}")
        return "\n".join(parts)
    if isinstance(msg, ModelResponse):
        parts = []
        for p in msg.parts:
            if isinstance(p, TextPart):
                parts.append(f"AI: {p.content}")
            elif isinstance(p, ToolCallPart):
                parts.append(f"AI Tool Call: {p.tool_name}({p.args})")
        return "\n".join(parts)
    return str(msg)


async def summarize_history(
    messages: "list[ModelMessage]",
    agent: Any = None,
    summary_window: int | None = None,
) -> "list[ModelMessage]":
    """
    Summarizes the history, keeping the last `summary_window` messages intact.
    Returns a new list of messages where older messages are replaced by a summary.
    """
    from pydantic_ai.messages import ModelRequest, UserPromptPart

    if summary_window is None:
        summary_window = CFG.LLM_HISTORY_SUMMARIZATION_WINDOW
    if len(messages) <= summary_window:
        summary_window = len(messages) - 1

    # Determine split index
    # We want to keep at least summary_window messages.
    # So split_idx <= len(messages) - summary_window.
    # We search backwards from there for a clean turn start.
    start_search_idx = max(0, len(messages) - summary_window)
    split_idx = -1

    # Iterate backwards from start_search_idx to 0
    for i in range(start_search_idx, -1, -1):
        if is_turn_start(messages[i]):
            split_idx = i
            break

    if split_idx <= 0:
        zrb_print(stylize_error("  Cannot split conversation..."), plain=True)
        return messages

    to_summarize = messages[:split_idx]
    to_keep = messages[split_idx:]

    # Use improved text representation
    history_text = "\n".join([message_to_text(m) for m in to_summarize])

    summarizer_agent = agent or create_summarizer_agent()

    # Run the summarizer agent
    result = await summarizer_agent.run(
        f"Summarize this conversation history:\n{history_text}"
    )
    summary_text = result.output

    # Create a summary message injected as user context
    try:
        zrb_print(stylize_yellow("Compressiong conversation..."), plain=True)
        summary_message = ModelRequest(
            parts=[
                UserPromptPart(
                    content=make_markdown_section(
                        "Previous conversation summary", summary_text
                    )
                )
            ]
        )
    except Exception as e:
        zrb_print(stylize_error(f"  Failed to compress conversation: {e}"), plain=True)
        return messages
    return [summary_message] + to_keep


def create_summarizer_history_processor(
    agent: Any = None,
    limiter: LLMLimiter | None = None,
    token_threshold: int | None = None,
    summary_window: int | None = None,
) -> "Callable[[list[ModelMessage]], Awaitable[list[ModelMessage]]]":
    """
    Creates a history processor that auto-summarizes history when it exceeds `token_threshold`.
    """
    llm_limiter = limiter or default_llm_limiter
    if token_threshold is None:
        token_threshold = CFG.LLM_HISTORY_SUMMARIZATION_TOKEN_THRESHOLD

    async def process_history(messages: "list[ModelMessage]") -> "list[ModelMessage]":
        current_tokens = llm_limiter.count_tokens(messages)
        if current_tokens <= token_threshold:
            return messages
        zrb_print(
            stylize_yellow(
                (
                    f"\n  Token Threshold exceeded ({current_tokens}/{token_threshold}). "
                    "Compressing conversation..."
                )
            ),
            plain=True,
        )
        result = await summarize_history(messages, agent, summary_window)
        if result != messages:
            new_tokens = llm_limiter.count_tokens(result)
            zrb_print(
                stylize_yellow(
                    f"  Conversation compressed ({new_tokens}/{token_threshold})"
                ),
                plain=True,
            )
        else:
            zrb_print(stylize_error("  Cannot compress conversation..."), plain=True)
        return result

    return process_history
