"""
Poping Agent Builder and Session Management

[File Overview]
===============
- Purpose: Provides OpenAI-style agent builder pattern and session management
- Data Flow:
    Builder → create agent on backend → Agent instance → Session context → chat/stream
- Core Data Structures:
    - AgentConfig: Configuration dataclass (llm, tools, memory, knowledge, datasets)
    - AgentBuilder: Fluent builder for configuring and creating agents
    - Agent: Agent instance with agent_id
    - Session: Chat session context manager with local tool execution
- Main Functions:
    1. AgentBuilder.with_tools() → register local/cloud tools
    2. AgentBuilder.with_knowledge() → attach knowledge bases
    3. AgentBuilder.build() → create agent on backend
    4. Agent.session() → create session context
    5. Session.chat() → send message and execute local tools
- Related Files:
    @/poping_sdk/poping/client.py → PopingClient HTTP methods
    @/poping_sdk/poping/tool.py → @tool decorator and metadata extraction
    @/poping_sdk/poping/context.py → Context operations
    @/poping_sdk/poping/storage.py → Storage operations
"""

from __future__ import annotations

import json
from dataclasses import dataclass, field
from typing import Any, Callable, Dict, List, Optional, Type, TypeVar, Union

from .client import PopingClient
from .context import Context
from .storage import Storage
from .tool import get_tool_definition

# Type variable for structured output
T = TypeVar("T")


# ==============================================================================
# StructuredResponse - Response wrapper for structured output
# ==============================================================================

class StructuredResponse:
    """
    Response wrapper that holds both text and parsed structured output.

    When using with_structured_output(), the response from session.chat()
    will be a StructuredResponse object with a `parsed` attribute containing
    the Pydantic model instance.

    Attributes:
        text: The raw text response from the LLM
        parsed: The parsed Pydantic model instance (if response_model was set)
        raw: The raw response dict from the backend
    """

    def __init__(
        self,
        text: str = "",
        parsed: Any = None,
        raw: Optional[Dict[str, Any]] = None,
    ):
        self.text = text
        self.parsed = parsed
        self.raw = raw

    def __str__(self) -> str:
        return self.text

    def __repr__(self) -> str:
        if self.parsed:
            return f"StructuredResponse(text={self.text!r}, parsed={self.parsed!r})"
        return f"StructuredResponse(text={self.text!r})"


# ==============================================================================
# AgentConfig - Lightweight Configuration Dataclass
# ==============================================================================

@dataclass
class AgentConfig:
    """
    Agent configuration (lightweight version)

    Stores agent behavior settings without over-engineering.
    Maps to backend agent creation payload via _build_backend_payload().
    """

    # Core LLM
    llm: str

    # Tools
    local_tools: List[Dict[str, Any]] = field(default_factory=list)  # Tool metadata (not callables)
    cloud_tools: List[str] = field(default_factory=list)  # Cloud tool IDs

    # Knowledge bases
    knowledge_bases: List[str] = field(default_factory=list)
    knowledge_toolset: bool = True  # Auto-add knowledge tools

    # Memory
    memory_session: bool = False  # Enable session memory
    memory_toolset: bool = True   # Auto-add memory tools

    # Datasets
    datasets: List[str] = field(default_factory=list)
    data_toolset: bool = True  # Auto-add dataset tools

    # LLM parameters
    system_prompt: Optional[str] = None
    max_tokens: int = 4096
    temperature: float = 0.7

    # Project scope
    project_id: Optional[str] = None


# ==============================================================================
# AgentBuilder - Fluent Builder Pattern
# ==============================================================================

class AgentBuilder:
    """
    Agent builder with fluent API

    [Class: AgentBuilder]
    ====================
    - Input: client (PopingClient), llm or agent_id
    - Output: Agent instance after build()
    - Role in Flow: Configuration phase before agent creation
    - Logic:
        1. Accumulate configuration via with_*() methods
        2. Store local tool callables separately (not in config)
        3. build() creates agent on backend
        4. Return Agent instance with agent_id

    Usage:
        agent = (
            poping.agent(llm="claude-3-5-sonnet-20241022")
            .with_tools(local=[my_tool], cloud=["oplus.text_to_image"])
            .with_knowledge(["kb_001"], toolset=True)
            .with_memory(toolset=True, session=False)
            .with_llm_config(system_prompt="You are helpful", temperature=0.7)
            .build(name="assistant")
        )
    """

    def __init__(
        self,
        client: PopingClient,
        llm: Optional[str] = None,
        agent_id: Optional[str] = None,
        project_id: Optional[str] = None,
    ):
        """
        Initialize builder

        Args:
            client: PopingClient instance
            llm: LLM model name (for creating new agent)
            agent_id: Agent ID (for loading existing agent)
            project_id: Project ID scope
        """
        self.client = client
        self.agent_id = agent_id
        self.project_id = project_id

        # Configuration
        if llm:
            self.config = AgentConfig(llm=llm, project_id=project_id)
        else:
            self.config = None

        # Local tool callables (name -> function)
        # Stored separately because they can't be serialized to backend
        self.local_tool_funcs: Dict[str, Callable] = {}

        # Subagents list (for multi-agent orchestration)
        self._subagents: List[Dict[str, Any]] = []

        # Subagent Agent objects (for local execution of subagents with local tools)
        # agent_id -> Agent object
        self._subagent_objects: Dict[str, "Agent"] = {}

        # Structured output response model (Pydantic BaseModel)
        self._response_model: Optional[Type] = None

    def with_tools(
        self,
        local: Optional[List[Any]] = None,
        cloud: Optional[List[str]] = None
    ) -> AgentBuilder:
        """
        Register tools

        Args:
            local: List of @tool decorated functions
            cloud: List of cloud tool IDs (e.g., ["oplus.text_to_image"])

        Returns:
            Self for chaining

        Example:
            @tool()
            def calculator(a: float, b: float) -> float:
                return a + b

            builder.with_tools(
                local=[calculator],
                cloud=["oplus.text_to_image"]
            )
        """
        if local:
            for fn in local:
                # Extract metadata (for backend)
                definition = get_tool_definition(fn)
                if definition:
                    self.config.local_tools.append(definition)

                # Store callable (for local execution)
                if callable(fn) and hasattr(fn, "_poping_tool"):
                    data = getattr(fn, "_poping_tool")
                    self.local_tool_funcs[data["name"]] = data["callable"]

        if cloud:
            self.config.cloud_tools.extend(cloud)

        return self

    def with_knowledge(
        self,
        kb_ids: List[str],
        toolset: bool = True
    ) -> AgentBuilder:
        """
        Attach knowledge bases

        Args:
            kb_ids: List of knowledge base IDs
            toolset: Auto-register knowledge tools (search, query)

        Returns:
            Self for chaining
        """
        self.config.knowledge_bases = kb_ids
        self.config.knowledge_toolset = toolset
        return self

    def with_memory(
        self,
        toolset: bool = True,
        session: bool = False
    ) -> AgentBuilder:
        """
        Configure memory

        Args:
            toolset: Auto-register memory tools (add, search)
            session: Enable session-level memory

        Returns:
            Self for chaining
        """
        self.config.memory_toolset = toolset
        self.config.memory_session = session
        return self

    def with_data(
        self,
        dataset_ids: List[str],
        toolset: bool = True
    ) -> AgentBuilder:
        """
        Attach datasets

        Args:
            dataset_ids: List of dataset IDs
            toolset: Auto-register dataset tools (query, search)

        Returns:
            Self for chaining
        """
        self.config.datasets = dataset_ids
        self.config.data_toolset = toolset
        return self

    def with_subagent(
        self,
        subagent: Union[str, "Agent"],
        name: Optional[str] = None,
        description: Optional[str] = None,
        streaming: bool = False,
    ) -> AgentBuilder:
        """
        Register a subagent that this agent can call

        When registered, this agent can call the subagent using the call_subagent tool.
        The subagent will be invoked server-side with the specified task.

        Args:
            subagent: Agent instance or agent_id string
            name: Optional name for the subagent (defaults to agent_id)
            description: Optional description of what the subagent does
            streaming: If True, subagent streams events like main agent.
                      If False (default), subagent returns final result only.
                      Note: streaming=True only works when main agent uses streaming.

        Returns:
            Self for chaining

        Example:
            # Create a helper agent
            helper = poping.agent(llm="claude-haiku-4-5-20251001").build(name="helper")

            # Create main agent with non-streaming subagent (default)
            main_agent = (
                poping.agent(llm="claude-sonnet-4-20250514")
                .with_subagent(helper, description="A helper for simple tasks")
                .build(name="orchestrator")
            )

            # Create main agent with streaming subagent
            main_agent_streaming = (
                poping.agent(llm="claude-sonnet-4-20250514")
                .with_subagent(helper, description="A helper", streaming=True)
                .build(name="orchestrator_streaming")
            )
        """
        # Get agent_id from Agent instance or use string directly
        if hasattr(subagent, 'agent_id'):
            agent_id = subagent.agent_id
            agent_name = getattr(subagent, 'name', agent_id)
            # Store the Agent object for local execution of subagents with local tools
            self._subagent_objects[agent_id] = subagent
        else:
            agent_id = str(subagent)
            agent_name = name or agent_id

        self._subagents.append({
            'agent_id': agent_id,
            'name': name or agent_name,
            'description': description or f"Subagent: {agent_name}",
            'streaming': streaming,
        })

        return self

    def with_llm_config(
        self,
        system_prompt: Optional[str] = None,
        max_tokens: Optional[int] = None,
        temperature: Optional[float] = None,
    ) -> AgentBuilder:
        """
        Configure LLM parameters

        Args:
            system_prompt: System prompt override
            max_tokens: Maximum tokens per response
            temperature: Sampling temperature (0.0-1.0)

        Returns:
            Self for chaining
        """
        if system_prompt is not None:
            self.config.system_prompt = system_prompt
        if max_tokens is not None:
            self.config.max_tokens = max_tokens
        if temperature is not None:
            self.config.temperature = temperature
        return self

    def with_structured_output(self, model: Type[T]) -> AgentBuilder:
        """
        Configure structured output using a Pydantic model.

        When set, the agent's chat response will include a `parsed` field
        containing an instance of the specified Pydantic model.

        Args:
            model: Pydantic BaseModel class defining the response schema

        Returns:
            Self for chaining

        Example:
            from pydantic import BaseModel

            class PersonInfo(BaseModel):
                name: str
                age: int
                occupation: str

            agent = (
                poping.agent(llm="claude-sonnet-4-5-20250929")
                .with_structured_output(PersonInfo)
                .build(name="extractor")
            )

            with agent.session() as session:
                response = session.chat("Extract: John is 30, works as engineer")
                # response.parsed is a PersonInfo instance
                print(response.parsed.name)  # "John"

        Note:
            - Requires provider support (Anthropic beta.messages.parse, OpenAI)
            - Streaming with structured output is not supported
        """
        self._response_model = model
        return self

    def build(self, name: Optional[str] = None) -> Agent:
        """
        Create or reuse agent on backend

        [Function: build]
        =================
        - Input: name (str, optional)
        - Output: Agent instance
        - Role in Flow: Final step - creates/reuses agent on backend
        - Logic:
            1. If agent_id exists, load existing agent from backend
            2. Check if agent with same name exists in project
            3. If exists, compare config and update if different
            4. If not exists, create new agent
            5. Return Agent instance with agent_id

        Args:
            name: Agent name (required for new agents)

        Returns:
            Agent instance

        Raises:
            ValueError: If name not provided for new agent
        """
        # Load existing agent by ID
        if self.agent_id:
            data = self.client.get_agent(self.agent_id, project_id=self.project_id)
            return Agent(
                client=self.client,
                agent_id=data["agent_id"],
                name=data["name"],
                project_id=self.project_id,
                local_tool_funcs=self.local_tool_funcs,
                response_model=self._response_model,
                subagent_objects=self._subagent_objects,
            )

        # Name required for new/lookup
        if not name:
            raise ValueError("Agent name required for new agent")

        # Check if agent with same name exists
        existing_agents = self.client.list_agents(project_id=self.project_id, name=name)

        if existing_agents:
            # Agent exists - check if config needs update
            existing = existing_agents[0]
            existing_id = existing["agent_id"]

            # Get full config for comparison
            existing_config = self.client.get_agent(existing_id, project_id=self.project_id)
            new_payload = self._build_backend_payload(name)

            # Compare relevant fields
            if self._configs_differ(existing_config, new_payload):
                # Update existing agent
                self.client.update_agent(existing_id, **new_payload)

            return Agent(
                client=self.client,
                agent_id=existing_id,
                name=name,
                project_id=self.project_id,
                local_tool_funcs=self.local_tool_funcs,
                response_model=self._response_model,
                subagent_objects=self._subagent_objects,
            )

        # Create new agent
        payload = self._build_backend_payload(name)
        data = self.client.create_agent(**payload)

        return Agent(
            client=self.client,
            agent_id=data["agent_id"],
            name=data.get("name", name),
            project_id=self.project_id,
            local_tool_funcs=self.local_tool_funcs,
            response_model=self._response_model,
            subagent_objects=self._subagent_objects,
        )

    def _configs_differ(self, existing: Dict[str, Any], new: Dict[str, Any]) -> bool:
        """
        Compare existing config with new payload

        Checks if key configuration fields differ.

        Args:
            existing: Existing agent config from backend
            new: New payload to be sent

        Returns:
            True if configs differ, False if same
        """
        # Fields to compare
        compare_fields = ["model", "system", "max_tokens", "temperature", "tools"]

        for field in compare_fields:
            existing_val = existing.get(field)
            new_val = new.get(field)

            # Normalize None vs missing
            if existing_val is None and new_val is None:
                continue
            if existing_val != new_val:
                return True

        return False

    def _build_backend_payload(self, name: str) -> Dict[str, Any]:
        """
        Transform SDK config to backend API payload

        Maps AgentConfig fields to backend /api/v1/agents POST schema.
        Backend expects: name, model, provider, tools, knowledge_bases, etc.

        Args:
            name: Agent name

        Returns:
            Payload dict for client.create_agent()
        """
        payload: Dict[str, Any] = {
            "name": name,
            "model": self.config.llm,
            "provider": "anthropic",  # Default provider
        }

        # Optional project scope
        if self.project_id:
            payload["project_id"] = self.project_id

        # LLM config
        if self.config.system_prompt:
            payload["system"] = self.config.system_prompt
        if self.config.max_tokens:
            payload["max_tokens"] = self.config.max_tokens
        if self.config.temperature:
            payload["temperature"] = self.config.temperature

        # Tools - merge local + cloud
        all_tools = []

        # Local tools (metadata only - callables stay in SDK)
        for tool_def in self.config.local_tools:
            all_tools.append({
                "name": tool_def["name"],
                "description": tool_def["description"],
                "input_schema": tool_def["input_schema"],
            })

        # Cloud tools (backend knows these)
        for cloud_id in self.config.cloud_tools:
            all_tools.append({"name": cloud_id, "type": "cloud"})

        # Subagent tool (if subagents registered)
        if self._subagents:
            # Build description listing all available subagents
            subagent_list = "\n".join(
                f"  - {s['agent_id']}: {s['description']}"
                for s in self._subagents
            )
            all_tools.append({
                "name": "call_subagent",
                "description": f"Call a subagent to perform a task. Available subagents:\n{subagent_list}",
                "input_schema": {
                    "type": "object",
                    "properties": {
                        "agent_id": {
                            "type": "string",
                            "description": "The ID of the subagent to call",
                            "enum": [s['agent_id'] for s in self._subagents],
                        },
                        "task": {
                            "type": "string",
                            "description": "The task or message to send to the subagent",
                        },
                    },
                    "required": ["agent_id", "task"],
                },
            })

        if all_tools:
            payload["tools"] = all_tools

        # Knowledge bases
        if self.config.knowledge_bases:
            payload["knowledge_bases"] = self.config.knowledge_bases

        # Datasets
        if self.config.datasets:
            payload["datasets"] = self.config.datasets

        # Subagents metadata (for backend to know streaming settings)
        if self._subagents:
            payload["metadata"] = payload.get("metadata", {})
            payload["metadata"]["subagents"] = {
                s['agent_id']: {
                    'name': s['name'],
                    'description': s['description'],
                    'streaming': s.get('streaming', False),
                }
                for s in self._subagents
            }

        return payload


# ==============================================================================
# Agent - Agent Instance
# ==============================================================================

class Agent:
    """
    Agent instance with agent_id

    [Class: Agent]
    ==============
    - Input: agent_id from backend
    - Output: Session context via session()
    - Role in Flow: Represents a persisted agent, creates sessions for chat
    - Logic:
        1. Hold agent metadata (agent_id, name)
        2. Hold client reference for API calls
        3. Hold local_tool_funcs for session execution
        4. session() creates Session context manager

    Attributes:
        agent_id: Unique agent identifier
        name: Agent name
        client: PopingClient instance
        project_id: Project scope
        local_tool_funcs: Registered local tool callables
    """

    def __init__(
        self,
        client: PopingClient,
        agent_id: str,
        name: str,
        project_id: Optional[str] = None,
        local_tool_funcs: Optional[Dict[str, Callable]] = None,
        response_model: Optional[Type] = None,
        subagent_objects: Optional[Dict[str, "Agent"]] = None,
    ):
        self.client = client
        self.agent_id = agent_id
        self.name = name
        self.project_id = project_id
        self.local_tool_funcs = local_tool_funcs or {}
        self.response_model = response_model
        self.subagent_objects = subagent_objects or {}

    def session(
        self,
        end_user_id: Optional[str] = None,
        session_id: Optional[str] = None,
    ) -> Session:
        """
        Create session context for chatting

        Args:
            end_user_id: End user identifier (for tracking)
            session_id: Existing session ID (optional - will create new if not provided)

        Returns:
            Session context manager

        Example:
            with agent.session(end_user_id="user_123") as conv:
                response = conv.chat("Hello!")
        """
        return Session(
            client=self.client,
            agent_id=self.agent_id,
            project_id=self.project_id,
            end_user_id=end_user_id,
            session_id=session_id,
            local_tool_funcs=self.local_tool_funcs,
            response_model=self.response_model,
            subagent_objects=self.subagent_objects,
        )

    def __repr__(self) -> str:
        return f"Agent(agent_id={self.agent_id!r}, name={self.name!r})"


# ==============================================================================
# Session - Chat Session Context Manager
# ==============================================================================

class Session:
    """
    Chat session with context management and local tool execution

    [Class: Session]
    ================
    - Input: agent_id, session_id (optional), end_user_id
    - Output: Chat responses
    - Role in Flow: Manages conversation state and executes local tools
    - Logic:
        1. __enter__ creates or loads session
        2. chat() sends message to backend
        3. _handle_local_tools() detects tool_use blocks
        4. Execute local tools and append tool_result
        5. Loop until no more tool calls
        6. Return final text response

    Attributes:
        client: PopingClient instance
        agent_id: Agent identifier
        session_id: Session identifier (created on __enter__)
        project_id: Project scope
        end_user_id: End user identifier
        local_tool_funcs: Local tool callables

    Usage:
        with agent.session(end_user_id="user_123") as conv:
            response = conv.chat("What is 5 + 3?")  # Executes local calculator tool
            print(response)
    """

    def __init__(
        self,
        client: PopingClient,
        agent_id: str,
        project_id: Optional[str],
        end_user_id: Optional[str],
        session_id: Optional[str],
        local_tool_funcs: Dict[str, Callable],
        response_model: Optional[Type] = None,
        subagent_objects: Optional[Dict[str, "Agent"]] = None,
    ):
        self.client = client
        self.agent_id = agent_id
        self.project_id = project_id
        self.end_user_id = end_user_id
        self.session_id = session_id
        self.local_tool_funcs = local_tool_funcs
        self.response_model = response_model
        self.subagent_objects = subagent_objects or {}

    def __enter__(self) -> Session:
        """
        Context manager entry - create or load session
        """
        if not self.session_id:
            # Create new session on backend
            data = self.client.create_session(
                project_id=self.project_id,
                end_user_id=self.end_user_id,
            )
            self.session_id = data.get("session_id")

        return self

    def __exit__(self, exc_type, exc_val, exc_tb) -> None:
        """
        Context manager exit - cleanup (optional)
        """
        # No cleanup needed - session persists on backend
        pass

    @property
    def context(self) -> Context:
        """
        Access session context (messages, artifacts)

        Returns:
            Context instance for this session
        """
        return Context(
            client=self.client,
            agent_id=self.agent_id,
            session_id=self.session_id,
            project_id=self.project_id,
            end_user_id=self.end_user_id,
        )

    @property
    def storage(self) -> Storage:
        """
        Access storage for this session

        Returns:
            Storage instance
        """
        return Storage(
            self.client,
            project_id=self.project_id,
            end_user_id=self.end_user_id,
        )

    def chat(
        self,
        message: Union[str, Dict[str, Any], List[Dict[str, Any]]],
        raw: bool = False,
    ) -> Any:
        """
        Send message and execute local tools

        [Function: chat]
        ================
        - Input: message (str or content blocks)
        - Output: Text response, StructuredResponse (if response_model set), or raw response
        - Role in Flow: Main chat interface with automatic tool execution
        - Logic:
            1. Normalize message to string
            2. If response_model is set, generate JSON schema
            3. Send to backend via client.chat()
            4. Check response for tool_use blocks
            5. If local tool → execute and append tool_result
            6. Loop up to 5 times to prevent infinite tool calls
            7. If response_model set, parse response and return StructuredResponse
            8. Return final text response

        Args:
            message: User message (string or content blocks)
            raw: Return raw response (default: False, returns text only)

        Returns:
            - If response_model set: StructuredResponse with .parsed attribute
            - If raw=True: Full response dict
            - Otherwise: Text response string

        Example:
            response = conv.chat("What is 42 + 58?")
            # SDK executes local calculator tool automatically
            print(response)  # "The result is 100"
        """
        message_text = self._normalize_message(message)

        # Generate JSON schema if response_model is set
        response_model_schema = None
        if self.response_model is not None:
            # Get JSON schema from Pydantic model
            response_model_schema = self.response_model.model_json_schema()

        # Tool execution loop (max 5 iterations to prevent infinite loops)
        for _ in range(5):
            resp = self.client.chat(
                self.agent_id,
                message=message_text,
                session_id=self.session_id,
                project_id=self.project_id,
                end_user_id=self.end_user_id,
                response_model_schema=response_model_schema,
            )

            # Update session_id if backend created one
            if not self.session_id and isinstance(resp, dict):
                self.session_id = resp.get("session_id")

            # Check for tool_use blocks
            tool_executed = self._handle_local_tools(resp)

            if not tool_executed:
                # No more tools to execute - return response
                break

            # Tools executed - ask agent to continue with results
            message_text = "continue"

        # Return raw response or extract text
        if raw:
            return resp

        # Handle structured output
        if self.response_model is not None:
            text = self._extract_text(resp)
            parsed = None

            # Check if backend returned parsed output
            if isinstance(resp, dict) and "parsed" in resp and resp["parsed"] is not None:
                try:
                    # Validate and create Pydantic model instance
                    parsed = self.response_model.model_validate(resp["parsed"])
                except Exception:
                    # If validation fails, return raw parsed dict
                    parsed = resp["parsed"]

            return StructuredResponse(text=text, parsed=parsed, raw=resp)

        return self._extract_text(resp)

    def stream(
        self,
        message: Union[str, Dict[str, Any], List[Dict[str, Any]]]
    ):
        """
        Stream chat response with local tool execution

        Args:
            message: User message

        Yields:
            Streaming events (NDJSON)

        Example:
            for event in conv.stream("Tell me a story"):
                if event.get("type") == "content_block_delta":
                    print(event["delta"]["text"], end="")
        """
        message_text = self._normalize_message(message)

        # Local tool execution loop (max 5 iterations to prevent infinite loops)
        for iteration in range(5):
            # Collect events for this iteration
            tool_uses = []
            current_tool_use = None
            # Track tool_use IDs that will be handled locally
            local_tool_use_ids = set()

            for event in self.client.chat_stream_events(
                self.agent_id,
                message=message_text,
                session_id=self.session_id,
                project_id=self.project_id,
                end_user_id=self.end_user_id,
            ):
                # Track tool_use blocks for later execution
                event_type = event.get("type")
                if event_type == "content_block_start":
                    cb = event.get("content_block", {})
                    if cb.get("type") == "tool_use":
                        current_tool_use = {
                            "id": cb.get("id"),
                            "name": cb.get("name"),
                            "input": cb.get("input") or {},
                        }
                elif event_type == "content_block_delta":
                    delta = event.get("delta", {})
                    if delta.get("type") == "input_json_delta" and current_tool_use:
                        # Accumulate partial JSON
                        partial = delta.get("partial_json", "")
                        if "partial_input" not in current_tool_use:
                            current_tool_use["partial_input"] = ""
                        current_tool_use["partial_input"] += partial
                elif event_type == "content_block_stop" and current_tool_use:
                    # Parse accumulated JSON if available
                    if "partial_input" in current_tool_use:
                        try:
                            current_tool_use["input"] = json.loads(current_tool_use["partial_input"])
                        except json.JSONDecodeError:
                            pass
                        del current_tool_use["partial_input"]

                    # Check if this tool will be handled locally
                    tool_name = current_tool_use.get("name")
                    tool_input = current_tool_use.get("input") or {}

                    # Check for call_subagent with local tools
                    if tool_name == "call_subagent" and isinstance(tool_input, dict):
                        subagent_id = tool_input.get("agent_id")
                        if subagent_id and subagent_id in self.subagent_objects:
                            subagent = self.subagent_objects[subagent_id]
                            if subagent.local_tool_funcs:
                                local_tool_use_ids.add(current_tool_use.get("id"))

                    # Check for regular local tools
                    if tool_name in self.local_tool_funcs:
                        local_tool_use_ids.add(current_tool_use.get("id"))

                    tool_uses.append(current_tool_use)
                    current_tool_use = None
                elif event_type == "tool_result":
                    # Filter out backend tool_result events for tools we'll handle locally
                    tool_use_id = event.get("tool_use_id")
                    if tool_use_id in local_tool_use_ids:
                        # Skip this event - we'll emit our own result after local execution
                        continue

                yield event

            # Check if we need to handle local tools
            if not tool_uses:
                return

            # Check if any tools need local execution (already tracked during streaming)
            if not local_tool_use_ids:
                # No local tools to execute - backend handled everything
                return

            # Execute local tools
            for tool in tool_uses:
                tool_name = tool.get("name")
                tool_use_id = tool.get("id")
                tool_input = tool.get("input") or {}

                result = None
                is_handled = False

                # Check if this is a call_subagent - SDK handles ALL subagent execution
                if tool_name == "call_subagent" and isinstance(tool_input, dict):
                    subagent_id = tool_input.get("agent_id")
                    task = tool_input.get("task") or tool_input.get("message") or ""

                    if subagent_id and subagent_id in self.subagent_objects:
                        subagent = self.subagent_objects[subagent_id]

                        # Execute subagent locally (SDK handles all subagent execution)
                        try:
                            with subagent.session(
                                end_user_id=self.end_user_id
                            ) as sub_session:
                                sub_response = sub_session.chat(task)
                                result = sub_response

                                # Emit a synthetic tool_result event
                                yield {
                                    "type": "tool_result",
                                    "tool_use_id": tool_use_id,
                                    "content": result,
                                    "is_error": False,
                                }
                        except Exception as exc:
                            result = f"Subagent execution error: {exc}"
                            yield {
                                "type": "tool_result",
                                "tool_use_id": tool_use_id,
                                "content": result,
                                "is_error": True,
                            }

                        is_handled = True
                    elif subagent_id:
                        # Subagent not found in SDK - return error
                        result = f"Subagent '{subagent_id}' not found. Make sure to register it with with_subagent()."
                        yield {
                            "type": "tool_result",
                            "tool_use_id": tool_use_id,
                            "content": result,
                            "is_error": True,
                        }
                        is_handled = True

                # Check if local tool (not subagent)
                if not is_handled and tool_name in self.local_tool_funcs:
                    func = self.local_tool_funcs[tool_name]

                    try:
                        if isinstance(tool_input, dict):
                            result = func(**tool_input)
                        else:
                            result = func(tool_input)

                        # Emit a synthetic tool_result event
                        yield {
                            "type": "tool_result",
                            "tool_use_id": tool_use_id,
                            "content": result,
                            "is_error": False,
                        }
                    except Exception as exc:
                        result = f"Tool execution error: {exc}"
                        yield {
                            "type": "tool_result",
                            "tool_use_id": tool_use_id,
                            "content": result,
                            "is_error": True,
                        }

                    is_handled = True

                # Append tool_result to session for handled tools
                if is_handled and result is not None:
                    self.client.append_session_message(
                        self.agent_id,
                        self.session_id,
                        role="user",
                        content=json.dumps({
                            "type": "tool_result",
                            "tool_use_id": tool_use_id,
                            "name": tool_name,
                            "result": result,
                        }),
                        project_id=self.project_id,
                        end_user_id=self.end_user_id,
                    )

            # Continue conversation with "continue" message
            message_text = "continue"

    def _normalize_message(
        self,
        message: Union[str, Dict[str, Any], List[Dict[str, Any]]]
    ) -> str:
        """
        Normalize message to string format

        Supports:
        - String: returned as-is
        - Dict: single content block
        - List: multiple content blocks

        Args:
            message: User message in various formats

        Returns:
            Normalized string message
        """
        if isinstance(message, str):
            return message

        # Handle content blocks
        blocks: List[Dict[str, Any]] = []
        if isinstance(message, dict):
            blocks = [message]
        elif isinstance(message, list):
            blocks = message
        else:
            raise TypeError("message must be str, dict, or list")

        # Extract text from blocks
        parts: List[str] = []
        for block in blocks:
            block_type = block.get("type")
            if block_type == "text":
                parts.append(str(block.get("text", "")))
            elif block_type == "image":
                # Placeholder for image blocks
                src = block.get("source", {})
                url = src.get("url") or src.get("data", "")
                parts.append(f"[image:{url}]")
            elif block_type == "tool_result":
                # Placeholder for tool result blocks
                content = block.get("content", "")
                parts.append(f"[tool_result]: {content}")
            else:
                parts.append(str(block))

        return "\n".join(parts)

    def _handle_local_tools(self, response: Any) -> bool:
        """
        Detect and execute local tools

        [Function: _handle_local_tools]
        ==============================
        - Input: Backend chat response
        - Output: True if tools executed, False otherwise
        - Role in Flow: Executes local tools and appends results to session
        - Logic:
            1. Extract content blocks from response
            2. Find tool_use blocks
            3. Check if tool name is in local_tool_funcs
            4. Execute local tool with provided arguments
            5. Append tool_result to session via client.append_session_message()
            6. Return True if any tools executed

        Args:
            response: Backend response dict

        Returns:
            True if local tools were executed, False otherwise
        """
        if not isinstance(response, dict):
            return False

        msg = response.get("message", {})
        if not isinstance(msg, dict):
            return False

        content = msg.get("content", [])
        if not isinstance(content, list):
            return False

        # Find tool_use blocks
        tool_uses = [
            block for block in content
            if isinstance(block, dict) and block.get("type") == "tool_use"
        ]

        if not tool_uses:
            return False

        # Execute local tools
        executed = False
        for block in tool_uses:
            tool_name = block.get("name")
            tool_use_id = block.get("id") or block.get("tool_use_id")
            args = block.get("input") or {}

            result = None
            is_handled = False

            # Check if this is a call_subagent - SDK handles ALL subagent calls
            if tool_name == "call_subagent" and isinstance(args, dict):
                subagent_id = args.get("agent_id")
                task = args.get("task") or args.get("message") or ""

                # Check if we have this subagent object
                if subagent_id and subagent_id in self.subagent_objects:
                    subagent = self.subagent_objects[subagent_id]

                    # Execute subagent locally (SDK handles all subagent execution)
                    try:
                        with subagent.session(
                            end_user_id=self.end_user_id
                        ) as sub_session:
                            sub_response = sub_session.chat(task)
                            result = sub_response
                    except Exception as exc:
                        result = f"Subagent execution error: {exc}"

                    is_handled = True
                elif subagent_id:
                    # Subagent not found in SDK - return error
                    result = f"Subagent '{subagent_id}' not found. Make sure to register it with with_subagent()."
                    is_handled = True

            # Check if local tool (not subagent)
            if not is_handled and tool_name in self.local_tool_funcs:
                # Execute tool
                func = self.local_tool_funcs[tool_name]

                try:
                    if isinstance(args, dict):
                        result = func(**args)
                    else:
                        result = func(args)
                except Exception as exc:
                    result = f"Tool execution error: {exc}"

                is_handled = True

            # If we handled the tool, append result
            if is_handled and result is not None:
                # Append tool_result to session
                self.client.append_session_message(
                    self.agent_id,
                    self.session_id,
                    role="user",
                    content=json.dumps({
                        "type": "tool_result",
                        "tool_use_id": tool_use_id,
                        "name": tool_name,
                        "result": result,
                    }),
                    project_id=self.project_id,
                    end_user_id=self.end_user_id,
                )

                executed = True

        return executed

    def _extract_text(self, response: Any) -> str:
        """
        Extract text from response

        Args:
            response: Backend response

        Returns:
            Text content (or full response if no text found)
        """
        if not isinstance(response, dict):
            return str(response)

        msg = response.get("message", {})
        if not isinstance(msg, dict):
            return str(response)

        content = msg.get("content", [])
        if not isinstance(content, list):
            return str(response)

        # Find first text block
        for block in content:
            if isinstance(block, dict) and block.get("type") == "text":
                return block.get("text", "")

        # No text block found - return full response
        return str(response)


__all__ = ["AgentConfig", "AgentBuilder", "Agent", "Session"]
