#!/usr/bin/env python
# -*- coding: utf-8 -*-
#
# multitasking: Non-blocking Python methods using decorators
# https://github.com/ranaroussi/multitasking
#
# Copyright 2016-2025 Ran Aroussi
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
#

__version__ = "0.0.12"

# Core imports for multitasking functionality
import time as _time
from functools import wraps
from sys import exit as sysexit
from os import _exit as osexit
from typing import Any, Callable, Dict, List, Optional, TypedDict, Union

# Threading and multiprocessing imports
from threading import Thread, Semaphore
from multiprocessing import Process, cpu_count


class PoolConfig(TypedDict):
    """Type definition for execution pool configuration.

    This defines the structure of each pool in the POOLS dictionary,
    containing the semaphore for limiting concurrent tasks, the engine
    type (Thread or Process), pool name, and thread count.
    """
    pool: Optional[Semaphore]  # Controls concurrent task execution
    engine: Union[type[Thread], type[Process]]  # Execution engine
    name: str  # Human-readable pool identifier
    threads: int  # Maximum concurrent tasks (0 = unlimited)


class Config(TypedDict):
    """Type definition for global multitasking configuration.

    This structure holds all global state including CPU info, engine
    preferences, task tracking, and pool management. It serves as the
    central configuration store for the entire library.
    """
    CPU_CORES: int  # Number of CPU cores detected
    ENGINE: str  # Default engine type ("thread" or "process")
    MAX_THREADS: int  # Global maximum thread/process count
    DAEMON: bool  # Whether new tasks spawn as daemon threads/processes
    KILL_RECEIVED: bool  # Signal to stop accepting new tasks
    TASKS: List[Union[Thread, Process]]  # All created tasks
    POOLS: Dict[str, PoolConfig]  # Named execution pools
    POOL_NAME: str  # Currently active pool name


# Global configuration dictionary - this is the central state store
# for all multitasking operations. It tracks pools, tasks, and settings.
config: Config = {
    "CPU_CORES": cpu_count(),  # Auto-detect available CPU cores
    "ENGINE": "thread",  # Default to threading (safer than processes)
    "MAX_THREADS": cpu_count(),  # Start with one thread per CPU core
    # Default to non-daemon to preserve the historical contract that
    # Python waits for in-flight tasks to finish before exiting. Callers
    # that run long-lived workers and need the interpreter to exit
    # without joining them should opt in with `set_daemon(True)`.
    "DAEMON": False,
    "KILL_RECEIVED": False,  # Not in shutdown mode initially
    "TASKS": [],  # No tasks created yet
    "POOLS": {},  # No pools created yet
    "POOL_NAME": "Main"  # Default pool name
}


def set_max_threads(threads: Optional[int] = None) -> None:
    """Configure the maximum number of concurrent threads/processes.

    This function allows users to override the default CPU-based thread
    count. Setting this affects new pools but not existing ones.

    Args:
        threads: Maximum concurrent tasks. If None, uses CPU count.
                 Must be positive integer or None.

    Example:
        set_max_threads(4)  # Limit to 4 concurrent tasks
        set_max_threads()   # Reset to CPU count
    """
    if threads is not None:
        # User provided explicit thread count
        config["MAX_THREADS"] = threads
    else:
        # Reset to system default (one per CPU core)
        config["MAX_THREADS"] = cpu_count()


def set_engine(kind: str = "") -> None:
    """Configure the execution engine for new pools.

    This determines whether new tasks run in threads or separate
    processes. Threads share memory but processes are more isolated.
    Only affects pools created after this call.

    Args:
        kind: Engine type. Contains "process" for multiprocessing,
              anything else defaults to threading.

    Note:
        Threading: Faster startup, shared memory, GIL limitations
        Processing: Slower startup, isolated memory, true parallelism
    """
    if "process" in kind.lower():
        # Use multiprocessing for CPU-bound tasks
        config["ENGINE"] = "process"
    else:
        # Use threading for I/O-bound tasks (default)
        config["ENGINE"] = "thread"


def set_daemon(daemon: bool = False) -> None:
    """Configure whether new tasks spawn as daemon threads/processes.

    Daemon tasks are terminated automatically when the interpreter
    exits; non-daemon tasks block interpreter exit until they finish
    (the long-standing default, which lets ``wait_for_tasks`` or a
    plain script-end serve as an implicit join point).

    Call this once, alongside :func:`set_engine`, when your code
    contains long-lived ``@task`` workers (schedulers, pollers,
    infinite cleanup loops) that must NOT block process exit — for
    example in test suites, CLIs that shut down on SIGTERM, or any
    application where tasks are fire-and-forget by design.

    Args:
        daemon: True to spawn daemon tasks, False to preserve the
                historical non-daemon behavior (default).

    Note:
        This setting affects only tasks spawned after the call.
        Existing tasks retain whatever daemon flag they were created
        with. Also affects Processes as well as Threads; for
        Processes the daemon flag additionally prevents them from
        spawning child processes of their own.

    Example:
        import multitasking

        multitasking.set_engine("thread")
        multitasking.set_daemon(True)  # long-lived worker opt-in

        @multitasking.task
        def cleanup():
            while True:
                do_cleanup()
                time.sleep(60)

        cleanup()  # fires and forgets; exits cleanly with the process
    """
    config["DAEMON"] = bool(daemon)


def getPool(name: Optional[str] = None) -> Dict[str, Union[str, int]]:
    """Retrieve information about an execution pool.

    Returns a dictionary with pool metadata including engine type,
    name, and thread count. Useful for debugging and monitoring.

    Args:
        name: Pool name to query. If None, uses current active pool.

    Returns:
        Dictionary with keys: 'engine', 'name', 'threads'

    Raises:
        KeyError: If the specified pool doesn't exist
    """
    # Default to currently active pool if no name specified
    if name is None:
        name = config["POOL_NAME"]

    # Determine engine type from the pool configuration
    engine = "thread"  # Default assumption
    if config["POOLS"][config["POOL_NAME"]]["engine"] == Process:
        engine = "process"

    # Return pool information as a dictionary
    return {
        "engine": engine,
        "name": name,
        "threads": config["POOLS"][config["POOL_NAME"]]["threads"]
    }


def createPool(
    name: str = "main",
    threads: Optional[int] = None,
    engine: Optional[str] = None
) -> None:
    """Create a new execution pool with specified configuration.

    Pools manage concurrent task execution using semaphores. Each pool
    has its own thread/process limit and engine type. Creating a pool
    automatically makes it the active pool for new tasks.

    Args:
        name: Unique identifier for this pool
        threads: Max concurrent tasks. None uses global MAX_THREADS.
                 Values < 2 create unlimited pools (no semaphore).
        engine: "process" or "thread". None uses global ENGINE setting.

    Note:
        Setting threads=0 or threads=1 creates an unlimited pool where
        all tasks run immediately without queuing.
    """
    # Switch to this pool as the active one
    config["POOL_NAME"] = name

    # Parse and validate thread count
    try:
        threads = (
            int(threads) if threads is not None
            else config["MAX_THREADS"]
        )
    except (ValueError, TypeError):
        # Invalid input, fall back to global setting
        threads = config["MAX_THREADS"]

    # Thread counts less than 2 mean unlimited execution
    if threads < 2:
        threads = 0  # 0 is our internal code for "unlimited"

    # Determine engine type (default to global setting)
    engine = engine if engine is not None else config["ENGINE"]

    # Update global settings to match this pool
    config["MAX_THREADS"] = threads
    config["ENGINE"] = engine

    # Create the pool configuration
    config["POOLS"][config["POOL_NAME"]] = {
        # Semaphore controls concurrent execution (None = unlimited)
        "pool": Semaphore(threads) if threads > 0 else None,
        # Engine class determines Thread vs Process execution
        "engine": Process if "process" in engine.lower() else Thread,
        "name": name,
        "threads": threads
    }


def task(
    callee: Callable[..., Any]
) -> Callable[..., Optional[Union[Thread, Process]]]:
    """Decorator that converts a function into an asynchronous task.

    This is the main decorator of the library. It wraps any function
    to make it run asynchronously in the background using the current
    pool's configuration (threads or processes).

    The spawned Thread/Process inherits its ``daemon`` flag from the
    global :data:`config` via :func:`set_daemon` — by default tasks
    are non-daemon (Python waits for them on interpreter exit). For
    fire-and-forget background workers, call
    ``multitasking.set_daemon(True)`` once before declaring the task.

    Args:
        callee: The function to be made asynchronous

    Returns:
        Decorated function that returns Thread/Process object or None

    Example:
        @task
        def my_function(x, y):
            return x + y

        result = my_function(1, 2)  # Returns Thread/Process object
        wait_for_tasks()  # Wait for completion
    """
    # Ensure we have at least one pool available for task execution
    if not config["POOLS"]:
        createPool()  # Create default pool if none exists

    def _run_via_pool(*args: Any, **kwargs: Any) -> Any:
        """Internal wrapper that handles semaphore-controlled execution.

        This function is what actually runs in the background thread/process.
        It acquires the pool's semaphore (if any) before executing the
        original function, ensuring we don't exceed the concurrent limit.
        """
        pool = config["POOLS"][config["POOL_NAME"]]['pool']
        if pool is not None:
            # Limited pool: acquire semaphore before execution
            with pool:
                return callee(*args, **kwargs)
        else:
            # Unlimited pool: execute immediately
            return callee(*args, **kwargs)

    @wraps(callee)  # Preserve original function metadata
    def async_method(
        *args: Any, **kwargs: Any
    ) -> Optional[Union[Thread, Process]]:
        """The actual decorated function that users call.

        This decides whether to run synchronously (for 0-thread pools)
        or asynchronously (for normal pools). It handles the creation
        and startup of Thread/Process objects.
        """
        # Check if this pool runs synchronously (threads=0)
        if config["POOLS"][config["POOL_NAME"]]['threads'] == 0:
            # No threading: execute immediately and return None
            callee(*args, **kwargs)
            return None

        # Check if we're in shutdown mode
        if not config["KILL_RECEIVED"]:
            # Normal operation: create background task
            try:
                # Get the engine class (Thread or Process)
                engine_class = config["POOLS"][config["POOL_NAME"]]['engine']

                # Create the task. The daemon flag is pulled from the
                # global config (see `set_daemon`) rather than hardcoded
                # so callers can opt in to daemon semantics for
                # fire-and-forget workers without breaking the historical
                # "wait for tasks at exit" contract that non-daemon users
                # (e.g. yfinance) depend on.
                single = engine_class(
                    target=_run_via_pool,
                    args=args,
                    kwargs=kwargs,
                    daemon=config["DAEMON"],
                )
            except Exception:
                # Fallback for older Python versions without daemon param
                single = engine_class(
                    target=_run_via_pool,
                    args=args,
                    kwargs=kwargs
                )

            # Track this task for monitoring and cleanup
            config["TASKS"].append(single)

            # Start the task execution
            single.start()

            # Return the task object for user control
            return single

        # Shutdown mode: don't create new tasks
        return None

    return async_method


def get_list_of_tasks() -> List[Union[Thread, Process]]:
    """Retrieve all tasks ever created by this library.

    This includes both currently running tasks and completed ones.
    Useful for debugging and monitoring task history.

    Returns:
        List of all Thread/Process objects created by @task decorator

    Note:
        Completed tasks remain in this list until program termination.
        Use get_active_tasks() to see only currently running tasks.
    """
    return config["TASKS"]


def get_active_tasks() -> List[Union[Thread, Process]]:
    """Retrieve only the currently running tasks.

    Filters the complete task list to show only tasks that are still
    executing. This is more useful than get_list_of_tasks() for
    monitoring current system load.

    Returns:
        List of Thread/Process objects that are still running

    Example:
        active = get_active_tasks()
        print(f"Currently running {len(active)} tasks")
    """
    return [task for task in config["TASKS"] if task.is_alive()]


def wait_for_tasks(sleep: float = 0) -> bool:
    """Block until all background tasks complete execution.

    This is the primary synchronization mechanism. It prevents new
    tasks from being created and waits for existing ones to finish.
    Essential for ensuring all work is done before program exit.

    Args:
        sleep: Seconds to sleep between checks. 0 means busy-wait.
               Higher values reduce CPU usage but increase latency.

    Returns:
        Always returns True when all tasks are complete

    Note:
        Sets KILL_RECEIVED=True during execution to prevent new tasks,
        then resets it to False when done.
    """
    # Signal that we're in shutdown mode - no new tasks allowed
    config["KILL_RECEIVED"] = True

    # Handle synchronous pools (threads=0) - nothing to wait for
    if config["POOLS"][config["POOL_NAME"]]['threads'] == 0:
        return True

    try:
        # Main waiting loop
        while True:
            # Find all tasks that are still running
            running_tasks = [
                task for task in config["TASKS"]
                if task is not None and task.is_alive()
            ]

            # Attempt to join each running task with timeout
            # This gives each task a chance to complete cleanly
            for task in running_tasks:
                task.join(1)  # Wait up to 1 second per task

            # Recheck which tasks are still running after join attempts
            still_running = len([
                task for task in config["TASKS"]
                if task is not None and task.is_alive()
            ])

            # If no tasks are running, we're done
            if still_running == 0:
                break

            # Optional sleep to reduce CPU usage during waiting
            if sleep > 0:
                _time.sleep(sleep)

    except Exception:
        # Ignore any exceptions during cleanup (e.g., interrupted joins)
        pass

    # Re-enable task creation for future use
    config["KILL_RECEIVED"] = False
    return True


def killall(self: Any = None, cls: Any = None) -> None:
    """Emergency shutdown function that terminates the entire program.

    This is a last-resort function that immediately exits the program,
    potentially leaving tasks in an inconsistent state. It tries
    sys.exit() first, then os._exit() as a final measure.

    Args:
        self: Unused parameter kept for backward compatibility
        cls: Unused parameter kept for backward compatibility

    Warning:
        This function does NOT wait for tasks to complete cleanly.
        Use wait_for_tasks() for graceful shutdown instead.

    Note:
        The function attempts sys.exit(0) first (which allows cleanup
        handlers to run), falling back to os._exit(0) which terminates
        immediately without any cleanup.
    """
    # Signal shutdown mode to prevent new tasks
    config["KILL_RECEIVED"] = True

    try:
        # Attempt graceful exit (allows cleanup handlers)
        sysexit(0)
    except SystemExit:
        # Force immediate termination if graceful exit fails
        osexit(0)

    # This line should never be reached, but reset flag just in case
    config["KILL_RECEIVED"] = False