from __future__ import annotations

import json
import time
from dataclasses import dataclass
from enum import Enum
from typing import Any, Dict, Union

from pyinstrument import processors
from pyinstrument.frame import Frame
from pyinstrument.renderers.base import FrameRenderer, ProcessorList
from pyinstrument.session import Session

# pyright: strict


@dataclass(frozen=True, eq=True)
class SpeedscopeFrame:
    """
    Data class to store data needed for speedscope's concept of a
    frame, hereafter referred to as a "speedscope frame", as opposed to
    a "pyinstrument frame". This type must be hashable in order to use
    it as a dictionary key; a dictionary will be used to track unique
    speedscope frames.
    """

    name: str | None
    file: str | None
    line: int | None


class SpeedscopeEventType(Enum):
    """Enum representing the only two types of speedscope frame events"""

    OPEN = "O"
    CLOSE = "C"


@dataclass
class SpeedscopeEvent:
    """
    Data class to store speedscope's concept of an "event", which
    corresponds to opening or closing stack frames as functions or
    methods are entered or exited.
    """

    type: SpeedscopeEventType
    at: float
    frame: int


@dataclass
class SpeedscopeProfile:
    """
    Data class to store speedscope's concept of a "profile".
    """

    name: str
    events: list[SpeedscopeEvent]
    end_value: float
    start_value: float = 0.0
    type: str = "evented"
    unit: str = "seconds"


@dataclass
class SpeedscopeFile:
    """
    Data class encoding fields in speedscope's JSON file schema
    """

    name: str
    profiles: list[SpeedscopeProfile]
    shared: dict[str, list[SpeedscopeFrame]]
    schema: str = "https://www.speedscope.app/file-format-schema.json"
    active_profile_index: None = None
    exporter: str = "pyinstrument"


SpeedscopeFrameDictType = Dict[str, Union[str, int, None]]
SpeedscopeEventDictType = Dict[str, Union[SpeedscopeEventType, float, int]]


class SpeedscopeEncoder(json.JSONEncoder):
    """
    Encoder class used by json.dumps to serialize the various
    speedscope data classes.
    """

    def default(self, o: Any) -> Any:
        if isinstance(o, SpeedscopeFile):
            return {
                "$schema": o.schema,
                "name": o.name,
                "activeProfileIndex": o.active_profile_index,
                "exporter": o.exporter,
                "profiles": o.profiles,
                "shared": o.shared,
            }
        if isinstance(o, SpeedscopeProfile):
            return {
                "type": o.type,
                "name": o.name,
                "unit": o.unit,
                "startValue": o.start_value,
                "endValue": o.end_value,
                "events": o.events,
            }
        if isinstance(o, (SpeedscopeFrame, SpeedscopeEvent)):
            d: SpeedscopeFrameDictType | SpeedscopeEventDictType = o.__dict__
            return d
        if isinstance(o, SpeedscopeEventType):
            return o.value
        return json.JSONEncoder.default(self, o)


class SpeedscopeRenderer(FrameRenderer):
    """
    Outputs a tree of JSON conforming to the speedscope schema documented at

    wiki: https://github.com/jlfwong/speedscope/wiki/Importing-from-custom-sources
    schema: https://www.speedscope.app/file-format-schema.json
    spec: https://github.com/jlfwong/speedscope/blob/main/src/lib/file-format-spec.ts
    example: https://github.com/jlfwong/speedscope/blob/main/sample/profiles/speedscope/0.0.1/simple.speedscope.json

    """

    output_file_extension = "speedscope.json"

    def __init__(self, **kwargs: Any):
        super().__init__(**kwargs)

        # Member holding a running total of wall clock time needed to
        # compute the times at which events occur
        self._event_time: float = 0.0

        # Map of speedscope frames to speedscope frame indices, needed
        # to construct evented speedscope profiles; exploits LIFO
        # property of popinfo method in Python 3.7+ dictionaries. This
        # dictionary is used to build up the "shared" JSON array in
        # speedscope's schema.
        self._frame_to_index: dict[SpeedscopeFrame, int] = {}

    def render_frame(self, frame: Frame | None) -> list[SpeedscopeEvent]:
        """
        Builds up a list of speedscope events that are used to populate the
        "events" array in speedscope-formatted JSON.

        This method has two notable side effects:

        * it populates the self._frame_to_index dictionary that matches
          speedscope frames with their positions in the "shared" array of
          speedscope output; this dictionary will be used to write this
          "shared" array in the render method

        * it accumulates a running total of time elapsed by
          accumulating the self_time spent in each pyinstrument frame;
          this running total is used by speedscope events to construct
          a flame chart.
        """

        # if frame is None, recursion bottoms out; no event frames
        # need to be added
        if frame is None:
            return []

        # Otherwise, form a speedscope frame and add it to the frame
        # to index map if the frame is not already a key in that map.
        sframe = SpeedscopeFrame(frame.function, frame.file_path, frame.line_no)
        if sframe not in self._frame_to_index:
            self._frame_to_index[sframe] = len(self._frame_to_index)

        # Get the frame index and add a speedscope event corresponding
        # to opening a stack frame.
        sframe_index = self._frame_to_index[sframe]
        open_event = SpeedscopeEvent(SpeedscopeEventType.OPEN, self._event_time, sframe_index)
        events_array: list[SpeedscopeEvent] = [open_event]

        # Add stack frame open and close events for all child frames
        # of this frame.
        for child in frame.children:
            events_array.extend(self.render_frame(child))

        # Update event time for closing this stack frame.
        #
        # If number of frames approaches 1e16 * desired accuracy
        # level, consider using Neumaier-Kahan summation; improves
        # worst-case relative accuracy of sum from O(num_summands *
        # eps) to (2 * eps + O(num_summands * eps * eps)), where eps
        # is IEEE-754 double precision unit roundoff, approximately
        # 1e-16. Average case relative accuracy expressions replace
        # num_summands with sqrt(num_summands). However, Kahan
        # summation quadruples operation count of sum, and Neumaier
        # variant also adds a branch & swap for each summand. Pairwise
        # summation isn't an option here because a running total is
        # needed.

        self._event_time += frame.absorbed_time
        if frame.is_synthetic_leaf:
            # only time contained within leaf nodes is real time i.e. not the sum of children
            self._event_time += frame.time

        # Add event closing this stack frame.
        close_event = SpeedscopeEvent(SpeedscopeEventType.CLOSE, self._event_time, sframe_index)
        events_array.append(close_event)

        return events_array

    def render(self, session: Session):
        frame = self.preprocess(session.root_frame())

        id_: str = time.strftime("%Y-%m-%dT%H-%M-%S", time.localtime(session.start_time))
        name: str = f"CPU profile for '{session.target_description}' at {id_}"

        sprofile_list: list[SpeedscopeProfile] = [
            SpeedscopeProfile(name, self.render_frame(frame), session.duration)
        ]

        # Exploits Python 3.7+ dictionary property of iterating over
        # keys in insertion order to build the list of speedscope
        # frames.
        sframe_list: list[SpeedscopeFrame] = [sframe for sframe in iter(self._frame_to_index)]

        shared_dict = {"frames": sframe_list}
        speedscope_file = SpeedscopeFile(name, sprofile_list, shared_dict)

        return "%s\n" % json.dumps(speedscope_file, cls=SpeedscopeEncoder)

    def default_processors(self) -> ProcessorList:
        """
        Default Processors for speedscope renderer; note that
        processors.aggregate_repeated_calls is removed because
        speedscope is a timeline-based format.
        """
        return [
            processors.remove_importlib,
            processors.remove_tracebackhide,
            processors.merge_consecutive_self_time,
            processors.remove_irrelevant_nodes,
            processors.remove_unnecessary_self_time_nodes,
            processors.remove_first_pyinstrument_frames_processor,
        ]