"""
AbuLang Async Support - Async/await functionality
"""

import asyncio
from typing import Callable, Any, Coroutine, List, Dict, Optional


class AsyncRunner:
    """Runs async code in AbuLang"""
    
    def __init__(self):
        self.loop = None
        self.tasks: Dict[str, asyncio.Task] = {}
    
    def _get_loop(self) -> asyncio.AbstractEventLoop:
        """Get or create event loop"""
        try:
            loop = asyncio.get_running_loop()
        except RuntimeError:
            loop = asyncio.new_event_loop()
            asyncio.set_event_loop(loop)
        return loop
    
    def run(self, coro: Coroutine) -> Any:
        """Run coroutine"""
        loop = self._get_loop()
        return loop.run_until_complete(coro)
    
    def run_sync(self, async_func: Callable, *args, **kwargs) -> Any:
        """Run async function synchronously"""
        coro = async_func(*args, **kwargs)
        return self.run(coro)
    
    def create_task(self, name: str, coro: Coroutine) -> asyncio.Task:
        """Create async task"""
        loop = self._get_loop()
        task = loop.create_task(coro)
        self.tasks[name] = task
        return task
    
    def get_task(self, name: str) -> Optional[asyncio.Task]:
        """Get task by name"""
        return self.tasks.get(name)
    
    def wait_task(self, name: str, timeout: Optional[float] = None) -> Any:
        """Wait for task to complete"""
        task = self.get_task(name)
        if task:
            loop = self._get_loop()
            return loop.run_until_complete(asyncio.wait_for(task, timeout=timeout))
    
    def cancel_task(self, name: str) -> bool:
        """Cancel task"""
        task = self.get_task(name)
        if task:
            task.cancel()
            return True
        return False
    
    def get_all_tasks(self) -> List[str]:
        """Get all task names"""
        return list(self.tasks.keys())
    
    def get_task_status(self, name: str) -> str:
        """Get task status"""
        task = self.get_task(name)
        if not task:
            return "not_found"
        if task.done():
            return "done"
        if task.cancelled():
            return "cancelled"
        return "running"


class AsyncContext:
    """Context for async operations"""
    
    def __init__(self):
        self.runner = AsyncRunner()
    
    async def sleep(self, seconds: float):
        """Async sleep"""
        await asyncio.sleep(seconds)
    
    async def gather(self, *coros) -> List[Any]:
        """Run multiple coroutines concurrently"""
        return await asyncio.gather(*coros)
    
    async def wait_for(self, coro: Coroutine, timeout: float) -> Any:
        """Wait for coroutine with timeout"""
        return await asyncio.wait_for(coro, timeout=timeout)
    
    def run_async(self, async_func: Callable, *args, **kwargs) -> Any:
        """Run async function"""
        return self.runner.run_sync(async_func, *args, **kwargs)


class AsyncDecorator:
    """Decorator for async functions"""
    
    @staticmethod
    def async_func(func: Callable) -> Callable:
        """Decorator to make function async"""
        async def wrapper(*args, **kwargs):
            return func(*args, **kwargs)
        return wrapper
    
    @staticmethod
    def async_timeout(timeout: float):
        """Decorator to add timeout to async function"""
        def decorator(func: Callable) -> Callable:
            async def wrapper(*args, **kwargs):
                return await asyncio.wait_for(func(*args, **kwargs), timeout=timeout)
            return wrapper
        return decorator
    
    @staticmethod
    def async_retry(max_retries: int = 3, delay: float = 1.0):
        """Decorator to retry async function"""
        def decorator(func: Callable) -> Callable:
            async def wrapper(*args, **kwargs):
                for attempt in range(max_retries):
                    try:
                        return await func(*args, **kwargs)
                    except Exception as e:
                        if attempt == max_retries - 1:
                            raise
                        await asyncio.sleep(delay)
            return wrapper
        return decorator


# Global async runner
_async_runner = AsyncRunner()
_async_context = AsyncContext()


def get_async_runner() -> AsyncRunner:
    """Get global async runner"""
    return _async_runner


def get_async_context() -> AsyncContext:
    """Get global async context"""
    return _async_context
