import cProfile
import pstats
import io
import os
import logging
import time
from typing import Callable, Optional
import flamegraph

# Configure logging
logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')

class PerformanceProfiler:
    """
    A class for profiling Python code to identify performance bottlenecks.
    """

    def __init__(self, output_dir: str = "profiles"):
        """
        Initializes the PerformanceProfiler.

        Args:
            output_dir (str): The directory to store profiling results. Defaults to "profiles".
        """
        self.output_dir = output_dir
        os.makedirs(self.output_dir, exist_ok=True)
        self.profiler: Optional[cProfile.Profile] = None

    def start(self) -> None:
        """
        Starts the profiler.
        """
        self.profiler = cProfile.Profile()
        self.profiler.enable()
        logging.info("Profiler started.")

    def stop(self, filename_prefix: str = "profile") -> None:
        """
        Stops the profiler and saves the results.

        Args:
            filename_prefix (str): The prefix for the output files. Defaults to "profile".
        """
        if self.profiler is None:
            logging.warning("Profiler was not started.")
            return

        self.profiler.disable()
        stats_filename = os.path.join(self.output_dir, f"{filename_prefix}.stats")
        self.flamegraph_filename = os.path.join(self.output_dir, f"{filename_prefix}.html")

        try:
            self.save_stats(stats_filename)
            self.generate_flamegraph(stats_filename, self.flamegraph_filename)
            logging.info(f"Profiler stopped. Stats saved to {stats_filename}, Flamegraph saved to {self.flamegraph_filename}")

        except Exception as e:
            logging.error(f"Error saving profiling results: {e}")

        finally:
            self.profiler = None

    def run(self, func: Callable, filename_prefix: str = "profile") -> None:
        """
        Runs the profiler on a given function.

        Args:
            func (Callable): The function to profile.
            filename_prefix (str): The prefix for the output files. Defaults to "profile".
        """
        self.start()
        try:
            func()
        except Exception as e:
            logging.error(f"Error during profiled function execution: {e}")
        finally:
            self.stop(filename_prefix)

    def save_stats(self, filename: str) -> None:
        """
        Saves the profiling statistics to a file.

        Args:
            filename (str): The name of the file to save the statistics to.
        """
        if self.profiler is None:
            logging.warning("Profiler was not started.")
            return

        try:
            with open(filename, "wb") as f:
                stats = pstats.Stats(self.profiler, stream=f)
                stats.sort_stats("tottime")
                stats.dump_stats(filename)
            logging.info(f"Profiling stats saved to {filename}")
        except Exception as e:
            logging.error(f"Error saving stats: {e}")

    def generate_flamegraph(self, stats_file: str, output_file: str) -> None:
        """
        Generates a flamegraph from the profiling statistics.

        Args:
            stats_file (str): The path to the profiling statistics file.
            output_file (str): The path to save the flamegraph HTML file.
        """
        try:
            flamegraph.FlameGraph().create(stats_file, output_file)
            logging.info(f"Flamegraph generated at {output_file}")
        except Exception as e:
            logging.error(f"Error generating flamegraph: {e}")

    def profile_memory(self, func: Callable, filename_prefix: str = "memory_profile"):
        """
        Placeholder for memory profiling functionality.  This would ideally use a library
        like `memory_profiler` to track memory usage.  For demonstration, it logs a message.

        Args:
            func (Callable): The function to profile.
            filename_prefix (str): The prefix for the output files.
        """
        logging.info(f"Memory profiling for {func.__name__} requested.  Memory profiling implementation is a placeholder.")
        # In a real implementation, this would use memory_profiler or similar to
        # track memory usage during the function execution and save the results.
        # Example (requires memory_profiler):
        # from memory_profiler import profile
        # @profile(filename=os.path.join(self.output_dir, f"{filename_prefix}.dat"))
        # def wrapped_func():
        #     func()
        # wrapped_func()
        func() # Execute the function anyway
        logging.info(f"Memory profiling for {func.__name__} complete (placeholder).  Results would be in {self.output_dir}/{filename_prefix}.dat (if memory_profiler was properly integrated).")


if __name__ == '__main__':
    # Example Usage:
    def example_function():
        """
        A simple function to demonstrate the profiler.
        """
        time.sleep(0.1)
        for i in range(100000):
            _ = i * i

    profiler = PerformanceProfiler()
    profiler.run(example_function, filename_prefix="example")
    profiler.profile_memory(example_function, filename_prefix="example_memory")

    print(f"Profiling results saved to {profiler.output_dir}")