from __future__ import annotations

from turbo_agent_core.schema.agents import (
    Agent, Tool, Character, 
    APITool, LLMTool, AgentTool, 
    FrontendTool, BuiltinTool, MCPTool,
    LLMModel
)
from turbo_agent_core.schema.states import Message, MessageRole
from turbo_agent_runtime.executors.agent import AgentRuntime
from turbo_agent_runtime.executors.llm import LLMModelRuntime
from turbo_agent_runtime.executors.character import CharacterRuntime
from turbo_agent_runtime.executors.tool.api_tool import APIToolRuntime
from turbo_agent_runtime.executors.tool.llm_tool import LLMToolRuntime
from turbo_agent_runtime.executors.tool.agent_tool import AgentToolRuntime
from turbo_agent_runtime.executors.tool.frontend_tool import FrontendToolRuntime
from turbo_agent_runtime.executors.tool.builtin_tool import BuiltinToolRuntime
from turbo_agent_runtime.executors.tool.mcp_tool import MCPToolRuntime

def _make_proxy(runtime_cls, method_name):
    def proxy(self, *args, **kwargs):
        # Convert self (Entity) to Runtime instance
        # We use model_validate to copy data. 
        # Note: This assumes Runtime has compatible schema with Entity.
        runtime_instance = runtime_cls.model_validate(self.model_dump())
        method = getattr(runtime_instance, method_name)
        return method(*args, **kwargs)
    return proxy

def _make_async_proxy(runtime_cls, method_name):
    async def proxy(self, *args, **kwargs):
        # Convert self (Entity) to Runtime instance
        runtime_instance = runtime_cls.model_validate(self.model_dump())
        method = getattr(runtime_instance, method_name)
        return await method(*args, **kwargs)
    return proxy

def _make_generator_proxy(runtime_cls, method_name):
    def proxy(self, *args, **kwargs):
        # Convert self (Entity) to Runtime instance
        runtime_instance = runtime_cls.model_validate(self.model_dump())
        method = getattr(runtime_instance, method_name)
        yield from method(*args, **kwargs)
    return proxy

def _make_async_generator_proxy(runtime_cls, method_name):
    async def proxy(self, *args, **kwargs):
        # Convert self (Entity) to Runtime instance
        runtime_instance = runtime_cls.model_validate(self.model_dump())
        method = getattr(runtime_instance, method_name)
        async for item in method(*args, **kwargs):
            yield item
    return proxy

def patch_turbo_entity_methods():
    """
    Patch Core Entities with Runtime methods using delegation.
    This ensures that the Runtime class is instantiated with the Entity's data,
    allowing access to Runtime-specific methods and properties.
    """
    
    # Patch Agent
    Agent.run = _make_proxy(AgentRuntime, "run")
    Agent.stream = _make_generator_proxy(AgentRuntime, "stream")
    Agent.a_run = _make_async_proxy(AgentRuntime, "a_run")
    Agent.a_stream = _make_async_generator_proxy(AgentRuntime, "a_stream")
    
    # Patch Character
    Character.run = _make_proxy(CharacterRuntime, "run")
    Character.stream = _make_generator_proxy(CharacterRuntime, "stream")
    Character.a_run = _make_async_proxy(CharacterRuntime, "a_run")
    Character.a_stream = _make_async_generator_proxy(CharacterRuntime, "a_stream")
    
    # Patch LLMModel
    LLMModel.run = _make_proxy(LLMModelRuntime, "run")
    LLMModel.stream = _make_generator_proxy(LLMModelRuntime, "stream")
    LLMModel.a_run = _make_async_proxy(LLMModelRuntime, "a_run")
    LLMModel.a_stream = _make_async_generator_proxy(LLMModelRuntime, "a_stream")
    
    # Patch Tool Versions
    print(f"DEBUG: APIToolRuntime.a_run: {APIToolRuntime.a_run}")
    APITool.run = _make_proxy(APIToolRuntime, "run")
    APITool.stream = _make_generator_proxy(APIToolRuntime, "stream")
    APITool.a_run = _make_async_proxy(APIToolRuntime, "a_run")
    APITool.a_stream = _make_async_generator_proxy(APIToolRuntime, "a_stream")
    
    LLMTool.run = _make_proxy(LLMToolRuntime, "run")
    LLMTool.stream = _make_generator_proxy(LLMToolRuntime, "stream")
    LLMTool.a_run = _make_async_proxy(LLMToolRuntime, "a_run")
    LLMTool.a_stream = _make_async_generator_proxy(LLMToolRuntime, "a_stream")
    
    AgentTool.a_run = _make_async_proxy(AgentToolRuntime, "run")
    AgentTool.a_stream = _make_async_generator_proxy(AgentToolRuntime, "stream")
    
    # Patch FrontendTool
    FrontendTool.run = _make_proxy(FrontendToolRuntime, "run")
    FrontendTool.stream = _make_generator_proxy(FrontendToolRuntime, "stream")
    FrontendTool.a_run = _make_async_proxy(FrontendToolRuntime, "a_run")
    FrontendTool.a_stream = _make_async_generator_proxy(FrontendToolRuntime, "a_stream")
    
    # Patch BuiltinTool
    BuiltinTool.run = _make_proxy(BuiltinToolRuntime, "run")
    BuiltinTool.stream = _make_generator_proxy(BuiltinToolRuntime, "stream")
    BuiltinTool.a_run = _make_async_proxy(BuiltinToolRuntime, "a_run")
    BuiltinTool.a_stream = _make_async_generator_proxy(BuiltinToolRuntime, "a_stream")
    
    # Patch MCPTool
    MCPTool.run = _make_proxy(MCPToolRuntime, "run")
    MCPTool.stream = _make_generator_proxy(MCPToolRuntime, "stream")
    MCPTool.a_run = _make_async_proxy(MCPToolRuntime, "a_run")
    MCPTool.a_stream = _make_async_generator_proxy(MCPToolRuntime, "a_stream")

    # Tool (Base) patching removed as Tool no longer holds versions container
    # If we need to run a generic Tool, it should be one of the subclasses above.

