"""CPU profiling logic for Scalene."""

from __future__ import annotations

import math
from typing import TYPE_CHECKING, Callable

from scalene.runningstats import RunningStats
from scalene.scalene_funcutils import ScaleneFuncUtils
from scalene.scalene_statistics import (
    ByteCodeIndex,
    Filename,
    LineNumber,
    ScaleneStatistics,
)
from scalene.scalene_utility import _main_thread_id, add_stack, enter_function_meta
from scalene.time_info import TimeInfo

if TYPE_CHECKING:
    from types import FrameType


class ScaleneCPUProfiler:
    """Handles CPU profiling sample processing."""

    def __init__(
        self, stats: ScaleneStatistics, available_cpus: int, use_virtual_time: bool
    ) -> None:
        """Initialize the CPU profiler.

        Args:
            stats: The statistics object to update with CPU samples.
            available_cpus: Number of available CPUs for utilization calculations.
            use_virtual_time: Whether we're using virtual time (SIGVTALRM) or
                wall clock time (SIGALRM) for CPU sampling.
        """
        self._stats = stats
        self._available_cpus = available_cpus
        self._use_virtual_time = use_virtual_time

    def process_cpu_sample(
        self,
        new_frames: list[tuple[FrameType, int, FrameType]],
        now: TimeInfo,
        gpu_load: float,
        gpu_mem_used: float,
        prev: TimeInfo,
        is_thread_sleeping: dict[int, bool],
        should_trace: Callable[[Filename, str], bool],
        last_cpu_interval: float,
        stacks_enabled: bool,
    ) -> None:
        """Handle interrupts for CPU profiling.

        Args:
            new_frames: List of (frame, thread_id, original_frame) tuples.
            now: Current time information.
            gpu_load: Current GPU load (0.0-1.0).
            gpu_mem_used: Current GPU memory usage.
            prev: Previous time information.
            is_thread_sleeping: Dict mapping thread IDs to sleep status.
            should_trace: Function to check if a file/function should be traced.
            last_cpu_interval: The last CPU sampling interval.
            stacks_enabled: Whether stack collection is enabled.
        """
        if not new_frames:
            return

        elapsed = now - prev

        # Skip samples with negative values (can occur in multi-process settings)
        if any([elapsed.virtual < 0, elapsed.wallclock < 0, elapsed.user < 0]):
            return

        # Calculate CPU utilization
        cpu_utilization = 0.0
        if elapsed.wallclock != 0:
            cpu_utilization = elapsed.user / elapsed.wallclock

        core_utilization = cpu_utilization / self._available_cpus
        if cpu_utilization > 1.0:
            cpu_utilization = 1.0
            elapsed.wallclock = elapsed.user

        # Handle NaN GPU load
        if math.isnan(gpu_load):
            gpu_load = 0.0
        assert 0.0 <= gpu_load <= 1.0

        gpu_time = gpu_load * elapsed.wallclock
        self._stats.gpu_stats.total_gpu_samples += gpu_time

        # Compute Python vs C (native) time attribution.
        #
        # The interval-based formula assumes python_time ≈ last_cpu_interval
        # and c_time is the "excess" CPU time beyond that (from deferred signals
        # during C calls). This works when signals are deferred, but when the
        # signal happens to fire while Python is running, c_time ≈ 0 even if
        # significant C time occurred earlier in the interval.
        #
        # In wall clock mode, we augment this with instruction-based detection:
        # if we're at a CALL instruction (just returned from C), we attribute
        # ALL time to native since the signal was likely deferred during the call.
        python_time = last_cpu_interval
        c_time = max(elapsed.virtual - python_time, 0)
        total_time = python_time + c_time

        # Count non-sleeping frames
        total_frames = sum(
            not is_thread_sleeping[tident] for frame, tident, orig_frame in new_frames
        )
        if total_frames == 0:
            total_frames = 1

        normalized_time = total_time / total_frames
        average_python_time = python_time / total_frames
        average_c_time = c_time / total_frames
        average_cpu_time = (python_time + c_time) / total_frames

        # Process main thread
        main_thread_frame = new_frames[0][0]

        if stacks_enabled:
            add_stack(
                main_thread_frame,
                should_trace,
                self._stats.stacks,
                average_python_time,
                average_c_time,
                average_cpu_time,
            )

        enter_function_meta(main_thread_frame, should_trace, self._stats)
        fname = Filename(main_thread_frame.f_code.co_filename)
        lineno = (
            LineNumber(main_thread_frame.f_lineno)
            if main_thread_frame.f_lineno is not None
            else LineNumber(main_thread_frame.f_code.co_firstlineno)
        )

        if not is_thread_sleeping[_main_thread_id]:
            # Check for loop-top sampling bias.  CPython checks eval_breaker
            # at JUMP_BACKWARD, so signals are preferentially delivered at the
            # first body line of tight loops.  When detected, redistribute the
            # sample evenly across all loop lines.
            loop_lines = ScaleneFuncUtils.get_loop_body_lines(
                main_thread_frame.f_code, lineno
            )
            if loop_lines is not None and len(loop_lines) > 1:
                n = len(loop_lines)
                for ll in loop_lines:
                    self._update_main_thread_stats(
                        fname,
                        LineNumber(ll),
                        now,
                        average_python_time / n,
                        average_c_time / n,
                        average_cpu_time / n,
                        cpu_utilization,
                        core_utilization,
                        gpu_load,
                        gpu_mem_used,
                        elapsed,
                        gpu_weight=1.0 / n,
                    )
            else:
                # Determine Python vs C time attribution.
                #
                # The interval-based formula gives us:
                #   python_time = last_cpu_interval (timer interval)
                #   c_time = elapsed.virtual - python_time (excess from deferral)
                #
                # When c_time > 0, the signal was deferred during a C call, so
                # attribute C time to the preceding CALL line (if found).
                #
                # In wall clock mode, we also check if we're currently AT a CALL
                # instruction. If so, the signal was likely deferred during that
                # call, and ALL time (including python_time) should go to native.
                # This handles the case where elapsed.virtual ≈ last_cpu_interval
                # (c_time ≈ 0) but we're still returning from C code.
                is_at_call = ScaleneFuncUtils.is_call_function(
                    main_thread_frame.f_code,
                    ByteCodeIndex(main_thread_frame.f_lasti),
                )

                if not self._use_virtual_time and is_at_call:
                    # Wall clock mode at a CALL instruction: attribute all time
                    # to native on this line. The signal was deferred during the
                    # C call, so even python_time was spent in C, not Python.
                    self._update_main_thread_stats(
                        fname,
                        lineno,
                        now,
                        0.0,
                        average_cpu_time,
                        average_cpu_time,
                        cpu_utilization,
                        core_utilization,
                        gpu_load,
                        gpu_mem_used,
                        elapsed,
                    )
                elif average_c_time > 0:
                    # Signal was deferred (c_time > 0). Find the preceding CALL
                    # line to attribute C time correctly.
                    preceding_call_line = ScaleneFuncUtils.find_preceding_call_line(
                        main_thread_frame.f_code,
                        ByteCodeIndex(main_thread_frame.f_lasti),
                    )
                    if (
                        preceding_call_line is not None
                        and LineNumber(preceding_call_line) != lineno
                    ):
                        # Split: C time on the CALL line, Python time on f_lineno.
                        self._update_main_thread_stats(
                            fname,
                            LineNumber(preceding_call_line),
                            now,
                            0.0,
                            average_c_time,
                            average_c_time,
                            cpu_utilization,
                            core_utilization,
                            gpu_load,
                            gpu_mem_used,
                            elapsed,
                        )
                        self._update_main_thread_stats(
                            fname,
                            lineno,
                            now,
                            average_python_time,
                            0.0,
                            average_python_time,
                            cpu_utilization,
                            core_utilization,
                            gpu_load,
                            gpu_mem_used,
                            elapsed,
                        )
                    else:
                        # CALL is on the same line or not found; keep together.
                        self._update_main_thread_stats(
                            fname,
                            lineno,
                            now,
                            average_python_time,
                            average_c_time,
                            average_cpu_time,
                            cpu_utilization,
                            core_utilization,
                            gpu_load,
                            gpu_mem_used,
                            elapsed,
                        )
                else:
                    # No C time detected; attribute as computed.
                    self._update_main_thread_stats(
                        fname,
                        lineno,
                        now,
                        average_python_time,
                        average_c_time,
                        average_cpu_time,
                        cpu_utilization,
                        core_utilization,
                        gpu_load,
                        gpu_mem_used,
                        elapsed,
                    )

        # Process other threads
        for frame, tident, orig_frame in new_frames:
            if frame == main_thread_frame:
                continue

            add_stack(
                frame,
                should_trace,
                self._stats.stacks,
                average_python_time,
                average_c_time,
                average_cpu_time,
            )

            fname = Filename(frame.f_code.co_filename)
            lineno = (
                LineNumber(frame.f_lineno)
                if frame.f_lineno is not None
                else LineNumber(frame.f_code.co_firstlineno)
            )
            enter_function_meta(frame, should_trace, self._stats)

            if is_thread_sleeping[tident]:
                continue

            self._update_thread_stats(
                fname,
                lineno,
                orig_frame,
                normalized_time,
                cpu_utilization,
                core_utilization,
            )

        # Cleanup
        del new_frames[:]
        del new_frames
        del is_thread_sleeping
        self._stats.cpu_stats.total_cpu_samples += total_time

    def _update_main_thread_stats(
        self,
        fname: Filename,
        lineno: LineNumber,
        now: TimeInfo,
        average_python_time: float,
        average_c_time: float,
        average_cpu_time: float,
        cpu_utilization: float,
        core_utilization: float,
        gpu_load: float,
        gpu_mem_used: float,
        elapsed: TimeInfo,
        gpu_weight: float = 1.0,
    ) -> None:
        """Update statistics for the main thread."""
        cpu_stats = self._stats.cpu_stats
        gpu_stats = self._stats.gpu_stats

        cpu_stats.cpu_samples_list[fname][lineno].append(now.wallclock)
        cpu_stats.cpu_samples_python[fname][lineno] += average_python_time
        cpu_stats.cpu_samples_c[fname][lineno] += average_c_time
        cpu_stats.cpu_samples[fname] += average_cpu_time
        cpu_stats.cpu_utilization[fname][lineno].push(cpu_utilization)
        cpu_stats.core_utilization[fname][lineno].push(core_utilization)

        gpu_stats.gpu_samples[fname][lineno] += (
            gpu_load * elapsed.wallclock * gpu_weight
        )
        gpu_stats.n_gpu_samples[fname][lineno] += elapsed.wallclock * gpu_weight
        gpu_stats.gpu_mem_samples[fname][lineno].push(gpu_mem_used)

    def _update_thread_stats(
        self,
        fname: Filename,
        lineno: LineNumber,
        orig_frame: FrameType,
        normalized_time: float,
        cpu_utilization: float,
        core_utilization: float,
    ) -> None:
        """Update statistics for non-main threads."""
        cpu_stats = self._stats.cpu_stats

        # Check if the original caller is stuck inside a call
        if ScaleneFuncUtils.is_call_function(
            orig_frame.f_code,
            ByteCodeIndex(orig_frame.f_lasti),
        ):
            # Attribute time to native
            cpu_stats.cpu_samples_c[fname][lineno] += normalized_time
        else:
            # Attribute time to Python
            cpu_stats.cpu_samples_python[fname][lineno] += normalized_time

        cpu_stats.cpu_samples[fname] += normalized_time
        cpu_stats.cpu_utilization[fname][lineno].push(cpu_utilization)
        cpu_stats.core_utilization[fname][lineno].push(core_utilization)
