"""
Common utility functions for agents.
""" 
import json
import logging
from typing import Any, Dict, Optional
from agentscope.message import Msg
from agentscope.agent import ReActAgent

logger = logging.getLogger(__name__)


def create_error_chunk(message: str) -> str:
    """ 辅助函数 to create a standard error chunk for streaming responses.
    
    Args:
        message: Error message
    
    Returns:
        Formatted error chunk string
    """ 
    error_chunk = {
        "error": {
            "message": message,
            "type": "server_error"
        }
    }
    return f"data: {json.dumps(error_chunk)}\n\n"


def parse_json_from_msg(msg: Msg) -> Any:
    """ 辅助函数 to extract and parse JSON from agent response.
    
    Args:
        msg: AgentScope message object
    
    Returns:
        Parsed JSON data or None if parsing fails
    """
    try:
        content = msg.content
        
        # Handle list content (common in some AgentScope responses)
        if isinstance(content, list):
            text_parts = []
            for item in content:
                if isinstance(item, dict) and item.get("type") == "text":
                    text_parts.append(item.get("text", ""))
                elif isinstance(item, str):
                    text_parts.append(item)
            content = "\n".join(text_parts)
        
        # Ensure content is string
        if not isinstance(content, str):
            content = str(content)

        # Try to find JSON block in markdown code block
        if "```json" in content:
            json_str = content.split("```json")[1].split("```")[0].strip()
        elif "```" in content:
            json_str = content.split("```")[1].split("```")[0].strip()
        else:
            json_str = content.strip()
        
        return json.loads(json_str)
    except Exception as e:
        logger.error(f"Failed to parse JSON from message: {e}")
        logger.debug(f"Raw content was: {msg.content}")
        return None


async def extract_structured_data(response_msg: Msg, agent: ReActAgent) -> Dict[str, Any]:
    """ 提取 structured data from agent response using multiple strategies.
    
    Args:
        response_msg: The response message from the agent
        agent: The ReActAgent instance for memory access
    
    Returns:
        A dict with 'reply' and 'nodes', or None if extraction fails
    """
    data = None

    # Strategy 1: Check metadata (Direct return from structured_model)
    if response_msg.metadata:
        data = response_msg.metadata
        logger.info("Found structured data in response metadata.")

    # Strategy 2: Check agent memory for the last structured output tool call
    if not data:
        logger.info("No metadata in response, checking agent memory for structured tool call...")
        try:
            memories = await agent.memory.get_memory()
            # Iterate backwards to find the most recent 'generate_response' tool use
            for m in reversed(memories):
                if m.role == "assistant":
                    if isinstance(m.content, list):
                        for item in m.content:
                            if isinstance(item, dict) and item.get("type") == "tool_use" and item.get("name") == "generate_response":
                                data = item.get("input")
                                logger.info(f"Found structured data in memory (list item): {data}")
                                break
                    elif isinstance(m.content, dict):
                        if m.content.get("type") == "tool_use" and m.content.get("name") == "generate_response":
                            data = m.content.get("input")
                            logger.info(f"Found structured data in memory (dict): {data}")
                
                if data:
                    break
        except Exception as e:
            logger.warning(f"Failed to inspect agent memory for fallback extraction: {e}")

    return data


def format_response_chunk(
    content: str,
    model: str = "qwen-plus",
    finish_reason: Optional[str] = None
) -> str:
    """ 格式化 a response chunk for streaming output.
    
    Args:
        content: The content to include in the chunk
        model: The model name
        finish_reason: Optional finish reason (e.g., "stop")
    
    Returns:
        Formatted chunk string
    """ 
    import time
    
    chunk_data = {
        "id": "chatcmpl-react",
        "object": "chat.completion.chunk",
        "created": int(time.time()),
        "model": model,
        "choices": [{
            "index": 0,
            "delta": {"content": content},
            "finish_reason": finish_reason
        }]
    }
    return f"data: {json.dumps(chunk_data)}\n\n"
