from typing import List, Optional, Union
from azure.identity import DefaultAzureCredential
from azure.ai.projects.aio import AIProjectClient
from azure.ai.evaluation._converters._models import _BUILT_IN_DESCRIPTIONS, _BUILT_IN_PARAMS
from openai.types.responses import ResponseOutputItem
import json
from .types import AgentResponse

class FoundryAgentRunner:
    """
    Runner for Microsoft Foundry Agents.
    """
    def __init__(self, name: str, project_endpoint: str = None, version: Optional[str] = None):
        self.name = name
        self.project_endpoint = project_endpoint
        self.version = version
        
        if not self.project_endpoint:
            raise ValueError(
                "Foundry Agent connection requires a project endpoint. "
                "Please provide 'project_endpoint' in the marker or environment variables."
            )
            
        self.credential = DefaultAzureCredential()

    async def run(self, query: str) -> AgentResponse:
        """
        Run the agent with the given query.
        Returns a standardized AgentResponse object.
        Supports both prompt agents and hosted agents.
        """
        async with AIProjectClient(
            endpoint=self.project_endpoint,
            credential=self.credential
        ) as project_client:
            # Using the agent service client
            async with project_client.get_openai_client() as openai_client:
                # 1. Resolve Agent (by name, optionally by version)
                if self.version is not None:
                    version_details = await project_client.agents.get_version(
                        agent_name=self.name, agent_version=self.version
                    )
                    definition = version_details.definition
                    agent_name = version_details.name
                    agent_version = version_details.version
                else:
                    agent = await project_client.agents.get(agent_name=self.name)
                    definition = agent.versions.latest.definition
                    agent_name = agent.name
                    agent_version = None

                agent_kind = definition.kind

                # 2. Create Conversation
                conversation = await openai_client.conversations.create()

                # 3. Branch by agent kind
                if agent_kind == "prompt":
                    response, instructions, tool_definitions = await self._run_prompt_agent(
                        openai_client, conversation, definition, query
                    )
                elif agent_kind == "hosted":
                    response, instructions, tool_definitions = await self._run_hosted_agent(
                        openai_client, conversation, agent_name, agent_version, query
                    )
                else:
                    raise ValueError(
                        f"Unsupported agent kind '{agent_kind}'. "
                        "Only 'prompt' and 'hosted' agents are supported."
                    )

                # 4. Extract output (shared for all agent kinds)
                output_text = response.output_text
                output_items = response.output

                # Discover tool definitions from response output (e.g. MCP tools)
                for output_item in output_items:
                    if output_item.type == "mcp_list_tools":
                        for tool in output_item.tools:
                            tool_definitions.append({
                                "name": tool.name,
                                "description": tool.description,
                                "parameters": tool.input_schema,
                            })

                # Extract tool calls
                tool_calls = self._extract_tool_calls(output_items)

                return AgentResponse(
                    output_text=output_text,
                    tool_calls=tool_calls,
                    tool_definitions=tool_definitions,
                    output_items=output_items,
                    instructions=instructions,
                    raw=response,
                    type="foundry_agent"
                )

    async def _run_prompt_agent(self, openai_client, conversation, definition, query):
        """
        Run a prompt agent by explicitly providing instructions, model, and tools.
        Returns (response, instructions, tool_definitions).
        """
        instructions = definition.instructions

        # Extract tool definitions from agent metadata
        tool_definitions = []
        if definition.tools:
            for tool in definition.tools:
                if tool.type in _BUILT_IN_DESCRIPTIONS:
                    tool_definitions.append({
                        "type": tool.type,
                        "description": _BUILT_IN_DESCRIPTIONS[tool.type],
                        "name": tool.type,
                        "parameters": _BUILT_IN_PARAMS.get(tool.type, {}),
                    })

        # Set auto approval for MCP tools
        tools = []
        if definition.tools:
            for tool in definition.tools:
                if tool.type == "mcp":
                    tool.require_approval = "never"
                tools.append(tool)

        response = await openai_client.responses.create(
            conversation=conversation.id,
            model=definition.model,
            instructions=instructions,
            input=query,
            tools=tools,
        )

        return response, instructions, tool_definitions

    async def _run_hosted_agent(self, openai_client, conversation, agent_name, agent_version, query):
        """
        Run a hosted agent via agent_reference.
        Hosted agents manage their own instructions, model, and tools internally,
        so we delegate execution entirely to the service.
        Returns (response, instructions, tool_definitions).
        """
        agent_ref = {"name": agent_name, "type": "agent_reference"}
        if agent_version:
            agent_ref["version"] = agent_version

        response = await openai_client.responses.create(
            conversation=conversation.id,
            extra_body={"agent_reference": agent_ref},
            input=query,
        )

        # Hosted agents don't expose instructions or tool definitions in metadata
        instructions = None
        tool_definitions = []

        return response, instructions, tool_definitions

    @staticmethod
    def _extract_tool_calls(output_items) -> list:
        """Extract tool calls from response output items."""
        tool_calls = []
        for output_item in output_items:
            if output_item.type == "file_search_call":
                tool_calls.append({
                    "type": "tool_call",
                    "name": "file_search",
                    "tool_call_id": output_item.id
                })
            elif output_item.type == "code_interpreter_call":
                tool_calls.append({
                    "type": "tool_call",
                    "name": "code_interpreter",
                    "tool_call_id": output_item.id,
                    "arguments": {
                        "input": output_item.code
                    }
                })
            elif output_item.type == "mcp_call":
                tool_calls.append({
                    "type": "tool_call",
                    "name": output_item.name,
                    "tool_call_id": output_item.id,
                    "arguments": json.loads(output_item.arguments) if output_item.arguments else {}
                })
            elif output_item.type == "function_call":
                tool_calls.append({
                    "type": "tool_call",
                    "name": output_item.name,
                    "tool_call_id": output_item.call_id,
                    "arguments": json.loads(output_item.arguments) if output_item.arguments else {}
                })
        return tool_calls

    @staticmethod
    def convert_response_for_evaluator(response_obj: AgentResponse) -> Union[str, list]:
        """
        Returns a simplified conversation history for Foundry agents.
        """
        output_items: List[ResponseOutputItem] = response_obj.output_items
        
        ALLOWED_RESPONSE_TYPES = {
            "message", 
            "file_search_call", 
            "code_interpreter_call", 
            "mcp_call", 
            "mcp_list_tools",
            "function_call",
            "function_call_output"
        }
        
        can_support = True
        for output_item in output_items:
            if output_item.type not in ALLOWED_RESPONSE_TYPES:
                can_support = False
                break
        
        if can_support:
            output = []
            for output_item in output_items:
                if output_item.type == "file_search_call":
                    output.append({
                        "role": "assistant",
                        "content": [
                            {
                                "type": "tool_call",
                                "name": "file_search",
                                "tool_call_id": output_item.id,
                                "arguments": {}
                            }
                        ]
                    })
                    output.append({
                        "role": "tool",
                        "tool_call_id": output_item.id,
                        "content": [
                            {
                                "type": "tool_result",
                                "tool_result": output_item.results
                            }
                        ]
                    })
                elif output_item.type == "code_interpreter_call":
                    output.append({
                        "role": "assistant",
                        "content": [
                            {
                                "type": "tool_call",
                                "name": "code_interpreter",
                                "tool_call_id": output_item.id,
                                "arguments": {
                                    "input": output_item.code
                                }
                            }
                        ]
                    })
                elif output_item.type == "mcp_call":
                    output.append({
                        "role": "assistant",
                        "content": [
                            {
                                "type": "tool_call",
                                "name": output_item.name,
                                "tool_call_id": output_item.id,
                                "arguments": json.loads(output_item.arguments) if output_item.arguments else {}
                            }
                        ]
                    })
                    if output_item.output:
                        tool_result = output_item.output
                        try:
                            tool_result = json.loads(output_item.output)
                        except Exception:
                            pass
                        output.append({
                            "role": "tool",
                            "tool_call_id": output_item.id,
                            "content": [
                                {
                                    "type": "tool_result",
                                    "tool_result": tool_result
                                }
                            ]
                        })
                elif output_item.type == "function_call":
                    output.append({
                        "role": "assistant",
                        "content": [
                            {
                                "type": "tool_call",
                                "name": output_item.name,
                                "tool_call_id": output_item.call_id,
                                "arguments": json.loads(output_item.arguments) if output_item.arguments else {}
                            }
                        ]
                    })
                elif output_item.type == "function_call_output":
                    tool_result = output_item.output
                    try:
                        tool_result = json.loads(output_item.output)
                    except Exception:
                        pass
                    output.append({
                        "role": "tool",
                        "tool_call_id": output_item.call_id,
                        "content": [
                            {
                                "type": "tool_result",
                                "tool_result": tool_result
                            }
                        ]
                    })
                elif output_item.type == "message":
                    content = []
                    for content_item in output_item.content:
                        if content_item.type == "output_text":
                            content.append({
                                "type": "text",
                                "text": content_item.text
                            })
                        elif content_item.type == "refusal":
                            content.append({
                                "type": "text",
                                "text": content_item.refusal
                            })
                    output.append({
                        "role": "assistant",
                        "content": content
                    })
            return output
        else:
            return response_obj.output_text
