"""Factory functions for creating subagent tools.

This module provides:
- create_subagent_tool: Create BaseTool from a call function
- create_subagent_call_func: Create a BaseTool.call compatible function from a pydantic-ai Agent
"""

from __future__ import annotations

from collections.abc import Awaitable, Callable, Container
from contextlib import AbstractAsyncContextManager, nullcontext
from inspect import isawaitable
from typing import Annotated, Any, cast
from uuid import uuid4

from pydantic import Field
from pydantic_ai import Agent, AgentRunResult, RunContext, UsageLimits
from pydantic_ai.models import Model

from ya_agent_sdk.context import AgentContext, ModelConfig
from ya_agent_sdk.events import SubagentCompleteEvent, SubagentStartEvent
from ya_agent_sdk.toolsets.core.base import BaseTool
from ya_agent_sdk.usage import InternalUsage

# Type alias for instruction functions
InstructionFunc = Callable[[RunContext[AgentContext]], str | None]

# Type alias for availability check functions
AvailabilityCheckFunc = Callable[[RunContext[AgentContext]], bool]

# Type alias for BaseTool.call compatible function
SubagentCallFunc = Callable[..., Awaitable[str]]


def create_subagent_tool(
    name: str,
    description: str,
    call_func: SubagentCallFunc,
    *,
    instruction: str | InstructionFunc | None = None,
    availability_check: AvailabilityCheckFunc | None = None,
) -> type[BaseTool]:
    """Create a BaseTool subclass that wraps a subagent call function.

    This factory function creates a tool class that uses the provided call_func
    directly as the tool's call method. The call_func should have a signature
    compatible with BaseTool.call: (ctx: RunContext[AgentContext], **kwargs) -> str

    Use create_subagent_call_func() to create a compatible call_func from a
    pydantic-ai Agent.

    Args:
        name: Tool name used for invocation.
        description: Tool description shown to the model.
        call_func: Async function with signature (ctx: RunContext[AgentContext], **kwargs) -> str.
                   Use create_subagent_call_func() to create this from an Agent.
        instruction: Optional instruction for system prompt. Can be a string or
                     a callable that takes RunContext and returns a string.
        availability_check: Optional callable that returns True if the tool is available.
                            Called dynamically each time is_available() is invoked.

    Returns:
        A BaseTool subclass that can be used with Toolset.

    Example::

        from pydantic_ai import Agent

        # Create an agent
        search_agent: Agent[AgentContext, str] = Agent(...)

        # Create the call function using create_subagent_call_func
        search_call = create_subagent_call_func(search_agent)

        # Create the tool
        SearchTool = create_subagent_tool(
            name="search",
            description="Search the web for information",
            call_func=search_call,
            instruction="Use this tool to search for current information.",
        )
    """

    class DynamicSubagentTool(BaseTool):
        """Dynamically created subagent tool."""

        # These will be set by the closure
        name = ""  # Placeholder, will be overwritten
        description = ""  # Placeholder, will be overwritten

        def is_available(self, ctx: RunContext[AgentContext]) -> bool:
            if availability_check is None:
                return True
            return availability_check(ctx)

        async def get_instruction(self, ctx: RunContext[AgentContext]) -> str | None:
            if instruction is None:
                return None
            if callable(instruction):
                return instruction(ctx)
            return instruction

        async def call(self, ctx: RunContext[AgentContext], /, **kwargs: Any) -> str:
            # Placeholder - will be replaced by actual call_func
            raise NotImplementedError  # pragma: no cover

    # Set class attributes from closure variables
    DynamicSubagentTool.name = name
    DynamicSubagentTool.description = description

    # Use call_func directly as the call method
    # call_func should already have the correct signature from create_subagent_call_func
    DynamicSubagentTool.call = call_func  # type: ignore[method-assign]

    # Set a meaningful class name for debugging
    DynamicSubagentTool.__name__ = f"{_to_pascal_case(name)}Tool"
    DynamicSubagentTool.__qualname__ = DynamicSubagentTool.__name__

    return DynamicSubagentTool


def _to_pascal_case(name: str) -> str:
    """Convert snake_case or kebab-case to PascalCase."""
    parts = name.replace("-", "_").split("_")
    return "".join(part.capitalize() for part in parts)


async def _run_subagent_iter(
    agent: Agent[AgentContext, Any],
    sub_ctx: AgentContext,
    prompt: str,
    message_history: list[Any] | None,
) -> AgentRunResult:
    """Run subagent iteration and stream events to subagent's queue.

    Events are emitted to sub_ctx (subagent context) so they go to the
    subagent's queue keyed by agent_id, not the parent's queue.

    Args:
        agent: The subagent to run.
        sub_ctx: Subagent's context (events emitted here).
        prompt: The prompt to send to the subagent.
        message_history: Optional conversation history for resume.

    Returns:
        AgentRunResult from the subagent execution.
    """
    async with agent.iter(
        prompt,
        deps=sub_ctx,
        usage_limits=UsageLimits(request_limit=1000),
        message_history=message_history,
    ) as run:
        async for node in run:
            if Agent.is_user_prompt_node(node) or Agent.is_end_node(node):
                continue

            if Agent.is_model_request_node(node) or Agent.is_call_tools_node(node):
                async with node.stream(run.ctx) as request_stream:
                    async for event in request_stream:
                        await sub_ctx.emit_event(event)

    return cast(AgentRunResult, run.result)


def generate_unique_id(existing: Container[str], *, max_retries: int = 10) -> str:
    """Generate a unique 4-character ID with collision detection.

    First tries using the last 4 characters of run_id. If that collides
    with existing IDs, generates random UUIDs until a unique one is found.

    Args:
        run_id: The current run ID to derive initial ID from.
        existing: Container of existing IDs to check against.
        max_retries: Maximum number of UUID generation attempts (default 10).

    Returns:
        A unique 4-character ID string.

    Raises:
        RuntimeError: If unable to generate unique ID within max_retries.
    """
    for _ in range(max_retries):
        agent_id = uuid4().hex[:4]
        if agent_id not in existing:
            return agent_id

    raise RuntimeError(f"Failed to generate unique agent_id after {max_retries} retries")


def create_subagent_call_func(
    agent: Agent[AgentContext, Any],
    *,
    model_cfg: ModelConfig | None = None,
) -> SubagentCallFunc:
    """Create a BaseTool.call compatible function from a pydantic-ai Agent.

    This function creates a call method that:
    - Has the correct signature for BaseTool.call: (ctx: RunContext[AgentContext], **kwargs) -> str
    - Generates stable agent_id in format {agent.name}-{short_id}
    - Registers the agent in parent's agent_registry
    - Manages subagent_history for conversation continuity
    - Records usage in extra_usages
    - Streams events to parent context

    Args:
        agent: A pydantic-ai Agent with AgentContext as deps type.
        model_cfg: Optional ModelConfig to override in subagent context.
                   If None, subagent inherits parent's model_cfg.

    Returns:
        A function compatible with BaseTool.call signature.

    Example::

        from pydantic_ai import Agent

        search_agent: Agent[AgentContext, str] = Agent(...)

        # Create the call function
        search_call = create_subagent_call_func(search_agent)

        # Pass to create_subagent_tool
        SearchTool = create_subagent_tool("search", "Search the web", search_call)
    """
    agent_name = agent.name or "subagent"

    async def call_func(
        self: BaseTool,
        ctx: RunContext[AgentContext],
        prompt: Annotated[str, Field(description="The prompt to send to the subagent")],
        agent_id: Annotated[str | None, Field(description="Optional agent ID to resume")] = None,
    ) -> str:
        """Execute the agent with the given prompt."""
        deps = ctx.deps

        # Generate stable agent_id if not provided
        if not agent_id:
            short_id = generate_unique_id(deps.subagent_history)
            agent_id = f"{agent_name}-{short_id}"

        # Track whether this is a new agent (not a resume) for cleanup on failure
        is_new_agent = agent_id not in deps.agent_registry

        # Create subagent context (handles registration in agent_registry)
        override_kwargs: dict[str, Any] = {}
        if model_cfg is not None:
            override_kwargs["model_cfg"] = model_cfg

        error_msg = ""
        success = True
        result_output = ""
        request_count = 0

        async with deps.create_subagent_context(agent_name, agent_id=agent_id, **override_kwargs) as sub_ctx:
            # Set the subagent's initial prompt for compact
            sub_ctx.user_prompts = prompt

            # Emit start event to subagent's queue (inside context so sub_ctx.start_at is set)
            prompt_preview = prompt[:100] + "..." if len(prompt) > 100 else prompt
            await sub_ctx.emit_event(
                SubagentStartEvent(
                    event_id=agent_id,  # Use agent_id as event_id to correlate Start/Complete
                    agent_id=agent_id,
                    agent_name=agent_name,
                    prompt_preview=prompt_preview,
                )
            )

            # Wrap subagent execution with observability wrapper
            subagent_ctx: AbstractAsyncContextManager[None] = nullcontext()
            if deps.subagent_wrapper is not None:
                wrapper_metadata = deps.get_wrapper_metadata()
                subagent_ctx = deps.subagent_wrapper(agent_name, agent_id, wrapper_metadata)

            # Save original model before entering wrapper so cleanup always works,
            # even if subagent_wrapper.__aenter__ raises.
            original_model = agent.model

            try:
                async with subagent_ctx:
                    # Apply model wrapper if configured
                    if deps.model_wrapper is not None:
                        wrapper_metadata = deps.get_wrapper_metadata()
                        wrapped = deps.model_wrapper(cast(Model, original_model), agent_name, wrapper_metadata)
                        agent.model = await wrapped if isawaitable(wrapped) else wrapped

                    result = await _run_subagent_iter(agent, sub_ctx, prompt, deps.subagent_history.get(agent_id))
                    result_output = result.output
                    request_count = result.usage().requests

                    # Store message history for future resume
                    deps.subagent_history[agent_id] = result.all_messages()

                    # Record usage in extra_usages
                    if ctx.tool_call_id:
                        model_id = cast(Model, agent.model).model_name
                        deps.add_extra_usage(
                            agent=agent_id,
                            internal_usage=InternalUsage(model_id=model_id, usage=result.usage()),
                            uuid=ctx.tool_call_id,
                        )

            except Exception as e:
                success = False
                error_msg = str(e)
                # Clean up agent_registry for new agents that failed before
                # producing any history, to avoid ghost entries that show up
                # in known-subagents and SubagentInfoTool with no history.
                if is_new_agent and agent_id not in deps.subagent_history:
                    deps.agent_registry.pop(agent_id, None)
                raise

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

                # Emit complete event to subagent's queue (use sub_ctx.elapsed_time for duration)
                elapsed = sub_ctx.elapsed_time
                duration = elapsed.total_seconds() if elapsed else 0.0
                result_preview = result_output[:500] + "..." if len(result_output) > 500 else result_output
                await sub_ctx.emit_event(
                    SubagentCompleteEvent(
                        event_id=agent_id,  # Same event_id as Start for correlation
                        agent_id=agent_id,
                        agent_name=agent_name,
                        success=success,
                        request_count=request_count,
                        result_preview=result_preview,
                        error=error_msg,
                        duration_seconds=duration,
                    )
                )

        # Return formatted result
        return f"""<id>{agent_id}</id>
<response>{result.output}</response>
"""

    return call_func  # type: ignore[return-value]
