"""Agent command executor"""

from typing import AsyncIterator, Optional

from diona.project.cli.core.models import ParsedInput, ExecutionResult, StreamChunk
from diona.project.cli.executor import BaseExecutor


class AgentExecutor(BaseExecutor):
    """Executor for agent commands"""

    _instance: Optional["AgentExecutor"] = None

    def __new__(cls) -> "AgentExecutor":
        if cls._instance is None:
            cls._instance = super().__new__(cls)
            cls._instance._initialized = False
        return cls._instance

    def __init__(self):
        if self._initialized:
            return
        self._initialized = True
    
    def can_execute(self, parsed_input: ParsedInput) -> bool:
        return parsed_input.input_type.value == "agent"
    
    def execute(self, parsed_input: ParsedInput) -> ExecutionResult:
        try:
            from diona.project.agent import get_session_manager, BaseModelProvider
            
            session_mgr = get_session_manager()
            model_provider = BaseModelProvider.get_default()
            
            result = session_mgr.run_with_session(
                parsed_input.raw_input,
                model_provider
            )
            
            return ExecutionResult(
                success=True,
                return_code=0,
                stdout=result
            )
        except Exception as e:
            return ExecutionResult(
                success=False,
                return_code=-1,
                error_message=str(e)
            )
    
    async def execute_stream(self, parsed_input: ParsedInput) -> AsyncIterator[StreamChunk]:
        try:
            from diona.project.agent import get_session_manager, BaseModelProvider
            
            session_mgr = get_session_manager()
            model_provider = BaseModelProvider.get_default()
            
            # 异步迭代 run_with_session_streamed 返回的 token
            async for token in session_mgr.run_with_session_streamed(
                parsed_input.raw_input,
                model_provider
            ):
                yield StreamChunk(content=token)
            
            yield StreamChunk(content="", is_final=True, return_code=0)
            
        except Exception as e:
            yield StreamChunk(
                content=str(e),
                is_error=True,
                is_final=True,
                return_code=-1
            )
