import time
import tracemalloc
import logging
import functools
import platform
import sys
import inspect
import json
from datetime import datetime
from typing import Optional, Callable, Any, TypeVar, Dict, List

F = TypeVar('F', bound=Callable[..., Any])


class AuditConfig:
    def __init__(self) -> None:
        self.log_file: str = "audit.log"
        self.console_output: bool = True
        self.json_format: bool = False
        self._initialized: bool = False
        self._tracking_depth: int = 0  # Reference counter for nesting


_config = AuditConfig()


def format_size(bytes_size: float) -> str:
    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"


def _safe_start_tracing():
    """Starts tracemalloc if not already tracing; increments depth."""
    if _config._tracking_depth == 0 and not tracemalloc.is_tracing():
        tracemalloc.start()
    _config._tracking_depth += 1


def _safe_stop_tracing() -> tuple:
    """Decrements depth; only stops tracemalloc at the outermost level."""
    _config._tracking_depth -= 1
    current, peak = (0, 0)
    if tracemalloc.is_tracing():
        current, peak = tracemalloc.get_traced_memory()
        if _config._tracking_depth <= 0:
            tracemalloc.stop()
            _config._tracking_depth = 0
    return current, peak


def _log_event(data: Dict[str, Any], is_warning: bool):
    """Centralized logging for both Text and JSON formats."""
    logger = logging.getLogger("ezaudit")

    if _config.json_format:
        # Add timestamp to JSON data
        data['timestamp'] = datetime.now().isoformat()
        msg = json.dumps(data)
    else:
        # Standard Text Format
        prefix = "Func" if "name" in data else "Block"
        name = data.get("name") or data.get("block")
        msg = (f"{prefix}: {name} | Status: {data['status']} | "
               f"Duration: {data['duration']:.4f}s | Peak: {data['peak']}")
        if data.get('args'):
            msg += f" | Args: {data['args']}"
        if data.get('error'):
            msg += f" | Error: {data['error']}"

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

    if _config.console_output:
        print(final_msg)


def configure(log_file: Optional[str] = None,
              console_output: Optional[bool] = None,
              json_format: Optional[bool] = None) -> 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
    if json_format is not None:
        _config.json_format = json_format

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

    if logger.hasHandlers():
        for h in logger.handlers[:]:
            h.close()
            logger.removeHandler(h)

    try:
        handler = logging.FileHandler(_config.log_file, encoding='utf-8')
        handler.setFormatter(logging.Formatter('%(asctime)s | %(message)s'))
        logger.addHandler(handler)
    except Exception as e:
        print(f"[EZ-AUDIT ERROR] {e}")

    if not _config._initialized:
        logger.info(
            f"--- SESSION START | OS: {platform.system()} | Py: {sys.version.split()[0]} ---")
        _config._initialized = True


class audit_block:
    """Context manager to audit a specific block of code."""

    def __init__(self, name: str) -> None:
        self.name = name
        self.start_time: float = 0.0

    def __enter__(self) -> 'audit_block':
        if not logging.getLogger("ezaudit").hasHandlers():
            configure()
        _safe_start_tracing()
        self.start_time = time.perf_counter()
        return self

    def __exit__(self, exc_type, exc_val, exc_tb) -> None:
        _, peak = _safe_stop_tracing()
        dur = time.perf_counter() - self.start_time
        status = "SUCCESS" if exc_type is None else "FAILED"
        _log_event({
            "block": self.name, "status": status, "duration": dur,
            "peak": format_size(peak), "error": type(exc_val).__name__ if exc_val else ""
        }, False)


def monitor(threshold_sec: Optional[float] = None, lite: bool = False) -> Any:
    """Decorator to monitor performance of sync or async functions."""
    def decorator(func: F) -> F:
        @functools.wraps(func)
        async def async_wrapper(*args, **kwargs):
            if not logging.getLogger("ezaudit").hasHandlers():
                configure()
            if not lite:
                _safe_start_tracing()
            start = time.perf_counter()
            status, error = "SUCCESS", ""
            try:
                return await func(*args, **kwargs)
            except Exception as e:
                status, error = "FAILED", type(e).__name__
                raise e
            finally:
                _, peak = (0, 0) if lite else _safe_stop_tracing()
                dur = time.perf_counter() - start
                _log_event({
                    "name": func.__name__, "status": status, "duration": dur,
                    "peak": format_size(peak), "args": str(args)[:50], "error": error
                }, threshold_sec is not None and dur > threshold_sec)

        @functools.wraps(func)
        def sync_wrapper(*args, **kwargs):
            if not logging.getLogger("ezaudit").hasHandlers():
                configure()
            if not lite:
                _safe_start_tracing()
            start = time.perf_counter()
            status, error = "SUCCESS", ""
            try:
                return func(*args, **kwargs)
            except Exception as e:
                status, error = "FAILED", type(e).__name__
                raise e
            finally:
                _, peak = (0, 0) if lite else _safe_stop_tracing()
                dur = time.perf_counter() - start
                _log_event({
                    "name": func.__name__, "status": status, "duration": dur,
                    "peak": format_size(peak), "args": str(args)[:50], "error": error
                }, threshold_sec is not None and dur > threshold_sec)

        # type: ignore
        return async_wrapper if inspect.iscoroutinefunction(func) else sync_wrapper

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


def generate_report():
    """Generates a summary from the log file."""
    import os
    if not os.path.exists(_config.log_file):
        print(f"Log file {_config.log_file} not found.")
        return
    print(f"\n{'='*65}\n{'🚀 EZ-AUDIT v1.0.0 SUMMARY':^65}\n{'='*65}")
    with open(_config.log_file, "r", encoding="utf-8") as f:
        for line in f:
            # Support text and JSON
            if "Func:" in line or "Block:" in line or "{" in line:
                print(line.strip().split(" | ", 1)[-1])
    print('='*65 + "\n")


def _cli_main():
    """CLI entry point."""
    if "--report" in sys.argv:
        generate_report()
    else:
        print("EZ-Audit CLI: Use 'ezaudit --report'")
