"""
Sentrial wrapper for the Claude Agent SDK (Claude Code).

Wraps the ``query()`` async generator to automatically track sessions,
tool calls, tokens, and costs to Sentrial.

Example::

    from claude_agent_sdk import query
    from sentrial import AsyncSentrialClient
    from sentrial.claude_code import wrap_claude_agent

    client = AsyncSentrialClient(api_key="sentrial_live_xxx")

    tracked_query = wrap_claude_agent(query, client=client, default_agent="my-agent")

    async for message in tracked_query(prompt="Fix the tests"):
        # Messages pass through unchanged
        print(message.type)

Requires:
    - ``claude-agent-sdk`` (peer dependency, not imported at runtime)
    - ``httpx`` (for ``AsyncSentrialClient``)
"""

from __future__ import annotations

import asyncio
import logging
import time
from typing import Any, AsyncIterator, Callable, Dict, List, Optional

_logger = logging.getLogger("sentrial")

# ---------------------------------------------------------------------------
# Type aliases (Claude Agent SDK is a peer dep, not imported at runtime)
# ---------------------------------------------------------------------------

# query() signature: query(prompt, options?) -> AsyncIterator[Message]
QueryFunction = Callable[..., AsyncIterator[Any]]


# ---------------------------------------------------------------------------
# Public API
# ---------------------------------------------------------------------------


def wrap_claude_agent(
    query_fn: QueryFunction,
    *,
    client: Any,
    default_agent: str = "claude-agent",
    user_id: str = "anonymous",
    convo_id: Optional[str] = None,
    extra_metadata: Optional[Dict[str, Any]] = None,
) -> QueryFunction:
    """Wrap the Claude Agent SDK's ``query()`` to auto-track sessions and tool calls.

    Tool calls are tracked by observing the message stream (AssistantMessage
    with ToolUseBlock content followed by UserMessage with tool results).
    This approach works regardless of SDK hook support.

    Args:
        query_fn: The original ``query()`` from ``claude_agent_sdk``.
        client: An :class:`AsyncSentrialClient` instance.
        default_agent: Agent name for session creation (default ``"claude-agent"``).
        user_id: User ID to associate with sessions (default ``"anonymous"``).
        convo_id: Optional conversation ID to group related sessions.
        extra_metadata: Extra metadata merged into every session.

    Returns:
        A wrapped function with the same async-iterator signature as ``query()``.
    """

    def wrapped_query(
        prompt: Any = None,
        options: Optional[Dict[str, Any]] = None,
        **kwargs: Any,
    ) -> AsyncIterator[Any]:
        # Support both positional and keyword args to match SDK flexibility
        if prompt is None:
            prompt = kwargs.get("prompt")
        if options is None:
            options = kwargs.get("options", {}) or {}

        return _tracked_generator(
            query_fn,
            prompt=prompt,
            options=options,
            client=client,
            default_agent=default_agent,
            user_id=user_id,
            convo_id=convo_id,
            extra_metadata=extra_metadata,
        )

    return wrapped_query


# ---------------------------------------------------------------------------
# Internal generator
# ---------------------------------------------------------------------------


async def _tracked_generator(
    query_fn: QueryFunction,
    *,
    prompt: Any,
    options: Any,
    client: Any,
    default_agent: str,
    user_id: str,
    convo_id: Optional[str],
    extra_metadata: Optional[Dict[str, Any]],
) -> AsyncIterator[Any]:
    start_time = time.monotonic()
    session_id: Optional[str] = None

    # Build session name from prompt
    if isinstance(prompt, str):
        session_name = f"{default_agent}: {prompt[:100]}"
    else:
        session_name = f"{default_agent} session"

    # Collect pending tool-call tracking tasks so we can flush before completing
    pending_tool_calls: List[asyncio.Task[Any]] = []

    # Track tool uses observed in AssistantMessage content blocks.
    # When the next UserMessage arrives (tool results), we match and track them.
    pending_tool_uses: List[Dict[str, Any]] = []

    # Pass options through unchanged — no hook injection needed
    generator = query_fn(prompt=prompt, options=options)

    try:
        async for message in generator:
            # Detect message type. The real SDK uses dataclass instances
            # (SystemMessage, ResultMessage, etc.) while tests may use dicts.
            msg_type, msg_subtype = _classify(message)

            # --- Observe init message: create session ---
            if msg_type == "system" and msg_subtype == "init":
                # SystemMessage stores init data in .data dict; dicts have it top-level
                init_data = getattr(message, "data", message) if not isinstance(message, dict) else message
                metadata: Dict[str, Any] = {
                    "model": _get(init_data, "model"),
                    "tools": _get(init_data, "tools"),
                    "cwd": _get(init_data, "cwd"),
                    "mcp_servers": _get(init_data, "mcp_servers"),
                    "sdk_session_id": _get(init_data, "session_id"),
                    **(extra_metadata or {}),
                }

                try:
                    session_id = await client.create_session(
                        name=session_name,
                        agent_name=default_agent,
                        user_id=user_id,
                        convo_id=convo_id,
                        metadata=metadata,
                    )
                except Exception:
                    session_id = None

            # --- Observe assistant message: collect tool uses ---
            if msg_type == "assistant" and session_id:
                tool_uses = _extract_tool_uses(message)
                if tool_uses:
                    pending_tool_uses.extend(tool_uses)

            # --- Observe user message (tool results): track tool calls ---
            if msg_type == "user" and session_id and pending_tool_uses:
                tool_results = _extract_tool_results(message)
                _track_matched_tools(
                    client=client,
                    session_id=session_id,
                    pending_tool_uses=pending_tool_uses,
                    tool_results=tool_results,
                    pending_tasks=pending_tool_calls,
                )
                pending_tool_uses.clear()

            # --- Observe result message: complete session ---
            if msg_type == "result" and session_id:
                is_error = bool(_get(message, "is_error"))
                input_tokens = (_get(message, "usage") or {}).get("input_tokens", 0)
                output_tokens = (_get(message, "usage") or {}).get("output_tokens", 0)

                failure_reason: Optional[str] = None
                if is_error:
                    errors = _get(message, "errors")
                    if errors and len(errors) > 0:
                        failure_reason = "; ".join(str(e) for e in errors)
                    else:
                        failure_reason = msg_subtype

                # Flush pending tool-call tracking before completing
                if pending_tool_calls:
                    await asyncio.gather(*pending_tool_calls, return_exceptions=True)

                elapsed_ms = int((time.monotonic() - start_time) * 1000)
                try:
                    await client.complete_session(
                        session_id=session_id,
                        success=not is_error,
                        failure_reason=failure_reason,
                        estimated_cost=_get(message, "total_cost_usd"),
                        prompt_tokens=input_tokens,
                        completion_tokens=output_tokens,
                        total_tokens=input_tokens + output_tokens,
                        duration_ms=_get(message, "duration_ms") if _get(message, "duration_ms") is not None else elapsed_ms,
                        user_input=prompt if isinstance(prompt, str) else None,
                        assistant_output=_get(message, "result"),
                        custom_metrics={
                            "num_turns": _get(message, "num_turns") if _get(message, "num_turns") is not None else 0,
                            "duration_api_ms": _get(message, "duration_api_ms") if _get(message, "duration_api_ms") is not None else 0,
                        },
                    )
                except Exception:
                    pass  # Never break the agent due to tracking failures

            yield message

    except Exception as error:
        # If the generator throws, complete the session as failed
        if session_id:
            elapsed_ms = int((time.monotonic() - start_time) * 1000)
            try:
                await client.complete_session(
                    session_id=session_id,
                    success=False,
                    failure_reason=str(error),
                    duration_ms=elapsed_ms,
                )
            except Exception:
                pass
        raise


# ---------------------------------------------------------------------------
# Tool use extraction from messages
# ---------------------------------------------------------------------------


def _extract_tool_uses(message: Any) -> List[Dict[str, Any]]:
    """Extract tool use blocks from an AssistantMessage.

    Returns a list of dicts: [{"tool_name": str, "tool_input": dict, "tool_use_id": str}]

    Supports both real SDK dataclass instances (AssistantMessage with
    ToolUseBlock content) and dict-based messages (used in tests).
    """
    results: List[Dict[str, Any]] = []

    if isinstance(message, dict):
        # Dict-based: {"type": "assistant", "message": {"content": [...]}}
        msg_data = message.get("message") or message
        content = msg_data.get("content") if isinstance(msg_data, dict) else None
        if not content or not isinstance(content, list):
            return results
        for block in content:
            if isinstance(block, dict) and block.get("type") == "tool_use":
                results.append({
                    "tool_name": block.get("name", "unknown"),
                    "tool_input": block.get("input", {}),
                    "tool_use_id": block.get("id", ""),
                })
    else:
        # Real SDK: AssistantMessage with .content list of block objects
        content = getattr(message, "content", None)
        if not content or not isinstance(content, list):
            return results
        for block in content:
            cls_name = type(block).__name__
            if cls_name == "ToolUseBlock":
                results.append({
                    "tool_name": getattr(block, "name", "unknown"),
                    "tool_input": getattr(block, "input", {}),
                    "tool_use_id": getattr(block, "id", ""),
                })

    return results


def _extract_tool_results(message: Any) -> Dict[str, Any]:
    """Extract tool results from a UserMessage, keyed by tool_use_id.

    Returns: {tool_use_id: result_content, ...}

    Supports both real SDK dataclass instances and dict-based messages.
    """
    results: Dict[str, Any] = {}

    if isinstance(message, dict):
        # Dict-based: {"type": "user", "message": {"content": [{"type": "tool_result", ...}]}}
        msg_data = message.get("message") or message
        content = msg_data.get("content") if isinstance(msg_data, dict) else None
        if not content or not isinstance(content, list):
            return results
        for block in content:
            if isinstance(block, dict) and block.get("type") == "tool_result":
                tid = block.get("tool_use_id", "")
                results[tid] = block.get("content", "")
    else:
        # Real SDK: UserMessage with .content list
        content = getattr(message, "content", None)
        if not content or not isinstance(content, list):
            return results
        for block in content:
            cls_name = type(block).__name__
            if cls_name == "ToolResultBlock":
                tid = getattr(block, "tool_use_id", "")
                results[tid] = getattr(block, "content", "")
            elif isinstance(block, dict) and block.get("type") == "tool_result":
                tid = block.get("tool_use_id", "")
                results[tid] = block.get("content", "")

    return results


def _track_matched_tools(
    *,
    client: Any,
    session_id: str,
    pending_tool_uses: List[Dict[str, Any]],
    tool_results: Dict[str, Any],
    pending_tasks: List[asyncio.Task[Any]],
) -> None:
    """Match tool uses with their results and fire-and-forget track calls."""
    for tool_use in pending_tool_uses:
        tool_use_id = tool_use["tool_use_id"]
        raw_result = tool_results.get(tool_use_id)

        if isinstance(raw_result, dict):
            tool_output = raw_result
        elif raw_result is not None:
            tool_output = {"response": str(raw_result)[:2000]}
        else:
            tool_output = {}

        try:
            task = asyncio.create_task(
                _safe_track_tool(
                    client,
                    session_id=session_id,
                    tool_name=tool_use["tool_name"],
                    tool_input=tool_use["tool_input"],
                    tool_output=tool_output,
                    metadata={"tool_use_id": tool_use_id},
                )
            )
            pending_tasks.append(task)
        except Exception:
            pass


# ---------------------------------------------------------------------------
# Helpers
# ---------------------------------------------------------------------------


def _classify(message: Any) -> tuple[Optional[str], Optional[str]]:
    """Determine (type, subtype) of a message.

    Supports both dict-based messages (used in tests) and real SDK dataclass
    instances (SystemMessage, ResultMessage, AssistantMessage, UserMessage).
    """
    if isinstance(message, dict):
        return message.get("type"), message.get("subtype")

    # Real SDK: class name determines the type
    cls_name = type(message).__name__
    subtype = getattr(message, "subtype", None)

    if cls_name == "SystemMessage":
        return "system", subtype
    if cls_name == "ResultMessage":
        return "result", subtype
    if cls_name == "AssistantMessage":
        return "assistant", None
    if cls_name == "UserMessage":
        return "user", None

    # Fallback: try .type attribute
    return getattr(message, "type", None), subtype


def _get(obj: Any, key: str) -> Any:
    """Get attribute or dict key from a message object."""
    if isinstance(obj, dict):
        return obj.get(key)
    return getattr(obj, key, None)


async def _safe_track_tool(
    client: Any,
    *,
    session_id: str,
    tool_name: str,
    tool_input: Any,
    tool_output: Any,
    tool_error: Optional[Dict[str, Any]] = None,
    metadata: Optional[Dict[str, Any]] = None,
) -> None:
    """Track a tool call, swallowing any errors."""
    try:
        await client.track_tool_call(
            session_id=session_id,
            tool_name=tool_name,
            tool_input=tool_input or {},
            tool_output=tool_output or {},
            tool_error=tool_error,
            metadata=metadata,
        )
    except Exception:
        pass
