from typing import Dict, List, Optional, Any, TypedDict, Union
from pathlib import Path
from datetime import datetime
from dataclasses import dataclass
import base64
import mimetypes
from pydantic import BaseModel, Field, ConfigDict, model_validator
from .basic import CompletionUsage, ToolCall
from .outputs import OutputMode


@dataclass
class SourceDocument:
    """Enhanced source document information similar to Parrot's format."""
    source: str
    filename: str
    file_path: Optional[str] = None
    source_path: Optional[str] = None
    url: Optional[str] = None
    content_type: Optional[str] = None
    category: Optional[str] = None
    source_type: Optional[str] = None
    source_ext: Optional[str] = None
    page_number: Optional[int] = None
    chunk_id: Optional[str] = None
    parent_document_id: Optional[str] = None
    chunk_index: Optional[int] = None
    score: Optional[float] = None
    metadata: Dict[str, Any] = None

    def to_dict(self) -> Dict[str, Any]:
        """Convert to dictionary format."""
        return {
            "source": self.source,
            "filename": self.filename,
            "file_path": self.file_path,
            "source_path": self.source_path,
            "url": self.url,
            "content_type": self.content_type,
            "category": self.category,
            "source_type": self.source_type,
            "source_ext": self.source_ext,
            "page_number": self.page_number,
            "chunk_id": self.chunk_id,
            "parent_document_id": self.parent_document_id,
            "chunk_index": self.chunk_index,
            "score": self.score,
            "metadata": self.metadata or {}
        }


class StreamChunk(BaseModel):
    """Represents a chunk in a streaming response."""
    content: str
    is_complete: bool = False
    chunk_id: Optional[str] = None
    turn_id: Optional[str] = None


class MessageResponse(TypedDict):
    """Response structure for LLM messages."""
    id: str
    type: str
    role: str
    content: List[Dict[str, Any]]
    model: str
    stop_reason: Optional[str]
    stop_sequence: Optional[str]
    usage: Dict[str, int]


class AIMessage(BaseModel):
    """Unified AI message response that can handle various output types."""

    # Core response data
    input: str = Field(
        description="The original input/prompt sent to the LLM"
    )
    output: Any = Field(
        description="The response output - can be text, structured data, dataframe, etc."
    )
    response: Optional[str] = Field(
        default=None,
        description="The textual response from the model, if applicable"
    )
    data: Optional[Any] = Field(
        default=None,
        description="Structured data associated with the response, if applicable"
    )
    code: Optional[str] = Field(
        default=None,
        description="The Python code used for analysis OR the Code generated under request (e.g. JSON definition for a Altair/Vega Chart)."
    )
    images: Optional[List[Path]] = Field(
        default_factory=list,
        description="List of image file paths generated by the model"
    )
    media: Optional[List[Path]] = Field(
        default_factory=list,
        description="List of media files generated by the model"
    )
    files: Optional[List[Path]] = Field(
        default_factory=list,
        description="List of file paths associated with the response"
    )
    documents: Optional[List[Any]] = Field(
        default_factory=list,
        description="List of document file paths associated with the response"
    )
    # Metadata
    model: str = Field(
        description="The model used for generation"
    )
    provider: str = Field(
        description="The LLM provider (openai, groq, claude, etc.)"
    )
    # Usage and performance
    usage: CompletionUsage = Field(
        description="Token usage and timing information"
    )
    # Additional response metadata
    stop_reason: Optional[str] = Field(
        default=None, description="Why the generation stopped"
    )
    finish_reason: Optional[str] = Field(
        default=None, description="Finish reason from provider"
    )
    # Tool usage
    tool_calls: List[ToolCall] = Field(
        default_factory=list,
        description="Tools called during generation"
    )
    # Conversation context
    user_id: Optional[Union[str, int]] = Field(
        default=None, description="User ID for conversation tracking"
    )
    # Timestamps
    created_at: datetime = Field(
        default_factory=datetime.now, description="When the response was created"
    )
    response_time: Optional[float] = Field(
        default=None,
        description="Time taken to generate the response (in seconds)"
    )
    # Raw response for debugging
    raw_response: Optional[Dict[str, Any]] = Field(
        default=None, description="Original response from provider"
    )
    # Conversation turn info
    session_id: Optional[str] = Field(
        default=None, description="Session ID for conversation tracking"
    )
    turn_id: Optional[str] = Field(
        default=None,
        description="Unique ID for this conversation turn"
    )
    # Vector store and conversation history tracking
    used_vector_context: bool = Field(
        default=False,
        description="Whether vector store context was used in generating this response"
    )
    used_conversation_history: bool = Field(
        default=False,
        description="Whether conversation history was used in generating this response"
    )
    vector_context_length: int = Field(
        default=0,
        description="Length of vector context used (in characters)"
    )
    conversation_context_length: int = Field(
        default=0,
        description="Length of conversation context used (in characters)"
    )
    search_results_count: int = Field(
        default=0,
        description="Number of search results retrieved from vector store"
    )
    search_type: Optional[str] = Field(
        default=None,
        description="Type of search performed (similarity, mmr, ensemble, etc.)"
    )
    search_score_threshold: Optional[float] = Field(
        default=None,
        description="Score threshold used for vector search"
    )
    context_sources: Optional[List[str]] = Field(
        default_factory=list,
        description="List of source identifiers for context used"
    )
    source_documents: Optional[List[SourceDocument]] = Field(
        default_factory=list,
        description="List of detailed source documents used for context"
    )
    structured_output: Optional[Union[BaseModel, Any]] = Field(
        default=None,
        description="Immutable original structured output (BaseModel, dataclass, DataFrame)"
    )
    is_structured: bool = Field(
        default=False,
        description="Whether the output is structured"
    )
    metadata: Dict[str, Any] = Field(
        default_factory=dict,
        description="Additional metadata associated with the response"
    )
    output_mode: OutputMode = Field(
        default=OutputMode.DEFAULT,
        description="The output mode used for rendering (markdown, html, json, etc.)"
    )

    class Config:
        """Pydantic configuration for AIMessage."""
        # Allow arbitrary types for output field (pandas DataFrames, etc.)
        arbitrary_types_allowed = True

    @property
    def content(self) -> Any:
        """
        Get content as a string. This is an alias for to_text property
        that provides a more intuitive API and compatibility with standard
        AI message formats.

        Returns:
            str: The text representation of the output
        """
        return self.output

    @content.setter
    def content(self, value: Any) -> None:
        """
        Set content by updating the output field.

        Args:
            value: The content string to set
        """
        self.output = value

    @property
    def to_text(self) -> str:
        """Get text representation of output."""
        if self.response:
            return self.response
        if isinstance(self.output, str):
            return self.output
        elif isinstance(self.output, dict) and 'content' in self.output:
            # Handle MessageResponse-style format
            content = self.output['content']
            if isinstance(content, list) and content:
                return content[0].get('text', str(self.output))
            return str(content)
        elif hasattr(self.output, 'to_string'):
            # Handle pandas DataFrames
            return self.output.to_string()
        else:
            # Return the output *as is*
            return str(self.output)

    @property
    def has_tools(self) -> bool:
        """Check if tools were used."""
        return len(self.tool_calls) > 0

    def add_tool_call(self, tool_call: ToolCall) -> None:
        """Add a tool call to the response."""
        self.tool_calls.append(tool_call)  # pylint: disable=E1101 # noqa

    # Context Information:
    @property
    def has_context(self) -> bool:
        """Check if any context (vector or conversation) was used."""
        return self.used_vector_context or self.used_conversation_history

    @property
    def context_summary(self) -> Dict[str, Any]:
        """Get a summary of context usage."""
        return {
            'used_vector_context': self.used_vector_context,
            'used_conversation_history': self.used_conversation_history,
            'vector_context_length': self.vector_context_length,
            'conversation_context_length': self.conversation_context_length,
            'search_results_count': self.search_results_count,
            'search_type': self.search_type,
            'context_sources_count': len(self.context_sources) if self.context_sources else 0
        }

    def set_vector_context_info(
        self,
        used: bool,
        context_length: int = 0,
        search_results_count: int = 0,
        search_type: Optional[str] = None,
        score_threshold: Optional[float] = None,
        sources: Optional[List[str]] = None,
        source_documents: Optional[List[SourceDocument]] = None
    ) -> None:
        """Set vector context information."""
        self.used_vector_context = used
        self.vector_context_length = context_length
        self.search_results_count = search_results_count
        self.search_type = search_type
        self.search_score_threshold = score_threshold
        if source_documents:
            self.source_documents.extend(source_documents)  # noqa
        if sources:
            self.context_sources.extend(sources)   # pylint: disable=E1101 # noqa

    def set_conversation_context_info(
        self,
        used: bool,
        context_length: int = 0
    ) -> None:
        """Set conversation context information."""
        self.used_conversation_history = used
        self.conversation_context_length = context_length

    def to_dict(self) -> Dict[str, Any]:
        """Convert to dictionary for serialization."""
        return self.model_dump()

    def get_context_metadata(self) -> Dict[str, Any]:
        """Get metadata about context usage for logging/analytics."""
        return {
            'context_usage': {
                'vector_context': {
                    'used': self.used_vector_context,
                    'length': self.vector_context_length,
                    'search_results': self.search_results_count,
                    'search_type': self.search_type,
                    'score_threshold': self.search_score_threshold,
                    'sources_count': len(self.context_sources) if self.context_sources else 0
                },
                'conversation_history': {
                    'used': self.used_conversation_history,
                    'length': self.conversation_context_length
                }
            },
            'tools': {
                'used': self.has_tools,
                'count': len(self.tool_calls)
            },
            'timing': {
                'created_at': self.created_at.isoformat()  # pylint: disable=E1101 # noqa
            }
        }

    def render_as(
        self,
        mode: OutputMode,
        formatter: Any,
        **kwargs
    ) -> 'AIMessage':
        """
        Create a new AIMessage with different rendering.
        Preserves structured_output.
        """
        source = self.structured_output if self.is_structured else self.output

        return AIMessage(
            **{
                **self.model_dump(exclude={'output', 'output_mode'}),
                'output': formatter.format(mode, source, **kwargs),
                'output_mode': mode
            }
        )

# Factory functions to create AIMessage from different providers
class AIMessageFactory:
    """Factory to create AIMessage from different provider responses."""

    @staticmethod
    def from_completion(
        response: Any,
        input_text: str,
        structured_output: Optional[Any] = None,
        output_mode: OutputMode = OutputMode.DEFAULT,
        formatter: Optional[Any] = None,
        **kwargs
    ) -> AIMessage:
        """Create AIMessage with proper separation of concerns"""

        # Determine the raw output
        raw_output = structured_output or response.content

        # Create base message
        message = AIMessage(
            input=input_text,
            output=raw_output,  # Initially set to raw
            structured_output=structured_output,
            is_structured=structured_output is not None,
            output_mode=output_mode,
            **kwargs
        )

        # Apply formatting if needed
        if output_mode != OutputMode.DEFAULT and formatter:
            # Render from structured_output if available, else from output
            source = message.structured_output if message.is_structured else message.output
            message.output = formatter.format(output_mode, source)

        return message

    @staticmethod
    def from_openai(
        response: Any,
        input_text: str,
        model: str,
        user_id: Optional[str] = None,
        session_id: Optional[str] = None,
        turn_id: Optional[str] = None,
        structured_output: Any = None
    ) -> AIMessage:
        """Create AIMessage from OpenAI response."""
        message = response.choices[0].message

        # Handle tool calls
        tool_calls = []
        if hasattr(message, 'tool_calls') and message.tool_calls:
            tool_calls.extend(
                ToolCall(
                    id=tc.id,
                    name=tc.function.name,
                    arguments=(
                        tc.function.arguments
                        if isinstance(tc.function.arguments, dict)
                        else eval(tc.function.arguments)
                    ),
                )
                for tc in message.tool_calls
            )

        finish_reason = getattr(response.choices[0], "finish_reason", None)
        stop_reason = getattr(response.choices[0], "stop_reason", None)

        return AIMessage(
            input=input_text,
            output=structured_output or message.content,
            is_structured=structured_output is not None,
            structured_output=structured_output,
            model=model,
            provider="openai",
            usage=CompletionUsage.from_openai(response.usage),
            stop_reason=stop_reason or finish_reason,
            finish_reason=finish_reason,
            tool_calls=tool_calls,
            user_id=user_id,
            session_id=session_id,
            turn_id=turn_id,
            raw_response=response.dict() if hasattr(response, 'dict') else response.__dict__
        )

    @staticmethod
    def from_groq(
        response: Any,
        input_text: str,
        model: str,
        user_id: Optional[str] = None,
        session_id: Optional[str] = None,
        turn_id: Optional[str] = None,
        structured_output: Any = None
    ) -> AIMessage:
        """Create AIMessage from Groq response."""
        message = response.choices[0].message

        # Handle tool calls
        tool_calls = []
        if hasattr(message, 'tool_calls') and message.tool_calls:
            tool_calls.extend(
                ToolCall(
                    id=tc.id,
                    name=tc.function.name,
                    arguments=(
                        tc.function.arguments
                        if isinstance(tc.function.arguments, dict)
                        else eval(tc.function.arguments)
                    ),
                )
                for tc in message.tool_calls
            )

        return AIMessage(
            input=input_text,
            output=structured_output or message.content,
            is_structured=structured_output is not None,
            structured_output=structured_output,
            model=model,
            provider="groq",
            usage=CompletionUsage.from_groq(response.usage),
            stop_reason=response.choices[0].finish_reason,
            finish_reason=response.choices[0].finish_reason,
            tool_calls=tool_calls,
            user_id=user_id,
            session_id=session_id,
            turn_id=turn_id,
            raw_response=response.dict() if hasattr(response, 'dict') else response.__dict__,
            response=message.content
        )

    @staticmethod
    def from_grok(
        response: Any,
        input_text: str,
        model: str,
        user_id: Optional[str] = None,
        session_id: Optional[str] = None,
        turn_id: Optional[str] = None,
        structured_output: Any = None
    ) -> AIMessage:
        """Create AIMessage from Grok response."""
        # xAI SDK response structure needs verification, assuming OpenAI-like for now based on 'usage'
        # If response is a Pydantic model from structured output, content might be parsed.
        
        content = ""
        tool_calls = []
        finish_reason = None

        # Check if it's a simple text response or object
        if hasattr(response, 'content'):
            content = response.content
        elif isinstance(response, dict) and 'content' in response:
            content = response['content']
            
        # Try to extract standard fields if present
        if hasattr(response, 'tool_calls') and response.tool_calls:
             # Map tool calls
             tool_calls.extend(
                ToolCall(
                    id=tc.id,
                    name=tc.function.name,
                    arguments=(
                        tc.function.arguments
                        if isinstance(tc.function.arguments, dict)
                        else eval(tc.function.arguments)
                    ),
                )
                for tc in response.tool_calls
            )
            
        return AIMessage(
            input=input_text,
            output=structured_output or content,
            is_structured=structured_output is not None,
            structured_output=structured_output,
            model=model,
            provider="grok",
            usage=CompletionUsage.from_grok(response.usage) if hasattr(response, 'usage') else CompletionUsage(),
            stop_reason=finish_reason,
            finish_reason=finish_reason,
            tool_calls=tool_calls,
            user_id=user_id,
            session_id=session_id,
            turn_id=turn_id,
            raw_response=response.__dict__ if hasattr(response, '__dict__') else response,
            response=content
        )

    @staticmethod
    def from_claude(
        response: Dict[str, Any],
        input_text: str,
        model: str,
        user_id: Optional[str] = None,
        session_id: Optional[str] = None,
        turn_id: Optional[str] = None,
        structured_output: Any = None,
        tool_calls: List[ToolCall] = None
    ) -> AIMessage:
        """Create AIMessage from Claude response."""
        # Extract text content
        content = ""
        for content_block in response.get("content", []):
            if content_block.get("type") == "text":
                content += content_block.get("text", "")

        return AIMessage(
            input=input_text,
            output=structured_output or content,
            is_structured=structured_output is not None,
            structured_output=structured_output,
            model=model,
            provider="claude",
            usage=CompletionUsage.from_claude(response.get("usage", {})),
            stop_reason=response.get("stop_reason"),
            finish_reason=response.get("stop_reason"),
            tool_calls=tool_calls or [],
            user_id=user_id,
            session_id=session_id,
            turn_id=turn_id,
            raw_response=response,
            response=content if isinstance(content, str) else str(content)
        )

    @staticmethod
    def _sanitize_for_json(data: Any) -> Any:
        """Recursively sanitize data for JSON serialization."""
        if isinstance(data, dict):
            return {str(k): AIMessageFactory._sanitize_for_json(v) for k, v in data.items()}
        elif isinstance(data, list):
            return [AIMessageFactory._sanitize_for_json(item) for item in data]
        elif isinstance(data, (str, int, float, bool, type(None))):
            return data
        elif hasattr(data, 'dict'):
            return AIMessageFactory._sanitize_for_json(data.dict())
        elif hasattr(data, 'model_dump'):
            return AIMessageFactory._sanitize_for_json(data.model_dump())
        elif hasattr(data, '__dict__'):
            return AIMessageFactory._sanitize_for_json(data.__dict__)
        else:
            return str(data)

    @staticmethod
    def from_gemini(
        response: Any,
        input_text: str,
        model: str,
        user_id: Optional[str] = None,
        session_id: Optional[str] = None,
        turn_id: Optional[str] = None,
        structured_output: Any = None,
        tool_calls: List[ToolCall] = None,
        # Add these new parameters:
        conversation_history: Optional[Any] = None,
        text_response: Optional[str] = None,
        files: Optional[List[Path]] = None,
        images: Optional[List[Any]] = None,
        code: Optional[str] = None
    ) -> AIMessage:
        """Create AIMessage from Gemini/Vertex AI response."""
        # Handle both direct text responses and response objects
        if text_response:
            content = text_response
        elif hasattr(response, 'text'):
            content = response.text
        else:
            content = str(response)

        # Extract usage information
        usage_dict = {}
        if hasattr(response, 'usage_metadata') and response.usage_metadata:
            # Vertex AI format
            usage_dict = {
                'prompt_token_count': response.usage_metadata.prompt_token_count,
                'candidates_token_count': response.usage_metadata.candidates_token_count,
                'total_token_count': response.usage_metadata.total_token_count
            }
        elif hasattr(response, 'usage'):
            # Standard Gemini API format
            usage_dict = response.usage.__dict__ if hasattr(response.usage, '__dict__') else {}

        # Sanitize raw response
        try:
             raw = response.__dict__ if hasattr(response, '__dict__') else str(response)
             sanitized_raw = AIMessageFactory._sanitize_for_json(raw)
        except Exception:
             sanitized_raw = str(response)

        # Encode images as base64 for inline rendering
        documents_with_images = []
        if images:
            for img_path in images:
                try:
                    img_path = Path(img_path) if not isinstance(img_path, Path) else img_path
                    if img_path.exists():
                        mime_type, _ = mimetypes.guess_type(str(img_path))
                        mime_type = mime_type or 'image/png'
                        with open(img_path, 'rb') as f:
                            image_bytes = f.read()
                        base64_data = base64.b64encode(image_bytes).decode('utf-8')
                        # Store as data URI for inline rendering
                        data_uri = f"data:{mime_type};base64,{base64_data}"
                        documents_with_images.append(data_uri)
                except Exception:
                    # If encoding fails, just skip this image
                    pass

        ai_message = AIMessage(
            input=input_text,
            output=structured_output or content,
            is_structured=structured_output is not None,
            structured_output=structured_output,
            model=model,
            provider="gemini",  # Will be overridden to "vertex_ai" in VertexAIClient
            usage=CompletionUsage.from_gemini(usage_dict),
            stop_reason="completed",
            finish_reason="completed",
            tool_calls=tool_calls or [],
            user_id=user_id,
            session_id=session_id,
            turn_id=turn_id,
            raw_response=sanitized_raw,
            response=content,
            files=files or [],
            images=images or [],
            documents=documents_with_images,
            code=code,
        )

        if conversation_history:
            ai_message.used_conversation_history = True
            ai_message.conversation_context_length = len(conversation_history.turns) if hasattr(conversation_history, 'turns') else 0

        return ai_message


    @staticmethod
    def create_message(
        response: Any,
        input_text: str,
        model: str,
        user_id: Optional[str] = None,
        session_id: Optional[str] = None,
        turn_id: Optional[str] = None,
        structured_output: Any = None,
        tool_calls: List[ToolCall] = None,
        # Add these new parameters:
        conversation_history: Optional[Any] = None,
        text_response: Optional[str] = None,
        usage: Optional[CompletionUsage] = None,
        response_time: Optional[float] = None
    ) -> AIMessage:
        """Create AIMessage from any provider response."""
        if hasattr(response, 'provider'):
            provider = response.provider.lower()
        else:
            provider = "unknown"

        ai_message = AIMessage(
            input=input_text,
            output=structured_output,
            is_structured=structured_output is not None,
            structured_output=structured_output,
            model=model,
            provider=provider,
            stop_reason="completed",
            finish_reason="completed",
            tool_calls=tool_calls or [],
            user_id=user_id,
            session_id=session_id,
            turn_id=turn_id,
            raw_response=response.__dict__ if hasattr(response, '__dict__') else {"result": text_response},
            response=text_response,
            usage=usage or CompletionUsage.from_response(response),
            response_time=response_time
        )
        if conversation_history:
            ai_message.used_conversation_history = True
            ai_message.conversation_context_length = len(conversation_history.turns) if hasattr(conversation_history, 'turns') else 0

        return ai_message

    @staticmethod
    def from_imagen(**kwargs):
        return AIMessage(**kwargs)

    @staticmethod
    def from_speech(**kwargs):
        return AIMessage(**kwargs)

    @staticmethod
    def from_video(**kwargs):
        return AIMessage(**kwargs)

class AgentResponse(BaseModel):
    """
    AgentResponse is a model that defines the structure of the response for Any Parrot agent.
    """
    session_id: Optional[str] = Field(
        default=None,
        description="Unique identifier for the session"
    )
    user_id: Optional[str] = Field(
        default=None,
        description="Unique identifier for the user"
    )
    agent_id: Optional[str] = Field(
        default=None,
        description="Unique identifier for the agent that processed the request"
    )
    agent_name: Optional[str] = Field(
        default="Agentic",
        description="Name of the agent that processed the request"
    )
    status: str = Field(default="success", description="Status of the response")
    question: Optional[str] = Field(
        default=None,
        description="Question made by User"
    )
    response: Optional[AIMessage] = Field(
        default=None,
        description="Response returned by the agent"
    )
    data: Optional[str] = Field(
        default=None,
        description="Data returned by the agent, can be text, JSON, etc."
    )
    # Optional output field for structured data
    output: Optional[Any] = Field(
        default=None,
        description="Output of the agent's processing"
    )
    attributes: Dict[str, str] = Field(
        default_factory=dict,
        description="Attributes associated with the response"
    )
    # Timestamp
    created_at: datetime = Field(
        default_factory=datetime.now, description="Timestamp when response was created"
    )
    # Optional file paths
    transcript: Optional[str] = Field(
        default=None, description="Transcript of the conversation with the agent"
    )
    script_path: Optional[str] = Field(
        default=None, description="Path to the conversational script associated with the session"
    )
    podcast_path: Optional[str] = Field(
        default=None, description="Path to the podcast associated with the session"
    )
    pdf_path: Optional[str] = Field(
        default=None, description="Path to the PDF associated with the session"
    )
    document_path: Optional[str] = Field(
        default=None, description="Path to any document generated during session"
    )
    images: Optional[List[Path]] = Field(
        default_factory=list,
        description="List of image file paths generated by the model"
    )
    media: Optional[List[Path]] = Field(
        default_factory=list,
        description="List of media file paths generated by the model"
    )
    documents: Optional[List[Any]] = Field(
        default_factory=list,
        description="List of document file paths associated with the response"
    )
    # complete list of generated files:
    files: List[str] = Field(
        default_factory=list,
        description="List of files generated during the session"
    )
    turn_id: Optional[str] = Field(
        default=None, description="Unique identifier for the conversation turn"
    )

    class Config:
        """Pydantic configuration for AgentResponse."""
        # Allow arbitrary types for output field (pandas DataFrames, etc.)
        arbitrary_types_allowed = False
        # Allow extra fields if needed
        extra = "allow"
        # Use enum values
        use_enum_values = True
        # Validate assignment
        validate_assignment = True

    @property
    def content(self) -> Any:
        """
        Get content as a string. This is an alias for to_text property
        that provides a more intuitive API and compatibility with standard
        AI message formats.

        Returns:
            str: The text representation of the output
        """
        return self.output

    @content.setter
    def content(self, value: Any) -> None:
        """
        Set content by updating the output field.

        Args:
            value: The content string to set
        """
        self.output = value

    @model_validator(mode='after')
    def sync_documents_and_paths(self):
        """
        Automatically populate documents list from individual path fields
        and sync documents from AIMessage if present.
        """
        # Sync from AIMessage if it exists
        if self.response and isinstance(self.response, AIMessage):
            if hasattr(self.response, 'documents') and self.response.documents:
                # Add AIMessage documents to our documents list
                for doc in self.response.documents:
                    if doc not in self.documents:
                        self.documents.append(doc)

            if hasattr(self.response, 'files') and self.response.files:
                # Add AIMessage files to our files list
                for file in self.response.files:
                    if file not in self.files:
                        self.files.append(file)

            if hasattr(self.response, 'images') and self.response.images:
                # Add AIMessage images to our images list
                for img in self.response.images:
                    if img not in self.images:
                        self.images.append(img)

            if hasattr(self.response, 'media') and self.response.media:
                # Add AIMessage media to our media list
                for med in self.response.media:
                    if med not in self.media:
                        self.media.append(med)

        # Auto-add individual paths to appropriate collections
        path_mappings = [
            (self.script_path, self.documents, 'script'),
            (self.podcast_path, self.files, 'podcast'),
            (self.pdf_path, self.documents, 'pdf'),
            (self.document_path, self.documents, 'document'),
        ]

        for path, collection, path_type in path_mappings:
            if path and path not in collection:
                collection.append(path)

        return self

    def add_document(self, path: Union[str, Path]) -> None:
        """
        Add a document to the documents list also document path.

        Args:
            path: Path to the document
        """
        path_str = str(path)
        if path_str not in self.documents:
            self.documents.append(path_str)
        self.document_path = path_str
        self.add_file(path_str)

    def add_file(self, path: Union[str, Path]) -> None:
        """
        Add a file to the files list.

        Args:
            path: Path to the file
        """
        path_str = str(path)
        if path_str not in self.files:
            self.files.append(path_str)

    def add_image(self, path: Union[str, Path]) -> None:
        """
        Add an image to the images list.

        Args:
            path: Path to the image
        """
        path_str = str(path)
        if path_str not in self.images:
            self.images.append(path_str)
        self.add_file(path_str)

    def add_media(self, path: Union[str, Path]) -> None:
        """
        Add a media file to the media list.

        Args:
            path: Path to the media file
        """
        path_str = str(path)
        if path_str not in self.media:
            self.media.append(path_str)
        self.add_file(path_str)

    def set_podcast_path(self, path: Union[str, Path]) -> None:
        """
        Set the podcast_path and automatically add to files list.

        Args:
            path: Path to the podcast
        """
        self.podcast_path = str(path)
        self.add_file(path)

    def set_pdf_path(self, path: Union[str, Path]) -> None:
        """
        Set the pdf_path and automatically add to documents list.

        Args:
            path: Path to the PDF
        """
        self.pdf_path = str(path)
        self.add_document(path)


class InvokeResult(BaseModel):
    """Lightweight result from a stateless invoke() call.

    Returned by ``AbstractClient.invoke()`` instead of the heavier
    ``AIMessage``. Carries only what is needed for structured extraction:
    the parsed output, the type class, model name, token usage, and the
    raw provider response for debugging.

    Attributes:
        output: Parsed result — Pydantic model instance, dataclass, or raw str.
        output_type: The type class used for structured output (stores the class
            itself so callers can use ``isinstance`` checks).
        model: Model identifier used for this invocation.
        usage: Token usage statistics.
        raw_response: Provider's raw response object, kept for debugging.
    """

    model_config = ConfigDict(arbitrary_types_allowed=True)

    output: Any = Field(
        description="Parsed result — Pydantic model instance, dataclass, or raw str"
    )
    output_type: Optional[type] = Field(
        default=None,
        description="The type class used for structured output"
    )
    model: str = Field(
        description="Model used for this invocation"
    )
    usage: CompletionUsage = Field(
        description="Token usage statistics"
    )
    raw_response: Optional[Any] = Field(
        default=None,
        description="Provider's raw response for debugging"
    )
