import os
import sys
from abc import ABC, abstractmethod
from typing import AsyncIterator
from diona.project.cli.engine.result import ParsedInput, StreamChunk, InputType
from diona.project.cli.commands.base import Command


class ExecutionContext:
    """Execution context for storing command results"""
    
    def __init__(self):
        self._variables = {
            "$0": "0",
            "$1": "",
            "$2": "",
            "$3": "",
        }
    
    def set_result(self, return_code: int, stdout: str, stderr: str, full_result: str = None):
        self._variables["$0"] = str(return_code)
        self._variables["$1"] = stdout
        self._variables["$2"] = stderr
        if full_result is not None:
            self._variables["$3"] = full_result
        else:
            self._variables["$3"] = stdout if stdout else stderr
    
    def get(self, name: str) -> str:
        return self._variables.get(name, "")


_context: ExecutionContext = None


def get_execution_context() -> ExecutionContext:
    global _context
    if _context is None:
        _context = ExecutionContext()
    return _context


class CommandExecutor(ABC):
    @abstractmethod
    def can_execute(self, parsed: ParsedInput) -> bool:
        pass

    @abstractmethod
    async def execute(self, parsed: ParsedInput) -> AsyncIterator[StreamChunk]:
        pass


class BuiltinExecutor(CommandExecutor):
    def can_execute(self, parsed: ParsedInput) -> bool:
        return parsed.input_type == InputType.BUILTIN

    async def execute(self, parsed: ParsedInput) -> AsyncIterator[StreamChunk]:
        from diona.project.cli.commands import get_command_registry
        registry = get_command_registry()
        
        input_with_vars = self._replace_variables(parsed.raw_input)
        
        cmd = registry.get(parsed.command_name)
        
        if cmd is None:
            yield StreamChunk(content=f"Command not found: {parsed.command_name}\n", is_error=True, is_final=True)
            return
        
        try:
            args = input_with_vars.split()[1:] if len(input_with_vars.split()) > 1 else []
            old_stdout = sys.stdout
            old_stderr = sys.stderr
            from io import StringIO
            stdout_capture = StringIO()
            stderr_capture = StringIO()
            sys.stdout = stdout_capture
            sys.stderr = stderr_capture
            
            try:
                result = cmd.execute(args)
            finally:
                sys.stdout = old_stdout
                sys.stderr = old_stderr
            
            stdout_output = stdout_capture.getvalue()
            stderr_output = stderr_capture.getvalue()
            
            context = get_execution_context()
            full_result = str(result) if result is not None else stdout_output
            context.set_result(0, stdout_output, stderr_output, full_result)
            
            if parsed.command_name == "exit":
                yield StreamChunk(content="Goodbye!\n", is_final=True)
            elif stdout_output:
                yield StreamChunk(content=stdout_output, is_final=True)
            elif stderr_output:
                yield StreamChunk(content=stderr_output, is_error=True, is_final=True)
        except Exception as e:
            context = get_execution_context()
            context.set_result(1, "", str(e), str(e))
            yield StreamChunk(content=f"Error: {str(e)}\n", is_error=True, is_final=True)

    def _replace_variables(self, raw_input: str) -> str:
        context = get_execution_context()
        for var_name in ["$0", "$1", "$2", "$3"]:
            if var_name in raw_input:
                raw_input = raw_input.replace(var_name, context.get(var_name))
        return raw_input


class RegisteredExecutor(CommandExecutor):
    """Executor for registered commands from CommandModel"""
    def can_execute(self, parsed: ParsedInput) -> bool:
        return parsed.input_type == InputType.REGISTERED

    async def execute(self, parsed: ParsedInput) -> AsyncIterator[StreamChunk]:
        from diona.project.cli.commands import get_command_registry
        registry = get_command_registry()
        
        input_with_vars = self._replace_variables(parsed.raw_input)
        
        model = registry.get_model(parsed.command_name)
        if model is None:
            yield StreamChunk(content=f"Registered command not found: {parsed.command_name}\n", is_error=True, is_final=True)
            return
        
        try:
            from io import StringIO
            import sys
            
            args = input_with_vars.split()[1:] if len(input_with_vars.split()) > 1 else []
            old_stdout = sys.stdout
            old_stderr = sys.stderr
            stdout_capture = StringIO()
            stderr_capture = StringIO()
            sys.stdout = stdout_capture
            sys.stderr = stderr_capture
            
            try:
                result = model.execute(args)
            finally:
                sys.stdout = old_stdout
                sys.stderr = old_stderr
            
            stdout_output = stdout_capture.getvalue()
            stderr_output = stderr_capture.getvalue()
            
            context = get_execution_context()
            context.set_result(0, stdout_output, stderr_output, str(result))
            
            yield StreamChunk(content=stdout_output, is_final=True)
        except Exception as e:
            yield StreamChunk(content=f"Error: {str(e)}\n", is_error=True, is_final=True)

    def _replace_variables(self, raw_input: str) -> str:
        context = get_execution_context()
        for var_name in ["$0", "$1", "$2", "$3"]:
            if var_name in raw_input:
                raw_input = raw_input.replace(var_name, context.get(var_name))
        return raw_input


class SystemExecutor(CommandExecutor):
    def can_execute(self, parsed: ParsedInput) -> bool:
        return parsed.input_type == InputType.SYSTEM

    async def execute(self, parsed: ParsedInput) -> AsyncIterator[StreamChunk]:
        import subprocess
        
        input_with_vars = self._replace_variables(parsed.raw_input)
        
        try:
            result = subprocess.run(
                input_with_vars,
                capture_output=True,
                text=True,
                shell=True,
                cwd=os.getcwd() if os.name == 'nt' else None
            )
            
            context = get_execution_context()
            full_result = result.stdout if result.stdout else result.stderr
            context.set_result(result.returncode, result.stdout, result.stderr, full_result)
            
            output = result.stdout or result.stderr or "(no output)"
            yield StreamChunk(content=output + "\n", is_final=True)
        except Exception as e:
            yield StreamChunk(content=f"Error: {str(e)}\n", is_error=True, is_final=True)

    def _replace_variables(self, raw_input: str) -> str:
        context = get_execution_context()
        for var_name in ["$0", "$1", "$2", "$3"]:
            if var_name in raw_input:
                raw_input = raw_input.replace(var_name, context.get(var_name))
        return raw_input


class AgentExecutor(CommandExecutor):
    def can_execute(self, parsed: ParsedInput) -> bool:
        return parsed.input_type == InputType.AGENT

    async def execute(self, parsed: ParsedInput) -> AsyncIterator[StreamChunk]:
        from diona.project.agent import get_agent_facade
        facade = get_agent_facade()
        
        try:
            async for token in facade.run_streamed(parsed.raw_input):
                yield StreamChunk(content=token, is_final=False)
            yield StreamChunk(content="", is_final=True)
        except Exception as e:
            yield StreamChunk(content=f"Error: {str(e)}\n", is_error=True, is_final=True)
