from pathlib import Path
import shutil
import time
import sys
import json
import importlib
import re
from typing import Optional, List, Dict, Any, Tuple, Set, Callable
from types import ModuleType
from concurrent.futures import ThreadPoolExecutor
import threading
import tracemalloc
import html
from datetime import datetime
from functools import wraps


class CallsTracker:
    """
    Advanced tracing and profiling tool for real-time Python program analysis.

    Features:
    - Function call tracing with timing and location information
    - Performance profiling with execution time and memory usage
    - Advanced filtering and customization options
    - Multiple output formats and reporting
    - Interactive controls and context-aware tracing
    """

    def __init__(self):
        """Initialize the advanced tracer with default configuration."""
        self.running = False
        self.log = []
        self.performance_data = {}
        self.call_count = {}
        self.filters = {
            "whitelist": set(),
            "blacklist": set(),
            "min_duration": 0,
            "max_duration": float("inf"),
            "ignore_system": True,
        }
        self.config = {
            "track_arguments": False,
            "track_return_values": False,
            "track_exceptions": True,
            "track_memory": True,
            "output_file": None,
            "log_level": "INFO",  # DEBUG, INFO, WARNING, ERROR
        }

        # Performance tracking
        self.start_times = {}
        self.memory_snapshots = {}

        # Thread safety
        self._lock = threading.Lock()

        # Start memory tracking if enabled
        if self.config["track_memory"]:
            tracemalloc.start()

    def trace_calls(self, frame, event, arg):
        """
        Main trace function called by sys.settrace.

        Args:
            frame: Current stack frame
            event: Trace event ('call', 'return', 'exception')
            arg: Event-specific argument

        Returns:
            Self reference for chaining traces
        """
        if event == "call":
            self._handle_call_event(frame, arg)
        elif event == "return":
            self._handle_return_event(frame, arg)
        elif event == "exception":
            self._handle_exception_event(frame, arg)

        return self.trace_calls

    def _handle_call_event(self, frame, arg):
        """Handle function call events."""
        code = frame.f_code
        func_name = code.co_name
        filename = code.co_filename

        # Apply filters
        if not self._should_trace(func_name, filename):
            return

        # Get call information
        lineno = frame.f_lineno
        call_id = f"{func_name}_{id(frame)}"

        # Record start time
        self.start_times[call_id] = time.time()

        # Track memory if enabled
        if self.config["track_memory"]:
            self.memory_snapshots[call_id] = tracemalloc.take_snapshot()

        # Build log entry
        entry = {
            "timestamp": time.time(),
            "time_str": time.strftime("%H:%M:%S"),
            "event": "call",
            "function": func_name,
            "filename": filename,
            "line_number": lineno,
            "call_id": call_id,
            "thread_id": threading.get_ident(),
        }

        # Track arguments if enabled
        if self.config["track_arguments"]:
            entry["arguments"] = self._get_function_arguments(frame)

        # Update call count
        self.call_count[func_name] = self.call_count.get(func_name, 0) + 1

        # Add to log and print
        self._add_log_entry(entry)

    def _handle_return_event(self, frame, arg):
        """Handle function return events."""
        code = frame.f_code
        func_name = code.co_name
        call_id = f"{func_name}_{id(frame)}"

        if call_id not in self.start_times:
            return

        # Calculate execution time
        start_time = self.start_times[call_id]
        execution_time = time.time() - start_time

        # Build return entry
        entry = {
            "timestamp": time.time(),
            "time_str": time.strftime("%H:%M:%S"),
            "event": "return",
            "function": func_name,
            "execution_time": execution_time,
            "call_id": call_id,
        }

        # Track return value if enabled
        if self.config["track_return_values"] and arg is not None:
            entry["return_value"] = str(arg)[:100]  # Limit length

        # Track memory usage if enabled
        if self.config["track_memory"] and call_id in self.memory_snapshots:
            current_snapshot = tracemalloc.take_snapshot()
            memory_usage = self._calculate_memory_usage(
                self.memory_snapshots[call_id], current_snapshot
            )
            entry["memory_usage"] = memory_usage
            del self.memory_snapshots[call_id]

        # Update performance data
        self._update_performance_data(
            func_name, execution_time, entry.get("memory_usage", 0)
        )

        # Add to log and print
        self._add_log_entry(entry)

        # Cleanup
        del self.start_times[call_id]

    def _handle_exception_event(self, frame, arg):
        """Handle exception events."""
        if not self.config["track_exceptions"]:
            return

        code = frame.f_code
        func_name = code.co_name

        entry = {
            "timestamp": time.time(),
            "time_str": time.strftime("%H:%M:%S"),
            "event": "exception",
            "function": func_name,
            "exception_type": type(arg[0]).__name__,
            "exception_message": str(arg[0]),
        }

        self._add_log_entry(entry, level="ERROR")

    def _should_trace(self, func_name: str, filename: str) -> bool:
        """
        Determine if a function should be traced based on filters.

        Args:
            func_name: Name of the function
            filename: File where function is defined

        Returns:
            Boolean indicating whether to trace this function
        """
        # Check blacklist
        if func_name in self.filters["blacklist"]:
            return False

        # Check whitelist (if not empty, only trace whitelisted functions)
        if self.filters["whitelist"] and func_name not in self.filters["whitelist"]:
            return False

        # Ignore system files if configured
        if self.filters["ignore_system"] and (
            "site-packages" in filename or "lib/python" in filename
        ):
            return False

        return True

    def _get_function_arguments(self, frame) -> Dict[str, Any]:
        """Extract function arguments from frame."""
        try:
            args = {}
            arg_info = inspect.getargvalues(frame)

            for arg in arg_info.args:
                if arg in frame.f_locals:
                    value = frame.f_locals[arg]
                    args[arg] = str(value)[:50]  # Limit length for display

            return args
        except:
            return {}

    def _calculate_memory_usage(self, start_snapshot, end_snapshot) -> int:
        """Calculate memory usage between two snapshots in bytes."""
        try:
            stats = end_snapshot.compare_to(start_snapshot, "lineno")
            return sum(stat.size_diff for stat in stats if stat.size_diff > 0)
        except:
            return 0

    def _update_performance_data(
        self, func_name: str, execution_time: float, memory_usage: int
    ):
        """Update performance statistics for a function."""
        with self._lock:
            if func_name not in self.performance_data:
                self.performance_data[func_name] = {
                    "total_time": 0,
                    "call_count": 0,
                    "avg_time": 0,
                    "max_time": 0,
                    "min_time": float("inf"),
                    "total_memory": 0,
                    "avg_memory": 0,
                }

            data = self.performance_data[func_name]
            data["total_time"] += execution_time
            data["call_count"] += 1
            data["avg_time"] = data["total_time"] / data["call_count"]
            data["max_time"] = max(data["max_time"], execution_time)
            data["min_time"] = min(data["min_time"], execution_time)
            data["total_memory"] += memory_usage
            data["avg_memory"] = data["total_memory"] / data["call_count"]

    def _add_log_entry(self, entry: Dict[str, Any], level: str = "INFO"):
        """
        Add entry to log and print based on log level.

        Args:
            entry: Log entry dictionary
            level: Log level (DEBUG, INFO, WARNING, ERROR)
        """
        entry["level"] = level
        self.log.append(entry)

        # Print based on log level configuration
        log_levels = {"DEBUG": 0, "INFO": 1, "WARNING": 2, "ERROR": 3}
        current_level = log_levels.get(self.config["log_level"], 1)
        entry_level = log_levels.get(level, 1)

        if entry_level >= current_level:
            print(self._format_log_entry(entry))

    def _format_log_entry(self, entry: Dict[str, Any]) -> str:
        """Format log entry for console output."""
        time_str = entry["time_str"]
        func_name = entry["function"]
        event = entry["event"]

        if event == "call":
            return f"[{time_str}] CALL {func_name}() at {entry['filename']}:{entry['line_number']}"
        elif event == "return":
            time_ms = entry["execution_time"] * 1000
            return f"[{time_str}] RETURN {func_name}() - {time_ms:.2f}ms"
        elif event == "exception":
            return f"[{time_str}] ERROR {func_name}() - {entry['exception_type']}: {entry['exception_message']}"

        return f"[{time_str}] {event.upper()} {func_name}()"

    # Public API Methods

    def start(self):
        """Start the tracing system."""
        if not self.running:
            self.running = True
            sys.settrace(self.trace_calls)
            print("\n[CallsTracker] Live trace started with advanced profiling...\n")

    def stop(self):
        """Stop the tracing system."""
        if self.running:
            sys.settrace(None)
            self.running = False

            # Stop memory tracking
            if self.config["track_memory"]:
                tracemalloc.stop()

            print("\n[CallsTracker] Live trace stopped.\n")

    def pause(self):
        """Pause tracing temporarily."""
        if self.running:
            sys.settrace(None)
            print("\n[CallsTracker] Tracing paused.\n")

    def resume(self):
        """Resume tracing after pause."""
        if self.running:
            sys.settrace(self.trace_calls)
            print("\n[CallsTracker] Tracing resumed.\n")

    # Filter Configuration Methods

    def add_to_whitelist(self, function_names: List[str]):
        """Add functions to whitelist for tracing."""
        self.filters["whitelist"].update(function_names)

    def add_to_blacklist(self, function_names: List[str]):
        """Add functions to blacklist to exclude from tracing."""
        self.filters["blacklist"].update(function_names)

    def set_duration_filter(
        self, min_duration: float = 0, max_duration: float = float("inf")
    ):
        """Set minimum and maximum duration filters in seconds."""
        self.filters["min_duration"] = min_duration
        self.filters["max_duration"] = max_duration

    def ignore_system_files(self, ignore: bool = True):
        """Configure whether to ignore system library files."""
        self.filters["ignore_system"] = ignore

    # Configuration Methods

    def configure(self, **kwargs):
        """Update tracer configuration."""
        for key, value in kwargs.items():
            if key in self.config:
                self.config[key] = value

        # Restart memory tracking if configuration changed
        if "track_memory" in kwargs:
            if kwargs["track_memory"] and not tracemalloc.is_tracing():
                tracemalloc.start()
            elif not kwargs["track_memory"] and tracemalloc.is_tracing():
                tracemalloc.stop()

    def set_output_file(self, filename: str):
        """Set output file for saving logs."""
        self.config["output_file"] = filename

    def set_log_level(self, level: str):
        """Set log level for output filtering."""
        valid_levels = ["DEBUG", "INFO", "WARNING", "ERROR"]
        if level in valid_levels:
            self.config["log_level"] = level

    # Analysis and Reporting Methods

    def get_performance_report(self) -> Dict[str, Any]:
        """Generate comprehensive performance report."""
        return {
            "total_functions_traced": len(self.performance_data),
            "total_calls": sum(self.call_count.values()),
            "performance_data": self.performance_data,
            "call_counts": self.call_count,
            "trace_duration": time.time()
            - (self.log[0]["timestamp"] if self.log else time.time()),
        }

    def find_bottlenecks(self, top_n: int = 10) -> List[Dict[str, Any]]:
        """Find top N performance bottlenecks."""
        bottlenecks = []
        for func_name, data in self.performance_data.items():
            if data["call_count"] > 0:
                bottlenecks.append(
                    {
                        "function": func_name,
                        "total_time": data["total_time"],
                        "avg_time": data["avg_time"],
                        "call_count": data["call_count"],
                        "memory_usage": data["avg_memory"],
                    }
                )

        return sorted(bottlenecks, key=lambda x: x["total_time"], reverse=True)[:top_n]

    def get_call_graph(self) -> Dict[str, Any]:
        """Generate call graph data for visualization."""
        # This is a simplified version - in practice you'd track call relationships
        call_graph = {"nodes": [], "edges": []}

        for func_name, data in self.performance_data.items():
            call_graph["nodes"].append(
                {
                    "id": func_name,
                    "call_count": data["call_count"],
                    "avg_time": data["avg_time"],
                }
            )

        return call_graph

    def export_logs(self, format: str = "json", filename: str = None):
        """Export logs to various formats."""
        filename = (
            filename
            or self.config["output_file"]
            or f'tracer_log_{datetime.now().strftime("%Y%m%d_%H%M%S")}'
        )

        if format == "json":
            with open(f"{filename}.json", "w") as f:
                json.dump(
                    {
                        "logs": self.log,
                        "performance_report": self.get_performance_report(),
                    },
                    f,
                    indent=2,
                )

        elif format == "html":
            self._export_html_report(filename)

        elif format == "text":
            with open(f"{filename}.txt", "w") as f:
                for entry in self.log:
                    f.write(self._format_log_entry(entry) + "\n")

        print(f"[CallsTracker] Logs exported to {filename}.{format}")

    def _export_html_report(self, filename: str):
        """Generate HTML report with charts and analysis."""
        html_content = f"""
        <!DOCTYPE html>
        <html>
        <head>
            <title>Calls Report</title>
            <script src="https://cdn.jsdelivr.net/npm/chart.js"></script>
            <style>
                body {{ font-family: Arial, sans-serif; margin: 20px; }}
                .section {{ margin-bottom: 30px; }}
                .log-entry {{ margin: 5px 0; padding: 5px; border-left: 3px solid #ccc; }}
                .call {{ border-color: #4CAF50; }}
                .return {{ border-color: #2196F3; }}
                .exception {{ border-color: #f44336; }}
            </style>
        </head>
        <body>
            <h1>Advanced Calls Report</h1>
            <div class="section">
                <h2>Performance Summary</h2>
                <div id="performanceChart">
                    <!-- Chart would be generated here -->
                </div>
            </div>
            <div class="section">
                <h2>Execution Log</h2>
                {self._generate_html_log()}
            </div>
        </body>
        </html>
        """

        with open(f"{filename}.html", "w") as f:
            f.write(html_content)

    def _generate_html_log(self) -> str:
        """Generate HTML formatted log entries."""
        html_entries = []
        for entry in self.log:
            class_name = entry["event"]
            formatted_entry = self._format_log_entry(entry)
            html_entries.append(
                f'<div class="log-entry {class_name}">{html.escape(formatted_entry)}</div>'
            )

        return "\n".join(html_entries)

    def clear_logs(self):
        """Clear all logs and reset counters."""
        self.log.clear()
        self.performance_data.clear()
        self.call_count.clear()
        self.start_times.clear()
        self.memory_snapshots.clear()
        print("[CallsTracker] All logs and counters cleared.")

    def get_statistics(self) -> Dict[str, Any]:
        """Get current tracing statistics."""
        return {
            "total_log_entries": len(self.log),
            "unique_functions_traced": len(self.performance_data),
            "total_function_calls": sum(self.call_count.values()),
            "current_status": "running" if self.running else "stopped",
            "memory_tracking_enabled": self.config["track_memory"],
        }


# Decorator for easy function tracing
def trace_function(tracer_instance: CallsTracker = None):
    """
    Decorator to trace specific functions without global tracing.

    Usage:
        @trace_function()
        def my_function():
            pass
    """

    def decorator(func):
        @wraps(func)
        def wrapper(*args, **kwargs):
            # Create tracer if not provided
            local_tracer = tracer_instance or CallsTracker()
            local_tracer.configure(track_arguments=True, track_return_values=True)

            # Trace this function call
            frame = inspect.currentframe()
            local_tracer._handle_call_event(frame, None)

            try:
                result = func(*args, **kwargs)
                local_tracer._handle_return_event(frame, result)
                return result
            except Exception as e:
                local_tracer._handle_exception_event(frame, (e, None))
                raise

        return wrapper

    return decorator
