import functools
import logging
import operator
import os
import sys
import threading
import time
from collections import defaultdict

import msgpack
import ulid

from .db import save_trace_in_sqlite
from .filters.core import LibraryPathFilter
from .filters.kolo import kolo_filter_filename
from .git import COMMIT_SHA
from .plugins import PluginProcessor, load_plugin_data
from .serialize import dump_msgpack, dump_msgpack_lightweight_repr, user_code_call_site
from .version import __version__

logger = logging.getLogger("kolo")


def frozen_filter(filename):
    return "<frozen " in filename


def pypy_filter(filename):
    return "<builtin>/" in filename


def exec_filter(filename):
    return filename == "<string>"


def attrs_filter(filename):
    if filename:
        return filename.startswith("<attrs generated")
    else:  # pragma: no cover
        # Index of the frame in the Python stack:
        # 0: This function
        # 1: KoloMonitor.ignore method
        # 2: KoloMonitor.monitor_ method
        # 3: Attrs code with empty filename
        # 4: Attrs code that may match attr/_make.py
        # 5: Attrs code that may match attr/_make.py
        frame = sys._getframe(4)
        if frame is not None and frame.f_code.co_filename == "":
            frame = sys._getframe(5)
        return frame is not None and frame.f_code.co_filename.endswith(
            os.path.normpath("attr/_make.py")
        )


def pytest_filter(filename):
    return filename == "<pytest match expression>"


def build_frame_data(event, name, filename, frame, frame_id, call_site):
    thread = threading.current_thread()

    frame_locals = {k: v for k, v in frame.f_locals.items() if k != "__builtins__"}
    return {
        "path": frame_path(filename, frame.f_lineno),
        "co_name": name,
        "qualname": get_qualname(frame),
        "event": event,
        "frame_id": frame_id,
        "locals": frame_locals,
        "thread": thread.name,
        # The operating system's thread ID.
        # Guaranteed to be unique for the lifetime of the thread.
        "thread_native_id": thread.native_id,
        "timestamp": time.time(),
        "type": "frame",
        "user_code_call_site": call_site,
    }


def frame_path(path, lineno):
    try:
        relative_path = os.path.relpath(path)
    except ValueError:
        relative_path = path
    return f"{relative_path}:{lineno}"


def get_qualname(frame):
    qualname = frame.f_code.co_qualname
    module = frame.f_globals["__name__"]
    return f"{module}.{qualname}"


class KoloLocals(threading.local):
    def __init__(self):
        self.call_frames = []
        self._frame_ids = {}


if sys.version_info >= (3, 12):

    class KoloMonitor:
        def __init__(self, db_path, *, config=None, one_trace_per_test=False, source):
            self.tool_id = sys.monitoring.PROFILER_ID
            self.active = False
            self.timestamp = None
            self.db_path = db_path
            self.config = config if config is not None else {}
            self.source = source
            self.one_trace_per_test = one_trace_per_test
            self.trace_id = f"trc_{ulid.new()}"

            filters = self.config.get("filters", {})
            self.include_frames = [
                os.path.normpath(f) for f in filters.get("include_frames", [])
            ]
            self.ignore_frames = [
                os.path.normpath(f) for f in filters.get("ignore_frames", [])
            ]

            self.default_include_frames = {}
            for plugin_data in load_plugin_data(self.config):
                processor = PluginProcessor(plugin_data, self.config)
                for co_name in plugin_data["co_names"]:
                    self.default_include_frames.setdefault(co_name, []).append(
                        processor
                    )

            self.default_ignore_frames = [
                LibraryPathFilter().filter,
                frozen_filter,
                pypy_filter,
                kolo_filter_filename,
                exec_filter,
                attrs_filter,
                pytest_filter,
            ]
            self.thread_locals = KoloLocals()
            self.event_map = {
                sys.monitoring.events.PY_START: self.monitor_pystart,
                sys.monitoring.events.PY_RETURN: self.monitor_pyreturn,
                sys.monitoring.events.PY_UNWIND: self.monitor_pyunwind,
                sys.monitoring.events.PY_RESUME: self.monitor_pyresume,
                sys.monitoring.events.PY_YIELD: self.monitor_pyyield,
                sys.monitoring.events.PY_THROW: self.monitor_pythrow,
            }
            self.main_frames = []
            self.thread_frames = defaultdict(list)
            self.main_thread_id = threading.main_thread().native_id
            if self.config.get("lightweight_repr", False):
                self.dump_msgpack = dump_msgpack_lightweight_repr
            else:
                self.dump_msgpack = dump_msgpack

        def monitor(self):
            if sys.monitoring.get_tool(self.tool_id):
                return
            self.active = True
            sys.monitoring.use_tool_id(self.tool_id, "kolo")

            all_events = functools.reduce(
                operator.or_, self.event_map, sys.monitoring.events.NO_EVENTS
            )
            sys.monitoring.set_events(self.tool_id, all_events)
            for event, callback in self.event_map.items():
                sys.monitoring.register_callback(self.tool_id, event, callback)

            self.timestamp = time.time()

        def disable(self):
            if self.active:
                self.active = False
                sys.monitoring.set_events(self.tool_id, sys.monitoring.events.NO_EVENTS)
                for event in self.event_map:
                    sys.monitoring.register_callback(self.tool_id, event, None)
                sys.monitoring.free_tool_id(self.tool_id)
                sys.monitoring.restart_events()

        def ignore(self, filename):
            for ignore in self.default_ignore_frames:
                if ignore(filename):
                    return True
            return False

        def include(self, processor, event, filename, name, arg):
            try:
                if not processor.matches(filename, name):
                    return None

                # Index of the frame in the Python stack:
                # 0: This method
                # 1: KoloMonitor.monitor_ method
                # 2: The frame we want
                frame = sys._getframe(2)
                if processor.call_extra is not None and not processor.call_extra(
                    frame, event, arg, processor.context
                ):
                    return None

                return processor.process(
                    frame, event, arg, self.thread_locals.call_frames
                )
            except Exception as e:
                logger.warning(
                    "Unexpected exception in default_include_frames: %s",
                    processor,
                    exc_info=e,
                )
                return None

        def start_test(self):
            self.trace_id = f"trc_{ulid.new()}"
            self.start_test_index = len(self.main_frames)
            self.start_test_indices = {
                thread_id: len(frames)
                for thread_id, frames in self.thread_frames.items()
            }

        def end_test(self):
            test_frames = {
                thread_id: frames[self.start_test_indices.get(thread_id, 0) :]
                for thread_id, frames in self.thread_frames.items()
            }
            self.save(
                self.main_frames[self.start_test_index :],
                test_frames,
            )

        def monitor_pystart(self, code, instruction_offset):  # pragma: no cover
            filename = code.co_filename
            name = code.co_name

            frames = []
            frame_types = []

            try:
                if name in self.default_include_frames:
                    processors = self.default_include_frames[name]
                    for processor in processors:
                        frame_data = self.include(
                            processor, "call", filename, name, None
                        )
                        if frame_data is not None:
                            frames.append(frame_data)
                            frame_types.append(frame_data["type"])

                for path in self.include_frames:
                    if path in filename:
                        frames.append(self.process_pystart(filename, name))
                        frame_types.append("frame")
                        return

                if self.ignore(filename):
                    if frames:
                        return  # Don't disable if a default processor matched
                    return sys.monitoring.DISABLE

                for path in self.ignore_frames:
                    if path in filename:
                        if frames:
                            return  # Don't disable if a default processor matched
                        return sys.monitoring.DISABLE

                frames.append(self.process_pystart(filename, name))
                frame_types.append("frame")
            finally:
                self.push_frames_call(frames, frame_types)

        def monitor_pyreturn(
            self, code, instruction_offset, retval
        ):  # pragma: no cover
            filename = code.co_filename
            name = code.co_name

            frames = []
            frame_types = []

            try:
                if name in self.default_include_frames:
                    processors = self.default_include_frames[name]
                    for processor in processors:
                        frame_data = self.include(
                            processor, "return", filename, name, retval
                        )
                        if frame_data is not None:
                            frames.append(frame_data)
                            frame_types.append(frame_data["type"])

                for path in self.include_frames:
                    if path in filename:
                        frames.append(self.process_pyreturn(filename, name, retval))
                        frame_types.append("frame")
                        return

                if self.ignore(filename):
                    if frames:
                        return  # Don't disable if a default processor matched
                    return sys.monitoring.DISABLE

                for path in self.ignore_frames:
                    if path in filename:
                        if frames:
                            return  # Don't disable if a default processor matched
                        return sys.monitoring.DISABLE

                frames.append(self.process_pyreturn(filename, name, retval))
                frame_types.append("frame")
            finally:
                self.push_frames_return(frames, frame_types)

        def monitor_pyunwind(
            self, code, instruction_offset, exception
        ):  # pragma: no cover
            filename = code.co_filename
            name = code.co_name

            frames = []
            frame_types = []

            try:
                if name in self.default_include_frames:
                    processors = self.default_include_frames[name]
                    for processor in processors:
                        frame_data = self.include(
                            processor, "unwind", filename, name, exception
                        )
                        if frame_data is not None:
                            frames.append(frame_data)
                            frame_types.append(frame_data["type"])

                for path in self.include_frames:
                    if path in filename:
                        frames.append(self.process_pyunwind(filename, name, exception))
                        frame_types.append("frame")
                        return

                if self.ignore(filename):
                    # We would like to return `sys.monitoring.DISABLE` here, but
                    # `PY_UNWIND` events cannot be disabled, so we do the next best
                    # thing and return asap.
                    return

                for path in self.ignore_frames:
                    if path in filename:
                        # We would like to return `sys.monitoring.DISABLE` here, but
                        # `PY_UNWIND` events cannot be disabled, so we do the next best
                        # thing and return asap.
                        return

                frames.append(self.process_pyunwind(filename, name, exception))
                frame_types.append("frame")
            finally:
                self.push_frames_return(frames, frame_types)

        def monitor_pyresume(self, code, instruction_offset):  # pragma: no cover
            filename = code.co_filename
            name = code.co_name

            frames = []
            frame_types = []

            try:
                if name in self.default_include_frames:
                    processors = self.default_include_frames[name]
                    for processor in processors:
                        frame_data = self.include(
                            processor, "resume", filename, name, None
                        )
                        if frame_data is not None:
                            frames.append(frame_data)
                            frame_types.append(frame_data["type"])

                for path in self.include_frames:
                    if path in filename:
                        frames.append(self.process_pyresume(filename, name))
                        frame_types.append("frame")
                        return

                if self.ignore(filename):
                    if frames:
                        return  # Don't disable if a default processor matched
                    return sys.monitoring.DISABLE

                for path in self.ignore_frames:
                    if path in filename:
                        if frames:
                            return  # Don't disable if a default processor matched
                        return sys.monitoring.DISABLE

                frames.append(self.process_pyresume(filename, name))
                frame_types.append("frame")
            finally:
                self.push_frames_call(frames, frame_types)

        def monitor_pyyield(self, code, instruction_offset, retval):  # pragma: no cover
            filename = code.co_filename
            name = code.co_name

            frames = []
            frame_types = []

            try:
                if name in self.default_include_frames:
                    processors = self.default_include_frames[name]
                    for processor in processors:
                        frame_data = self.include(
                            processor, "yield", filename, name, retval
                        )
                        if frame_data is not None:
                            frames.append(frame_data)
                            frame_types.append(frame_data["type"])

                for path in self.include_frames:
                    if path in filename:
                        frame_data = self.process_pyyield(filename, name, retval)
                        if frame_data is not None:
                            frames.append(frame_data)
                            frame_types.append("frame")
                            return

                if self.ignore(filename):
                    if frames:
                        return  # Don't disable if a default processor matched
                    return sys.monitoring.DISABLE

                for path in self.ignore_frames:
                    if path in filename:
                        if frames:
                            return  # Don't disable if a default processor matched
                        return sys.monitoring.DISABLE

                frame_data = self.process_pyyield(filename, name, retval)
                if frame_data is not None:
                    frames.append(frame_data)
                    frame_types.append("frame")
            finally:
                self.push_frames_return(frames, frame_types)

        def monitor_pythrow(
            self, code, instruction_offset, exception
        ):  # pragma: no cover
            filename = code.co_filename
            name = code.co_name

            frames = []
            frame_types = []

            try:
                if name in self.default_include_frames:
                    processors = self.default_include_frames[name]
                    for processor in processors:
                        frame_data = self.include(
                            processor, "throw", filename, name, exception
                        )
                        if frame_data is not None:
                            frames.append(frame_data)
                            frame_types.append(frame_data["type"])

                for path in self.include_frames:
                    if path in filename:
                        frames.append(self.process_pythrow(filename, name, exception))
                        frame_types.append("frame")
                        return

                if self.ignore(filename):
                    # We would like to return `sys.monitoring.DISABLE` here, but
                    # `PY_UNWIND` events cannot be disabled, so we do the next best
                    # thing and return asap.
                    return

                for path in self.ignore_frames:
                    if path in filename:
                        # We would like to return `sys.monitoring.DISABLE` here, but
                        # `PY_UNWIND` events cannot be disabled, so we do the next best
                        # thing and return asap.
                        return

                frames.append(self.process_pythrow(filename, name, exception))
                frame_types.append("frame")
            finally:
                self.push_frames_return(frames, frame_types)

        def process_pystart(self, filename, name):  # pragma: no cover
            # Index of the frame in the Python stack:
            # 0: This function
            # 1: KoloMonitor.monitor_pystart method
            # 2: The frame we want
            frame = sys._getframe(2)
            frame_id = f"frm_{ulid.new()}"
            self.thread_locals._frame_ids[id(frame)] = frame_id

            user_code_call_site_ = user_code_call_site(
                self.thread_locals.call_frames, frame_id
            )

            self.thread_locals.call_frames.append((frame, frame_id))

            data = build_frame_data(
                "call",
                name,
                filename,
                frame,
                frame_id,
                user_code_call_site_,
            )
            data["arg"] = None
            return data

        def process_pyreturn(self, filename, name, retval):  # pragma: no cover
            # Index of the frame in the Python stack:
            # 0: This function
            # 1: KoloMonitor.monitor_pyreturn method
            # 2: The frame we want
            frame = sys._getframe(2)
            frame_id = self.thread_locals._frame_ids[id(frame)]

            user_code_call_site_ = user_code_call_site(
                self.thread_locals.call_frames, frame_id
            )

            self.thread_locals.call_frames.pop()

            data = build_frame_data(
                "return",
                name,
                filename,
                frame,
                frame_id,
                user_code_call_site_,
            )
            data["arg"] = retval
            return data

        def process_pyunwind(self, filename, name, exception):  # pragma: no cover
            # Index of the frame in the Python stack:
            # 0: This function
            # 1: KoloMonitor.monitor_pyunwind method
            # 2: The frame we want
            frame = sys._getframe(2)
            frame_id = self.thread_locals._frame_ids[id(frame)]

            user_code_call_site_ = user_code_call_site(
                self.thread_locals.call_frames, frame_id
            )

            self.thread_locals.call_frames.pop()

            data = build_frame_data(
                "unwind",
                name,
                filename,
                frame,
                frame_id,
                user_code_call_site_,
            )
            data["exception"] = exception
            return data

        def process_pyresume(self, filename, name):  # pragma: no cover
            # Index of the frame in the Python stack:
            # 0: This function
            # 1: KoloMonitor.monitor_pyresume method
            # 2: The frame we want
            frame = sys._getframe(2)
            frame_id = f"frm_{ulid.new()}"
            self.thread_locals._frame_ids[id(frame)] = frame_id

            user_code_call_site_ = user_code_call_site(
                self.thread_locals.call_frames, frame_id
            )

            self.thread_locals.call_frames.append((frame, frame_id))

            data = build_frame_data(
                "resume",
                name,
                filename,
                frame,
                frame_id,
                user_code_call_site_,
            )
            data["arg"] = None
            return data

        def process_pyyield(self, filename, name, retval):  # pragma: no cover
            # Index of the frame in the Python stack:
            # 0: This function
            # 1: KoloMonitor.monitor_pyyield method
            # 2: The frame we want
            frame = sys._getframe(2)
            try:
                frame_id = self.thread_locals._frame_ids[id(frame)]
            except KeyError:
                # We enabled Kolo in a generator or async function, so missed the start/resume frame
                return

            user_code_call_site_ = user_code_call_site(
                self.thread_locals.call_frames, frame_id
            )

            self.thread_locals.call_frames.pop()

            data = build_frame_data(
                "yield",
                name,
                filename,
                frame,
                frame_id,
                user_code_call_site_,
            )
            data["arg"] = retval
            return data

        def process_pythrow(self, filename, name, exception):  # pragma: no cover
            # Index of the frame in the Python stack:
            # 0: This function
            # 1: KoloMonitor.monitor_pythrow method
            # 2: The frame we want
            frame = sys._getframe(2)
            frame_id = self.thread_locals._frame_ids[id(frame)]

            user_code_call_site_ = user_code_call_site(
                self.thread_locals.call_frames, frame_id
            )

            self.thread_locals.call_frames.pop()

            data = build_frame_data(
                "throw",
                name,
                filename,
                frame,
                frame_id,
                user_code_call_site_,
            )
            data["exception"] = exception
            return data

        def push_frame_data(self, data):
            native_id = data["thread_native_id"]
            data = self.dump_msgpack(data)
            if native_id == self.main_thread_id:
                self.main_frames.append(data)
            else:
                self.thread_frames[native_id].append(data)  # pragma: no cover

        def push_frames_call(self, frames, frame_types):
            if not frames:
                pass
            elif frame_types == ["frame"]:
                self.push_frame_data(frames[0])
            else:
                if self.one_trace_per_test:
                    for index, frame_type in enumerate(frame_types):
                        if frame_type == "start_test":
                            before, frames = frames[:index], frames[index:]
                            for frame_data in before:
                                self.push_frame_data(frame_data)

                            self.start_test()

                for frame_data in frames:
                    self.push_frame_data(frame_data)

        def push_frames_return(self, frames, frame_types):
            if not frames:
                pass
            elif frame_types == ["frame"]:
                self.push_frame_data(frames[0])
            else:
                frames.reverse()
                frame_types.reverse()

                if self.one_trace_per_test:
                    for index, frame_type in enumerate(frame_types):
                        if frame_type == "end_test":
                            before, frames = frames[: index + 1], frames[index + 1 :]
                            for frame_data in before:
                                self.push_frame_data(frame_data)

                            self.end_test()

                for frame_data in frames:
                    self.push_frame_data(frame_data)

        def save(self, main_frames=None, thread_frames=None):
            timeout = self.config.get("sqlite_busy_timeout", 60)
            data = {
                "command_line_args": sys.argv,
                "current_commit_sha": COMMIT_SHA,
                "main_thread_id": self.main_thread_id,
                "meta": {
                    "version": __version__,
                    "source": self.source,
                    "use_frame_boundaries": True,
                },
                "timestamp": self.timestamp,
                "trace_id": self.trace_id,
            }

            if main_frames is None:
                main_frames = self.main_frames
            if thread_frames is None:
                thread_frames = self.thread_frames

            data["frames_of_interest"] = [
                msgpack.unpackb(f, strict_map_key=False) for f in main_frames
            ]
            data["frames"] = {
                thread_id: [msgpack.unpackb(f, strict_map_key=False) for f in frames]
                for thread_id, frames in thread_frames.items()
            }
            serialized_data = self.dump_msgpack(data)
            save_trace_in_sqlite(
                self.db_path, self.trace_id, msgpack=serialized_data, timeout=timeout
            )
