"""Concurrent execution utilities for optimized performance."""

from __future__ import annotations

import asyncio
import concurrent.futures
import gc
import time
from collections.abc import Awaitable, Callable
from typing import Any, TypeVar

import psutil
from typing_extensions import Self

from pytola.core.logger import get_logger

T = TypeVar("T")

logger = get_logger(__name__)


class MemoryMonitor:
    """Memory usage monitoring and optimization."""

    def __init__(self, threshold_mb: int = 1000):
        """Initialize memory monitor.

        Args:
            threshold_mb: Memory threshold in MB to trigger cleanup
        """
        self.threshold_mb = threshold_mb
        self._initial_memory = self._get_memory_usage()
        self._peak_memory = self._initial_memory

    def _get_memory_usage(self) -> float:
        """Get current memory usage in MB."""
        process = psutil.Process()
        return process.memory_info().rss / 1024 / 1024

    def check_memory_usage(self) -> dict[str, float]:
        """Check current memory usage and statistics.

        Returns
        -------
            Memory usage statistics
        """
        current = self._get_memory_usage()
        self._peak_memory = max(self._peak_memory, current)

        return {
            "current_mb": current,
            "peak_mb": self._peak_memory,
            "initial_mb": self._initial_memory,
            "growth_mb": current - self._initial_memory,
        }

    def should_trigger_cleanup(self) -> bool:
        """Check if memory cleanup should be triggered."""
        stats = self.check_memory_usage()
        return stats["current_mb"] > self.threshold_mb

    def force_cleanup(self) -> None:
        """Force garbage collection and memory cleanup."""
        logger.debug("Forcing memory cleanup...")
        gc.collect()
        # Clear Python's internal caches
        import sys

        sys._clear_type_cache()

        stats = self.check_memory_usage()
        logger.debug(f"Memory after cleanup: {stats['current_mb']:.1f}MB")

    def log_memory_stats(self) -> None:
        """Log current memory statistics."""
        stats = self.check_memory_usage()
        logger.info(
            f"Memory usage - Current: {stats['current_mb']:.1f}MB, "
            f"Peak: {stats['peak_mb']:.1f}MB, "
            f"Growth: {stats['growth_mb']:.1f}MB"
        )


class PerformanceOptimizer:
    """Optimized concurrent execution manager."""

    def __init__(self, max_workers: int = 4, use_processes: bool = False):
        """Initialize performance optimizer.

        Args:
            max_workers: Maximum concurrent workers
            use_processes: Whether to use process pool instead of thread pool
        """
        self.max_workers = max_workers
        self.use_processes = use_processes
        self._executor = None
        self._loop = None
        self._memory_monitor = MemoryMonitor()

    def __enter__(self) -> Self:
        """Context manager entry."""
        self._loop = asyncio.get_event_loop()
        if self.use_processes:
            self._executor = concurrent.futures.ProcessPoolExecutor(max_workers=self.max_workers)
        else:
            self._executor = concurrent.futures.ThreadPoolExecutor(
                max_workers=self.max_workers, thread_name_prefix="pypack-worker"
            )
        return self

    def __exit__(self, exc_type, exc_val, exc_tb) -> None:
        """Context manager exit."""
        if self._executor:
            self._executor.shutdown(wait=True)
        # Log final memory stats
        self._memory_monitor.log_memory_stats()

    async def __aenter__(self) -> Self:
        """Async context manager entry."""
        return self.__enter__()

    async def __aexit__(self, exc_type, exc_val, exc_tb) -> None:
        """Async context manager exit."""
        return self.__exit__(exc_type, exc_val, exc_tb)

    async def run_cpu_bound_tasks(self, tasks: list[Callable[[], T]]) -> list[T]:
        """Run CPU-bound tasks using process pool for better performance.

        Args:
            tasks: List of CPU-intensive functions to execute

        Returns
        -------
            Results from all tasks
        """
        if not tasks:
            return []

        start_time = time.perf_counter()
        logger.debug(f"Running {len(tasks)} CPU-bound tasks with {self.max_workers} workers")

        # Use process pool for CPU-bound tasks
        loop = asyncio.get_event_loop()
        with concurrent.futures.ProcessPoolExecutor(max_workers=self.max_workers) as executor:
            futures = [loop.run_in_executor(executor, task) for task in tasks]
            results = await asyncio.gather(*futures, return_exceptions=True)

        elapsed = time.perf_counter() - start_time
        logger.debug(f"CPU-bound tasks completed in {elapsed:.3f}s")

        # Handle exceptions
        processed_results = []
        for i, result in enumerate(results):
            if isinstance(result, Exception):
                logger.warning(f"Task {i} failed: {result}")
                processed_results.append(None)
            else:
                processed_results.append(result)

        return processed_results

    async def run_io_bound_tasks(self, tasks: list[Callable[[], Awaitable[T]]]) -> list[T]:
        """Run I/O-bound tasks using asyncio for optimal performance.

        Args:
            tasks: List of async I/O functions to execute

        Returns
        -------
            Results from all tasks
        """
        if not tasks:
            return []

        start_time = time.perf_counter()
        logger.debug(f"Running {len(tasks)} I/O-bound tasks with asyncio")

        # Use asyncio.gather for I/O-bound tasks
        results = await asyncio.gather(*[task() for task in tasks], return_exceptions=True)

        elapsed = time.perf_counter() - start_time
        logger.debug(f"I/O-bound tasks completed in {elapsed:.3f}s")

        # Handle exceptions
        processed_results = []
        for i, result in enumerate(results):
            if isinstance(result, Exception):
                logger.warning(f"Task {i} failed: {result}")
                processed_results.append(None)
            else:
                processed_results.append(result)

        return processed_results

    async def run_mixed_workload(
        self,
        cpu_tasks: list[Callable[[], Any]],
        io_tasks: list[Callable[[], Awaitable[Any]]],
    ) -> tuple[list[Any], list[Any]]:
        """Run mixed CPU and I/O bound tasks optimally.

        Args:
            cpu_tasks: CPU-intensive synchronous tasks
            io_tasks: I/O-intensive async tasks

        Returns
        -------
            Tuple of (cpu_results, io_results)
        """
        start_time = time.perf_counter()

        # Run CPU and I/O tasks concurrently
        cpu_future = self.run_cpu_bound_tasks(cpu_tasks)
        io_future = self.run_io_bound_tasks(io_tasks)

        cpu_results, io_results = await asyncio.gather(cpu_future, io_future)

        elapsed = time.perf_counter() - start_time
        logger.info(f"Mixed workload completed in {elapsed:.3f}s with {self.max_workers} workers")

        return cpu_results, io_results

    def get_optimal_worker_count(self) -> int:
        """Get optimal worker count based on system resources.

        Returns
        -------
            Recommended number of workers
        """
        import multiprocessing

        import psutil

        # Get CPU count and available memory
        cpu_count = multiprocessing.cpu_count()
        available_memory = psutil.virtual_memory().available

        # Calculate based on memory requirements (assume 500MB per worker)
        memory_based_workers = max(1, int(available_memory / (500 * 1024 * 1024)))

        # Return minimum of CPU count and memory-based count
        optimal_count = min(cpu_count, memory_based_workers, 8)  # Cap at 8
        logger.debug(f"Optimal worker count: {optimal_count} (CPU: {cpu_count}, Memory-based: {memory_based_workers})")

        return optimal_count


# Performance monitoring utilities
class PerformanceMonitor:
    """Performance monitoring and profiling utilities."""

    def __init__(self):
        self.metrics = {}
        self._start_times = {}

    def start_timing(self, operation: str) -> None:
        """Start timing an operation."""
        self._start_times[operation] = time.perf_counter()

    def end_timing(self, operation: str) -> float:
        """End timing and record the duration."""
        if operation in self._start_times:
            duration = time.perf_counter() - self._start_times[operation]
            self.metrics[operation] = [*self.metrics.get(operation, []), duration]
            del self._start_times[operation]
            return duration
        return 0.0

    def get_average_time(self, operation: str) -> float:
        """Get average execution time for an operation."""
        times = self.metrics.get(operation, [])
        return sum(times) / len(times) if times else 0.0

    def get_stats(self) -> dict[str, dict[str, float]]:
        """Get performance statistics."""
        stats = {}
        for operation, times in self.metrics.items():
            if times:
                stats[operation] = {
                    "count": len(times),
                    "total": sum(times),
                    "average": sum(times) / len(times),
                    "min": min(times),
                    "max": max(times),
                }
        return stats

    def print_summary(self) -> None:
        """Print performance summary."""
        stats = self.get_stats()
        if not stats:
            logger.info("No performance metrics recorded")
            return

        logger.info("Performance Summary:")
        logger.info("-" * 50)
        for operation, data in stats.items():
            logger.info(f"{operation}:")
            logger.info(f"  Count: {data['count']}")
            logger.info(f"  Total: {data['total']:.3f}s")
            logger.info(f"  Average: {data['average']:.3f}s")
            logger.info(f"  Min: {data['min']:.3f}s")
            logger.info(f"  Max: {data['max']:.3f}s")
