import time
import tracemalloc
import logging
import functools
import platform
import sys
from datetime import datetime
from typing import Optional, Callable, Any

# --- Helper Logic ---


def format_size(bytes_size: float) -> str:
    """Converts bytes to a human-readable string."""
    for unit in ['B', 'KB', 'MB', 'GB']:
        if bytes_size < 1024:
            return f"{bytes_size:.2f}{unit}"
        bytes_size /= 1024
    return f"{bytes_size:.2f}GB"


class AuditConfig:
    def __init__(self):
        self.log_file = "audit.log"
        self.console_output = True
        self._initialized = False


_config = AuditConfig()


def configure(log_file: Optional[str] = None, console_output: Optional[bool] = None):
    """Global configuration for EZ-Audit."""
    if log_file is not None:
        _config.log_file = log_file
    if console_output is not None:
        _config.console_output = console_output

    logger = logging.getLogger("ezaudit")
    logger.setLevel(logging.INFO)

    # IMPROVED: Close handlers before removing them
    if logger.hasHandlers():
        for handler in logger.handlers[:]:
            handler.close()  # <--- Essential for Windows
            logger.removeHandler(handler)

    file_handler = logging.FileHandler(_config.log_file)
    formatter = logging.Formatter('%(asctime)s | %(levelname)s | %(message)s')
    file_handler.setFormatter(formatter)
    logger.addHandler(file_handler)

    # Log Session Header once per session
    if not _config._initialized:
        header = f"--- AUDIT SESSION START | OS: {platform.system()} | Py: {sys.version.split()[0]} ---"
        logger.info(header)
        _config._initialized = True

# --- The Context Manager (New Feature!) ---


class audit_block:
    """
    Allows auditing a specific block of code using 'with':
    with audit_block("Database Query"):
        do_query()
    """

    def __init__(self, name: str):
        self.name = name

    def __enter__(self):
        tracemalloc.start()
        self.start_time = time.perf_counter()
        return self

    def __exit__(self, exc_type, exc_val, exc_tb):
        duration = time.perf_counter() - self.start_time
        _, peak = tracemalloc.get_traced_memory()
        tracemalloc.stop()

        status = "SUCCESS" if exc_type is None else "FAILED"
        msg = f"Block: {self.name} | Status: {status} | Duration: {duration:.4f}s | Peak: {format_size(peak)}"

        logger = logging.getLogger("ezaudit")
        if _config.console_output:
            print(f"[EZ-AUDIT] {msg}")
        logger.info(msg)

# --- The Main Decorator ---


def monitor(threshold_sec: Optional[float] = None):
    """Decorator to monitor function performance."""
    def decorator(func: Callable):
        @functools.wraps(func)
        def wrapper(*args, **kwargs):
            # Ensure logger is set up if user didn't call configure()
            if not logging.getLogger("ezaudit").hasHandlers():
                configure()

            tracemalloc.start()
            start_time = time.perf_counter()
            status = "SUCCESS"
            error_msg = ""

            try:
                return func(*args, **kwargs)
            except Exception as e:
                status = "FAILED"
                error_msg = f" | Error: {type(e).__name__}"
                raise e
            finally:
                duration = time.perf_counter() - start_time
                _, peak = tracemalloc.get_traced_memory()
                tracemalloc.stop()

                size_str = format_size(peak)
                audit_msg = f"Func: {func.__name__} | Status: {status} | Duration: {duration:.4f}s | Peak: {size_str}{error_msg}"

                logger = logging.getLogger("ezaudit")
                is_over = threshold_sec and duration > threshold_sec

                if is_over:
                    final_msg = f"⚠️ [THRESHOLD EXCEEDED] {audit_msg}"
                    logger.warning(final_msg)
                else:
                    final_msg = f"[EZ-AUDIT] {audit_msg}"
                    logger.info(audit_msg)

                if _config.console_output:
                    print(final_msg)

        return wrapper

    if callable(threshold_sec):
        f = threshold_sec
        threshold_sec = None
        return decorator(f)
    return decorator
