"""Message and metadata extraction for Claude Code sessions.

Extracts ChatMessage objects, token usage, metadata, and other
information from Claude Code events.
"""

from logging import getLogger
from typing import Any

from inspect_ai.event import Event, ModelEvent
from inspect_ai.model import (
    Content,
    ContentImage,
    ContentReasoning,
    ContentText,
)
from inspect_ai.model._chat_message import (
    ChatMessageTool,
    ChatMessageUser,
)
from inspect_ai.tool import ToolCall

from .detection import (
    get_timestamp,
)
from .models import (
    AssistantEvent,
    BaseEvent,
    SystemEvent,
    UserEvent,
)
from .toolview import tool_view

logger = getLogger(__name__)


def extract_user_message(event: BaseEvent) -> ChatMessageUser | None:
    """Extract a ChatMessageUser from a user event.

    Args:
        event: Claude Code user event

    Returns:
        ChatMessageUser, or None if extraction fails
    """
    if not isinstance(event, UserEvent):
        return None

    content = event.message.content

    # Handle string content
    if isinstance(content, str):
        # Skip command messages
        if content.startswith("<command-name>"):
            return None
        return ChatMessageUser(content=content)

    # Handle list content (tool results, etc.)
    if isinstance(content, list):
        # Check if this is a tool result message
        for block in content:
            if isinstance(block, dict) and block.get("type") == "tool_result":
                # Tool results become ChatMessageTool, not ChatMessageUser
                return None

        # Extract text content
        text_parts = []
        for block in content:
            if isinstance(block, dict):
                if block.get("type") == "text":
                    text_parts.append(str(block.get("text", "")))
                elif "text" in block:
                    text_parts.append(str(block["text"]))
            elif isinstance(block, str):
                text_parts.append(block)

        if text_parts:
            return ChatMessageUser(content="\n".join(text_parts))

    return None


def _extract_content_blocks(
    items: list[dict[str, Any]],
) -> str | list[Content]:
    """Extract text and image content blocks from a tool result content list.

    Returns a plain string when there are no images (preserving existing
    behaviour), or a mixed list of ContentText/ContentImage when images
    are present.
    """
    blocks: list[Content] = []
    has_images = False

    for item in items:
        if not isinstance(item, dict):
            continue
        if item.get("type") == "text":
            text = str(item.get("text", ""))
            # Merge adjacent text blocks
            if blocks and isinstance(blocks[-1], ContentText):
                blocks[-1] = ContentText(text=blocks[-1].text + "\n" + text)
            else:
                blocks.append(ContentText(text=text))
        elif item.get("type") == "image":
            source = item.get("source", {})
            if isinstance(source, dict) and source.get("type") == "base64":
                media_type = source.get("media_type", "image/png")
                data = source.get("data", "")
                blocks.append(ContentImage(image=f"data:{media_type};base64,{data}"))
                has_images = True
            else:
                source_type = (
                    source.get("type", "unknown")
                    if isinstance(source, dict)
                    else "unknown"
                )
                logger.warning(
                    "Unhandled image source type '%s'; image dropped",
                    source_type,
                )

    if not has_images:
        # Return plain string when there are no images (preserving existing
        # behaviour for callers that expect str).
        return "\n".join(
            block.text for block in blocks if isinstance(block, ContentText)
        )

    return blocks


def extract_tool_result_messages(
    event: BaseEvent,
    tool_functions: dict[str, str] | None = None,
) -> list[ChatMessageTool]:
    """Extract ChatMessageTool objects from a user event with tool results.

    Args:
        event: Claude Code user event that may contain tool results
        tool_functions: Optional mapping of tool_use_id to function name,
            used to populate the ``function`` field on ChatMessageTool.

    Returns:
        List of ChatMessageTool objects
    """
    if not isinstance(event, UserEvent):
        return []

    content = event.message.content

    if not isinstance(content, list):
        return []

    tool_messages = []
    for block in content:
        if isinstance(block, dict) and block.get("type") == "tool_result":
            tool_use_id = str(block.get("tool_use_id", ""))
            result_content = block.get("content", "")
            # Note: is_error could be used to set error state in future

            # Handle content that might be a list or string
            if isinstance(result_content, list):
                tool_content = _extract_content_blocks(result_content)
            elif isinstance(result_content, str):
                tool_content = result_content
            else:
                tool_content = str(result_content)

            function_name = tool_functions.get(tool_use_id) if tool_functions else None
            tool_messages.append(
                ChatMessageTool(
                    content=tool_content,
                    tool_call_id=tool_use_id,
                    function=function_name,
                )
            )

    return tool_messages


def extract_assistant_content(
    message_content: list[dict[str, Any]],
) -> tuple[list[Content], list[ToolCall]]:
    """Extract content and tool calls from assistant message content blocks.

    Args:
        message_content: List of content blocks from assistant message

    Returns:
        Tuple of (content list, tool calls list)
    """
    content: list[Content] = []
    tool_calls: list[ToolCall] = []

    for block in message_content:
        if not isinstance(block, dict):
            continue

        block_type = block.get("type", "")

        if block_type == "text":
            text = block.get("text", "")
            if text:
                content.append(ContentText(text=str(text)))

        elif block_type == "thinking":
            thinking = block.get("thinking", "")
            signature = block.get("signature")
            if thinking:
                content.append(
                    ContentReasoning(
                        reasoning=str(thinking),
                        signature=str(signature) if signature else None,
                    )
                )

        elif block_type == "tool_use":
            tool_id = str(block.get("id", ""))
            tool_name = str(block.get("name", ""))
            tool_input = block.get("input", {})
            if not isinstance(tool_input, dict):
                tool_input = {}
            tool_calls.append(
                ToolCall(
                    id=tool_id,
                    function=tool_name,
                    arguments=tool_input,
                    type="function",
                    view=tool_view(tool_name, tool_input),
                )
            )

    return content, tool_calls


def extract_usage(event: BaseEvent) -> dict[str, int]:
    """Extract token usage from an assistant event.

    Args:
        event: Claude Code assistant event

    Returns:
        Dict with input_tokens, output_tokens, cache_creation_input_tokens,
        cache_read_input_tokens
    """
    if not isinstance(event, AssistantEvent):
        return {}

    usage = event.message.usage
    if not usage:
        return {}

    result: dict[str, int] = {}

    if usage.input_tokens:
        result["input_tokens"] = usage.input_tokens
    if usage.output_tokens:
        result["output_tokens"] = usage.output_tokens
    if usage.cache_creation_input_tokens:
        result["cache_creation_input_tokens"] = usage.cache_creation_input_tokens
    if usage.cache_read_input_tokens:
        result["cache_read_input_tokens"] = usage.cache_read_input_tokens

    return result


def sum_tokens(events: list[BaseEvent]) -> int:
    """Sum total tokens across all events.

    Args:
        events: List of events to sum

    Returns:
        Total token count (input + output)
    """
    total = 0
    for event in events:
        usage = extract_usage(event)
        total += usage.get("input_tokens", 0)
        total += usage.get("output_tokens", 0)
    return total


def extract_model_name(events: list[BaseEvent]) -> str | None:
    """Extract the model name from events.

    Args:
        events: List of events

    Returns:
        Model name, or None if not found
    """
    for event in events:
        if isinstance(event, AssistantEvent):
            model = event.message.model
            if model:
                return model
    return None


def extract_session_metadata(events: list[BaseEvent]) -> dict[str, Any]:
    """Extract session metadata from events.

    Args:
        events: List of events

    Returns:
        Dict with cwd, gitBranch, version, slug, etc.
    """
    metadata: dict[str, Any] = {}

    # Get from first event with these fields
    for event in events:
        if event.cwd and "cwd" not in metadata:
            metadata["cwd"] = event.cwd
        if event.gitBranch and "gitBranch" not in metadata:
            metadata["gitBranch"] = event.gitBranch
        if event.version and "version" not in metadata:
            metadata["version"] = event.version
        if event.slug and "slug" not in metadata:
            metadata["slug"] = event.slug

        # Stop once we have all common fields
        if all(k in metadata for k in ["cwd", "gitBranch", "version", "slug"]):
            break

    return metadata


def extract_compaction_info(event: BaseEvent) -> dict[str, Any] | None:
    """Extract compaction information from a system event.

    Args:
        event: System event with subtype=compact_boundary

    Returns:
        Dict with trigger, preTokens, or None if not a compaction event
    """
    if not isinstance(event, SystemEvent):
        return None

    if event.subtype != "compact_boundary":
        return None

    result: dict[str, Any] = {
        "trigger": "auto",
        "content": event.content or "Conversation compacted",
    }

    if event.compactMetadata:
        result["trigger"] = event.compactMetadata.trigger
        result["preTokens"] = event.compactMetadata.preTokens

    return result


def get_first_timestamp(events: list[BaseEvent]) -> str | None:
    """Get the earliest timestamp from events.

    Args:
        events: List of events

    Returns:
        ISO timestamp string, or None if no timestamps found
    """
    timestamps = []
    for event in events:
        ts = get_timestamp(event)
        if ts:
            timestamps.append(ts)

    if not timestamps:
        return None

    # Sort and return earliest
    timestamps.sort()
    return timestamps[0]


def sum_scout_tokens(events: list[Event]) -> int:
    """Sum total tokens from converted Scout ModelEvent objects.

    Unlike sum_tokens() which only counts main-session BaseEvent tokens,
    this counts tokens from all ModelEvents including loaded subagent events.

    Args:
        events: List of Inspect AI events

    Returns:
        Total token count across all ModelEvents
    """
    total = 0
    for event in events:
        if isinstance(event, ModelEvent) and event.output and event.output.usage:
            total += event.output.usage.total_tokens
    return total
