# coding=utf-8
#
# Copyright © Cloud Linux GmbH & Cloud Linux Software, Inc 2010-2023 All Rights Reserved
#
# Licensed under CLOUD LINUX LICENSE AGREEMENT
# http://cloudlinux.com/docs/LICENSE.TXT
import enum
import os
import contextlib
import gc
import time
import statistics
import typing
from pathlib import Path
from contextlib import ExitStack
from collections import defaultdict
from enum import StrEnum
from dataclasses import dataclass
from typing import NamedTuple, Generator, Callable, TypeAlias, ContextManager, Self, Sequence, Iterable, Protocol

import sqlalchemy as sa

import psutil
import pytest
from _pytest.terminal import TerminalReporter

_db_metadata = sa.MetaData()
_measurements_table = sa.Table(
    'measurements',
    _db_metadata,
    sa.Column('sequence_id', sa.String, primary_key=True),
    sa.Column('timestamp', sa.Float, primary_key=True),
    sa.Column('cpu_usage', sa.Float, nullable=False),
    sa.Column('run_time_seconds', sa.Float, nullable=False),
    sa.Column('memory_allocated_bytes', sa.Integer, nullable=False),
    sa.Column('total_memory_bytes', sa.Integer, nullable=False),
)


class Measurement(NamedTuple):
    timestamp: float
    cpu_usage: float
    run_time_seconds: float
    memory_allocated_bytes: int
    total_memory_bytes: int

    def __str__(self) -> str:
        memory_allocations_mb = _bytes_to_mb(self.memory_allocated_bytes)
        total_memory_mb = _bytes_to_mb(self.total_memory_bytes)
        return '\n'.join([
            f'Time             : {self.run_time_seconds} sec',
            f'CPU usage        : {self.cpu_usage:.2f}%',
            f'Memory allocated : {memory_allocations_mb:.2f} mb',
            f'Total memory     : {total_memory_mb:.2f} mb',
        ])


assert {c.name for c in list(_measurements_table.columns)}.issuperset(Measurement._fields)


@dataclass
class Statistic:
    @classmethod
    def from_data(cls, data: Sequence[float]) -> Self:
        return Statistic(
            mean=statistics.mean(data),
            std_dev=statistics.stdev(data) if len(data) > 1 else 0,
            max_value=max(data),
            min_value=min(data)
        )

    mean: float
    std_dev: float
    max_value: float
    min_value: float


@dataclass(frozen=True, slots=True)
class StatisticSummarySection:
    @classmethod
    def from_statistic(cls, title: str, statistic: Statistic, units: str) -> Self:
        return cls(
            title=title,
            mean=f'{statistic.mean:.2f}{units}',
            std_dev=f'{statistic.std_dev:.2f}{units}',
            max_value=f'{statistic.max_value:.2f}{units}',
            min_value=f'{statistic.min_value:.2f}{units}',
        )

    @classmethod
    def empty(cls, title: str) -> Self:
        return cls(
            title=title,
            mean='N/A',
            std_dev='N/A',
            max_value='N/A',
            min_value='N/A',
        )

    title: str
    mean: str
    std_dev: str
    max_value: str
    min_value: str


@enum.unique
class _SummarySection(StrEnum):
    CPU_USAGE = 'CPU Usage'
    RUN_TIME = 'Run Time'
    MEMORY_ALLOCATED = 'Memory Allocated'
    TOTAL_MEMORY = 'Total Memory'


def _create_summary(sections: Iterable[StatisticSummarySection]) -> str:
    summary = ['Measurement Statistics Summary']
    summary.append('----------------------------')
    for section in sections:
        summary.append(section.title)
        summary.append(f'  Mean: {section.mean}')
        summary.append(f'  Std Dev: {section.std_dev}')
        summary.append(f'  Max: {section.max_value}')
        summary.append(f'  Min: {section.min_value}')

    return "\n".join(summary)


def _measurements_to_summary(measurements: Iterable[Measurement]) -> str:
    title_to_data = defaultdict[_SummarySection, list[float]](list)
    for measurement in measurements:
        title_to_data[_SummarySection.CPU_USAGE].append(measurement.cpu_usage)
        title_to_data[_SummarySection.RUN_TIME].append(measurement.run_time_seconds)
        title_to_data[_SummarySection.MEMORY_ALLOCATED].append(_bytes_to_mb(measurement.memory_allocated_bytes))
        title_to_data[_SummarySection.TOTAL_MEMORY].append(_bytes_to_mb(measurement.total_memory_bytes))

    sections = [StatisticSummarySection.from_statistic(
        title=title,
        statistic=Statistic.from_data(data),
        units={
            _SummarySection.CPU_USAGE: '%',
            _SummarySection.RUN_TIME: ' sec',
            _SummarySection.MEMORY_ALLOCATED: ' mb',
            _SummarySection.TOTAL_MEMORY: ' mb',
        }[title],
    ) for title, data in title_to_data.items()]
    return _create_summary(sections)


_empty_summary = _create_summary(StatisticSummarySection.empty(s.name) for s in _SummarySection)


Profiled: TypeAlias = Callable[[], ContextManager[None]]


def _get_uss_memory(process: psutil.Process) -> int:
    # NOTE(vlebedev): On the background of why USS is used, have a look at
    #                 https://gmpy.dev/blog/2016/real-process-memory-and-environ-in-python
    key = 'memory_full_info'
    result = process.as_dict(attrs=[key])[key].uss
    return result


def _bytes_to_mb(bytes_: int) -> float:
    return bytes_ / 1000000


class _Profiler:
    def __init__(self, measurements: list[Measurement]) -> None:
        self._pid = os.getpid()
        self._measurements = measurements

    @contextlib.contextmanager
    def __call__(self) -> Generator[None, None, None]:
        process = psutil.Process(self._pid)

        gc.disable()
        start_time = time.perf_counter()
        memory_before = _get_uss_memory(process)
        process.cpu_percent()
        try:
            yield
        finally:
            end_time = time.perf_counter()
            memory_after = _get_uss_memory(process)
            cpu_usage_after = process.cpu_percent()
            memory_delta_bytes = memory_after - memory_before
            measurement = Measurement(
                timestamp=time.time(),
                cpu_usage=cpu_usage_after,
                run_time_seconds=end_time - start_time,
                memory_allocated_bytes=memory_delta_bytes,
                total_memory_bytes=memory_after,
            )
            self._measurements.append(measurement)
            gc.enable()
            gc.collect()


class _WriteMeasurements(Protocol):
    def __call__(self, measurements: Iterable[Measurement]) -> None:
        ...


_MEASUREMENTS_KEY = 'measurements'
_MEASUREMENTS_ENABLED_FLAG = '--with-measurements'
_SAVE_MEASUREMENTS_FLAG = '--save-measurements'


class PytestProfilerPlugin:
    def pytest_configure(self, config: pytest.Config):
        if not config.getoption(_MEASUREMENTS_ENABLED_FLAG):
            @self._fixture(scope='function', name='profiled')
            def create_profiler_stub() -> Profiled:
                @contextlib.contextmanager
                def stub() -> Generator[None, None, None]:
                    yield
                return stub
            return

        if config.getoption(_SAVE_MEASUREMENTS_FLAG):
            @self._fixture(scope='function', name='_write_measurements')
            def create_measurements_writer(
                request: pytest.FixtureRequest,
            ) -> Generator[_WriteMeasurements, None, None]:
                db_file = Path(request.path).parent / 'measurements.db'
                sequence_id = request.node.nodeid

                with ExitStack() as deffer:
                    engine = sa.create_engine(f'sqlite+pysqlite:///{db_file}', echo=__debug__)
                    deffer.callback(engine.dispose)

                    conn = deffer.enter_context(engine.connect())

                    _db_metadata.create_all(conn)

                    def write_measurements(measurements: Iterable[Measurement]) -> None:
                        stmt = _measurements_table.insert(None).values([{
                            'sequence_id': sequence_id,
                            **m._asdict(),
                        } for m in measurements])
                        conn.execute(stmt)

                    yield write_measurements
        else:
            @self._fixture(scope='function', name='_write_measurements')
            def create_measurements_writer_stub() -> _WriteMeasurements:
                def stub(measurements: Iterable[Measurement]) -> None:
                    # NOTE(vlebedev): Make pylance happy by "using" all arguments.
                    del measurements
                return stub

        @self._fixture(scope='function', name='profiled')
        def profiler_created(
            record_property: Callable[[str, object], None],
            _write_measurements: _WriteMeasurements,
        ) -> Generator[Profiled, None, None]:
            measurements = []
            record_property(_MEASUREMENTS_KEY, measurements)
            yield _Profiler(measurements)
            _write_measurements(measurements)

    def pytest_addoption(self, parser: pytest.Parser):
        parser.addoption(
            _MEASUREMENTS_ENABLED_FLAG,
            action='store_true',
            default=False,
            help='Enable measurements collection during each test',
        )
        parser.addoption(
            _SAVE_MEASUREMENTS_FLAG,
            action='store_true',
            default=False,
            help='Save measurements to the database',
        )

    def pytest_report_teststatus(self, report: pytest.TestReport, config: pytest.Config):
        if not (config.getoption(_MEASUREMENTS_ENABLED_FLAG) and config.option.verbose >= 1):
            return
        if report.when != 'teardown':
            return

        tr: TerminalReporter = config.pluginmanager.get_plugin('terminalreporter')
        if tr is None:
            return

        try:
            measurements = typing.cast(
                list[Measurement],
                next(v for k, v in report.user_properties if k == _MEASUREMENTS_KEY),
            )
        except StopIteration:
            return

        try:
            summary = _measurements_to_summary(measurements)
        except ValueError:
            summary = _empty_summary

        tr.ensure_newline()
        tr.write_line(summary)

    def _fixture(self, *args, **kwargs):
        def wrapper(func):
            fixture = pytest.fixture(*args, **kwargs)(func)
            # NOTE(vlebedev): Without assigning fixture to some property,
            #                 pytest won't be able to find and use it.
            setattr(self, kwargs.get('name', func.__name__), fixture)
            return fixture
        return wrapper