"""
Resource Pool for groknroll.

Provides generic resource pooling with automatic cleanup and health checks.
Used for connection pooling (LMHandler, MCP, LSP clients).

Example:
    # Define resource factory
    async def create_lsp_client():
        client = LSPClient()
        await client.start()
        return client

    # Create pool
    pool = ResourcePool(
        factory=create_lsp_client,
        max_size=10,
        idle_timeout=300,
        health_check=lambda client: client.is_alive()
    )

    # Acquire and release
    async with pool.acquire() as client:
        result = await client.go_to_definition(...)
"""

import asyncio
import time
from typing import Any, Callable, Optional


class Resource:
    """
    Wrapper for pooled resources.

    Tracks creation time, last use time, and health status.
    """

    def __init__(self, obj: Any, created_at: float):
        """
        Initialize resource wrapper.

        Args:
            obj: The actual resource object
            created_at: Timestamp when resource was created
        """
        self.obj = obj
        self.created_at = created_at
        self.last_used = created_at
        self.in_use = False

    def mark_used(self):
        """Mark resource as recently used."""
        self.last_used = time.time()


class ResourcePool:
    """
    Generic resource pool with health checks and automatic cleanup.

    Manages a pool of reusable resources (connections, clients, etc.) with:
    - Maximum pool size
    - Idle timeout for unused resources
    - Health checks to remove unhealthy resources
    - Automatic cleanup on pool closure
    """

    def __init__(
        self,
        factory: Callable[[], Any],
        max_size: int = 10,
        idle_timeout: float = 300,
        health_check: Optional[Callable[[Any], bool]] = None,
        cleanup: Optional[Callable[[Any], None]] = None,
    ):
        """
        Initialize resource pool.

        Args:
            factory: Async function to create new resources
            max_size: Maximum number of resources in pool
            idle_timeout: Seconds before idle resource is removed
            health_check: Function to check if resource is healthy
            cleanup: Function to cleanup resource when removed
        """
        self.factory = factory
        self.max_size = max_size
        self.idle_timeout = idle_timeout
        self.health_check = health_check
        self.cleanup = cleanup

        # Pool state
        self.pool: list[Resource] = []
        self.in_use: set[Resource] = set()
        self._lock = asyncio.Lock()
        self._closed = False

        # Start background cleanup task
        self._cleanup_task = asyncio.create_task(self._cleanup_loop())

    async def acquire(self) -> "ResourceContext":
        """
        Acquire a resource from the pool.

        Returns:
            ResourceContext that can be used with 'async with'

        Raises:
            RuntimeError: If pool is closed
        """
        if self._closed:
            raise RuntimeError("Resource pool is closed")

        async with self._lock:
            # Try to get an available resource
            resource = await self._get_or_create()
            resource.in_use = True
            resource.mark_used()
            self.in_use.add(resource)

        return ResourceContext(self, resource)

    async def _get_or_create(self) -> Resource:
        """Get an available resource or create a new one."""
        # Try to find a healthy idle resource
        for resource in self.pool[:]:
            if not resource.in_use:
                # Check health if health_check is provided
                if self.health_check and not self.health_check(resource.obj):
                    # Unhealthy - remove and cleanup
                    self.pool.remove(resource)
                    if self.cleanup:
                        try:
                            if asyncio.iscoroutinefunction(self.cleanup):
                                await self.cleanup(resource.obj)
                            else:
                                self.cleanup(resource.obj)
                        except Exception:
                            pass
                    continue

                # Found a healthy resource
                return resource

        # No available resources - create new if under max size
        if len(self.pool) + len(self.in_use) < self.max_size:
            obj = await self._create_resource()
            resource = Resource(obj, time.time())
            self.pool.append(resource)
            return resource

        # Pool is full - wait for a resource to become available
        # This is a simplified approach; a production version might use a queue
        raise RuntimeError("Resource pool exhausted")

    async def _create_resource(self) -> Any:
        """Create a new resource using the factory."""
        if asyncio.iscoroutinefunction(self.factory):
            return await self.factory()
        else:
            return self.factory()

    async def release(self, resource: Resource):
        """
        Release a resource back to the pool.

        Args:
            resource: Resource to release
        """
        if self._closed:
            # Pool is closed - cleanup resource
            if self.cleanup:
                try:
                    if asyncio.iscoroutinefunction(self.cleanup):
                        await self.cleanup(resource.obj)
                    else:
                        self.cleanup(resource.obj)
                except Exception:
                    pass
            return

        async with self._lock:
            resource.in_use = False
            self.in_use.discard(resource)

    async def _cleanup_loop(self):
        """Background task to remove idle resources."""
        while not self._closed:
            try:
                await asyncio.sleep(30)  # Check every 30 seconds
                await self._remove_idle_resources()
            except asyncio.CancelledError:
                break
            except Exception:
                pass

    async def _remove_idle_resources(self):
        """Remove resources that have been idle too long."""
        async with self._lock:
            now = time.time()
            for resource in self.pool[:]:
                if not resource.in_use and (now - resource.last_used) > self.idle_timeout:
                    # Remove idle resource
                    self.pool.remove(resource)
                    if self.cleanup:
                        try:
                            if asyncio.iscoroutinefunction(self.cleanup):
                                await self.cleanup(resource.obj)
                            else:
                                self.cleanup(resource.obj)
                        except Exception:
                            pass

    async def close(self):
        """Close the pool and cleanup all resources."""
        self._closed = True

        # Cancel cleanup task
        if self._cleanup_task:
            self._cleanup_task.cancel()
            try:
                await self._cleanup_task
            except asyncio.CancelledError:
                pass

        # Cleanup all resources
        async with self._lock:
            all_resources = list(self.pool) + list(self.in_use)
            for resource in all_resources:
                if self.cleanup:
                    try:
                        if asyncio.iscoroutinefunction(self.cleanup):
                            await self.cleanup(resource.obj)
                        else:
                            self.cleanup(resource.obj)
                    except Exception:
                        pass

            self.pool.clear()
            self.in_use.clear()

    def get_stats(self) -> dict[str, int]:
        """Get pool statistics."""
        return {
            "total": len(self.pool) + len(self.in_use),
            "available": len([r for r in self.pool if not r.in_use]),
            "in_use": len(self.in_use),
            "max_size": self.max_size,
        }


class ResourceContext:
    """
    Context manager for pool resources.

    Ensures resources are properly released back to the pool.
    """

    def __init__(self, pool: ResourcePool, resource: Resource):
        """
        Initialize resource context.

        Args:
            pool: Resource pool
            resource: Resource to manage
        """
        self.pool = pool
        self.resource = resource

    async def __aenter__(self):
        """Enter async context - return the resource object."""
        return self.resource.obj

    async def __aexit__(self, exc_type, exc_val, exc_tb):
        """Exit async context - release resource back to pool."""
        await self.pool.release(self.resource)
