import time
import os
import numpy as np
from scipy import stats
from dataclasses import dataclass, field

# ── Linux CPU affinity ──────────────────────────────────────────────────────
try:
    import ctypes
    libc = ctypes.CDLL("libc.so.6", use_errno=True)

    def pin_to_core(core: int = 0):
        """Pin this process to a single CPU core to reduce scheduling jitter."""
        cpu_set_size = 128
        cpu_set = ctypes.create_string_buffer(cpu_set_size)
        byte_idx = core // 8
        bit_idx  = core % 8
        cpu_set[byte_idx] = bytes([1 << bit_idx])
        ret = libc.sched_setaffinity(0, cpu_set_size, cpu_set)
        if ret != 0:
            print(f"  [!] Could not pin to core {core}. Continuing anyway.")
        else:
            print(f"  [*] Pinned to CPU core {core}")

except Exception:
    def pin_to_core(core: int = 0):
        print("  [!] CPU pinning not available. Continuing without it.")


@dataclass
class TVLAResult:
    operation:      str
    t_statistic:    float
    p_value:        float
    leaking:        bool
    severity:       str
    fixed_mean_ns:  float
    random_mean_ns: float
    difference_ns:  float
    num_traces:     int
    verdict:        str
    fixed_times:    list = field(default_factory=list, repr=False)
    random_times:   list = field(default_factory=list, repr=False)


class TimingAnalyzer:
    """
    Test Vector Leakage Assessment (TVLA) using Welch's t-test.
    ISO 17825 standard methodology.

    Production improvements:
    - Warm-up phase: stabilizes CPU caches and branch predictor
    - Outlier trimming: removes OS interrupt spikes (top 5%)
    - CPU affinity: reduces cross-core scheduling jitter
    - Interleaved measurements: controls for system load changes
    """

    TVLA_THRESHOLD  = 4.5
    WARMUP_TRACES   = 500
    OUTLIER_PERCENT = 95

    def __init__(self, num_traces: int = 10000):
        self.num_traces = num_traces

    def _time_once(self, fn) -> int:
        start = time.perf_counter_ns()
        fn()
        return time.perf_counter_ns() - start

    def _warmup(self, fixed_fn, random_fn):
        """
        Run 500 throwaway iterations before measurement.
        Stabilizes CPU instruction cache, branch predictor,
        memory allocator, and Python interpreter state.
        """
        print(f"    warming up ({self.WARMUP_TRACES} discarded traces)...")
        for _ in range(self.WARMUP_TRACES):
            fixed_fn()
            random_fn()

    def _trim_outliers(self, times: list) -> list:
        """
        Remove top 5% of measurements.
        These are OS interrupts, GC pauses, and scheduler preemptions.
        They are measurement artifacts, not cryptographic signal.
        """
        threshold = np.percentile(times, self.OUTLIER_PERCENT)
        trimmed   = [t for t in times if t <= threshold]
        removed   = len(times) - len(trimmed)
        if removed > 0:
            print(f"    trimmed {removed} outliers above {threshold/1e6:.3f}ms")
        return trimmed

    def collect(self, fixed_fn, random_fn) -> tuple[list, list]:
        self._warmup(fixed_fn, random_fn)

        fixed_times, random_times = [], []
        for i in range(self.num_traces):
            if i % 1000 == 0:
                print(f"    traces: {i}/{self.num_traces}")
            fixed_times.append(self._time_once(fixed_fn))
            random_times.append(self._time_once(random_fn))

        fixed_times  = self._trim_outliers(fixed_times)
        random_times = self._trim_outliers(random_times)
        return fixed_times, random_times

    def analyze(self, op_name: str, fixed_fn, random_fn) -> TVLAResult:
        fixed, random = self.collect(fixed_fn, random_fn)

        t_stat, p_value = stats.ttest_ind(fixed, random, equal_var=False)
        abs_t   = abs(t_stat)
        leaking = abs_t > self.TVLA_THRESHOLD

        return TVLAResult(
            operation=op_name,
            t_statistic=round(t_stat, 4),
            p_value=round(p_value, 8),
            leaking=leaking,
            severity=self._severity(abs_t),
            fixed_mean_ns=round(np.mean(fixed), 1),
            random_mean_ns=round(np.mean(random), 1),
            difference_ns=round(abs(np.mean(fixed) - np.mean(random)), 1),
            num_traces=len(fixed),
            verdict="LEAK DETECTED" if leaking else "PASS",
            fixed_times=fixed,
            random_times=random
        )

    def _severity(self, t: float) -> str:
        if t > 10:  return "CRITICAL"
        if t > 4.5: return "HIGH"
        return "PASS"
