#!/usr/libexec/platform-python

from __future__ import print_function
import ctypes as ct
import os
import stat
import argparse
import time
import signal
from collections import defaultdict
from operator import attrgetter
from bcc import BPF
from bcc.table import _stars

DISK_NAME_LEN = 32

# ==================== IO Trace事件定义（统一管理）====================
# 格式: (名称, 显示标签, 是否有处理器)
_EVENT_SPECS = [
    ('QUEUE',      'Q',     True),
    ('INSERT',     'I',     True),
    ('ISSUE',      'D',     True),
    ('COMPLETE',   'C',     True),
    ('SPLIT',      'SPLIT', True),
    ('FRONTMERGE', None,    False),
    ('BACKMERGE',  None,    False),
    ('MERGE',      None,    False),
]

# 批量定义所有ACTION常量（元组解包，确保类型检查工具可见）
(ACTION_QUEUE, ACTION_INSERT, ACTION_ISSUE, ACTION_COMPLETE,
 ACTION_SPLIT, ACTION_FRONTMERGE, ACTION_BACKMERGE, ACTION_MERGE) = range(1, 9)

# 批量定义所有COUNTER常量
(COUNTER_QUEUE, COUNTER_INSERT, COUNTER_ISSUE, COUNTER_COMPLETE,
 COUNTER_SPLIT, COUNTER_FRONTMERGE, COUNTER_BACKMERGE, COUNTER_MERGE) = range(8)

# 自动生成映射表
EVENT_TYPE_LABELS = {}
EVENT_COUNTER_NAMES = {}
EVENT_ACTION_NAMES = {}

for idx, (name, label, has_handler) in enumerate(_EVENT_SPECS):
    action_id = idx + 1
    counter_id = idx

    # 构建映射表
    if label:
        EVENT_TYPE_LABELS[action_id] = label
    EVENT_COUNTER_NAMES[counter_id] = name
    EVENT_ACTION_NAMES[action_id] = name

# ==================== BPF和运行时配置 ====================
# BPF配置索引，与内核侧保持一致
CFG_IDX_MAJOR = 0
CFG_IDX_MINOR = 1
CFG_IDX_DETAIL = 2
CFG_IDX_PID_FILTER = 3
CFG_IDX_DEBUG = 4

# detail 级别语义
DETAIL_LEVEL_LATENCY_FIRST = 1
DETAIL_LEVEL_LATENCY_SECOND = 2

# 设备号编码参数
DEVICE_MINOR_SHIFT = 20
DEVICE_MINOR_MASK = (1 << DEVICE_MINOR_SHIFT) - 1

# 时间单位常量
NS_PER_MS = 1_000_000
NS_PER_SEC = 1_000_000_000

SECTOR_SIZE = 512
# 运行时参数
CLEANUP_INTERVAL_SEC = 5
MAX_SEARCH_OFFSET_SECTORS = 4096
EVENT_COUNTER_SLOTS = 8
TOTAL_BYTES_INDEX = 0

# 输出格式化
SEPARATOR = "-" * 80

# 时间阈值配置
TIME_THRESHOLDS = {
    'Q2I_MAX': 1000 * NS_PER_MS,
    'I2D_MAX': 1000 * NS_PER_MS,
    'D2C_MAX': 5000 * NS_PER_MS,
}

def format_size(bytes_val, precision=2):
    if bytes_val >= 1024 * 1024:
        return f"{bytes_val / (1024*1024):.{precision}f}M"
    if bytes_val >= 1024:
        return f"{bytes_val / 1024:.{precision}f}K"
    return f"{bytes_val:.{precision}f}B"

def format_time(time_ns, precision=2):
    if time_ns >= 1_000_000_000:
        return f"{time_ns / 1_000_000_000:.{precision}f}s"
    if time_ns >= 1_000_000:
        return f"{time_ns / 1_000_000:.{precision}f}ms"
    if time_ns >= 1_000:
        return f"{time_ns / 1_000:.{precision}f}us"
    return f"{time_ns:.{precision}f}ns"

def find_matching_entry(tree, sector, validator_fn):
    """通用的树匹配函数"""
    stack = tree.get(sector, [])
    for i, entry in enumerate(reversed(stack)):
        if entry['matched']:
            continue
        result = validator_fn(entry)
        if result:
            index = len(stack) - 1 - i
            return result, index
    return None, None

def find_original_bio(Q_tree, event):
    """查找原始BIO (用于Split事件)
    Split可能从原始bio的任意位置开始，需要向前搜索
    """
    max_offset = min(MAX_SEARCH_OFFSET_SECTORS, event.sector)
    for offset in range(max_offset + 1):
        search_sector = event.sector - offset
        q_stack = Q_tree.get(search_sector, [])
        for q_info in reversed(q_stack):
            # 提前过滤不符合条件的
            if q_info['matched'] or q_info['bio_type'] != 'ORIGINAL':
                continue
            # 范围检查
            original_start = search_sector
            original_end = search_sector + (q_info['bytes'] // 512)
            if event.sector < original_start or event.sector >= original_end:
                continue
            # 字节检查
            remaining_bytes = q_info['bytes'] - q_info.get('split_consumed', 0)
            if remaining_bytes < event.bytes:
                continue
            return q_info
    return None

class LightBlktraceEvent(ct.Structure):
    _fields_ = [
        ("action", ct.c_uint32),
        ("cpu", ct.c_uint32),
        ("timestamp", ct.c_uint64),
        ("req_id", ct.c_uint64),
        ("sector", ct.c_uint64),
        ("bytes", ct.c_uint32),
        ("rwflag", ct.c_uint32),
        ("pid", ct.c_uint32),
        ("device", ct.c_uint32),
        ("comm", ct.c_char * 16),
        ("disk_name", ct.c_char * DISK_NAME_LEN),
    ]


class LatencyMetric(object):
    def __init__(self):
        self.lat_total_ns = 0  # 累计时延（纳秒）
        self.count = 0         # 事件计数

    def add_sample(self, latency_ns):
        if latency_ns > 0:
            self.lat_total_ns += latency_ns
            self.count += 1

    @property
    def average_ns(self):
        return self.lat_total_ns // self.count if self.count else 0


class ProcessStats(object):
    def __init__(self):
        self.pid = 0
        self.comm = ""
        self.read_bytes = 0
        self.write_bytes = 0
        self.total_bytes = 0
        self.bio_count = 0  # 记录bio数量，用于计算平均bio大小
        self.bio_size_sum = 0  # 记录bio大小总和

    def update(self, data_len, is_write=False):
        if is_write:
            self.write_bytes += data_len
        else:
            self.read_bytes += data_len
        self.total_bytes += data_len

    def add_bio(self, bio_size):
        """记录一个bio的大小"""
        self.bio_count += 1
        self.bio_size_sum += bio_size


class LatencyStats(object):
    """
    IO延迟统计数据类（纯数据）
    """

    def __init__(self):
        self.disk_name = ""

        # 时延统计: Q → S → I → D → C (使用LatencyMetric封装)
        self.q2s = LatencyMetric()  # Queue到Split时延
        self.s2i = LatencyMetric()  # Split到Insert时延
        self.q2i = LatencyMetric()  # Queue到Insert时延 (无Split的情况)
        self.i2d = LatencyMetric()  # Insert到Issue时延
        self.q2d = LatencyMetric()  # Queue到Issue时延 (none调度器，跳过Insert)
        self.d2c = LatencyMetric()  # Issue到Complete时延
        self.q2c = LatencyMetric()  # Queue到Complete (端到端)
        self.q2q = LatencyMetric()  # Queue到Queue间隔
        self.last_queue_time = 0

        # IO大小统计
        self.total_bio_bytes = 0
        self.total_request_bytes = 0

        # Merge/Split统计
        self.frontmerge_count = 0
        self.backmerge_count = 0
        self.merge_count = 0
        self.split_count = 0
        self.split_original_bytes = 0
        self.split_after_bytes = 0
        self.split_original_bio_count = 0

        # 调度器模式检测
        self.insert_count = 0

    def is_none_scheduler(self):
        """判断是否为none调度器模式（无Insert事件）"""
        return self.insert_count == 0

    def add_q2q(self, timestamp, disk_name):
        # 更新磁盘名称
        name = bytes(disk_name).split(b"\0", 1)[0].decode("utf-8", "replace")
        if name and not self.disk_name:
            self.disk_name = name

        if self.last_queue_time > 0:
            q2q = timestamp - self.last_queue_time
            self.q2q.add_sample(q2q)
        self.last_queue_time = timestamp

    def add_bio_bytes(self, bytes_val):
        self.total_bio_bytes += bytes_val

    def add_request_bytes(self, bytes_val):
        self.total_request_bytes += bytes_val

    def add_merge(self, action):
        if action == ACTION_FRONTMERGE:
            self.frontmerge_count += 1
        elif action == ACTION_BACKMERGE:
            self.backmerge_count += 1
        elif action == ACTION_MERGE:
            self.merge_count += 1


class LatencyAsciiRenderer(object):
    """以更直观的表格形式展示延迟数据"""

    INDENT = "    "

    @staticmethod
    def _format_metric(label, avg_ns, count, percent="--"):
        time_str = format_time(avg_ns, precision=0)
        percent_part = f"{percent:>8}"
        return f"{label:<5}{time_str:>10}{count:>10}{percent_part}"

    @classmethod
    def render(cls, stats, major, minor):
        device = stats.disk_name or "unknown"

        # 根据调度器模式选择显示指标
        if stats.is_none_scheduler():
            # none调度器: Q → D → C (无Insert阶段)
            labels = ("Q2D", "D2C", "Q2C", "Q2Q")
            total_latency = stats.q2d.lat_total_ns + stats.d2c.lat_total_ns
        else:
            # 传统调度器: Q → I → D → C (完整链路)
            labels = ("Q2I", "Q2S", "S2I", "I2D", "D2C", "Q2C", "Q2Q")
            total_latency = (stats.q2s.lat_total_ns + stats.s2i.lat_total_ns +
                             stats.q2i.lat_total_ns + stats.i2d.lat_total_ns +
                             stats.d2c.lat_total_ns)

        lines = [f"{device} {major}-{minor}:"]
        for label in labels:
            attr = label.lower()
            metric = getattr(stats, attr)
            avg_ns = metric.average_ns
            count = metric.count
            if label == "Q2C" or label == "Q2Q":
                percent = "--"
            else:
                percent = f"{(metric.lat_total_ns * 100.0 / total_latency):.2f}%" if total_latency else "--"
            lines.append(cls.INDENT + cls._format_metric(label, avg_ns, count, percent))

        total_merges = (stats.frontmerge_count + stats.backmerge_count + stats.merge_count)
        if total_merges > 0:
            lines.append("\n" + cls.INDENT + "Merge Statistics:")
            lines.append(
                cls.INDENT * 2 +
                f"Merges    total:{total_merges:5} (FM:{stats.frontmerge_count} "
                f"BM:{stats.backmerge_count} RM:{stats.merge_count})"
            )
        if stats.split_count > 0:
            lines.append("\n" + cls.INDENT + "Split Statistics:")
            avg_original = (stats.split_original_bytes /
                            stats.split_original_bio_count
                            if stats.split_original_bio_count > 0 else 0)
            avg_after = (stats.split_after_bytes / stats.split_count
                         if stats.split_count > 0 else 0)
            split_ratio = (stats.split_after_bytes * 100 //
                           stats.split_original_bytes
                           if stats.split_original_bytes > 0 else 0)
            avg_original_str = format_size(avg_original, precision=0)
            avg_after_str = format_size(avg_after, precision=0)
            lines.append(
                cls.INDENT * 2 +
                f"Splits    {stats.split_original_bio_count} bios -> "
                f"{stats.split_count} splits  "
                f"Avg: {avg_original_str} -> {avg_after_str} ({split_ratio}%)"
            )

        return "\n".join(lines)


class ResultPrinter(object):
    """统一的结果输出类"""

    @classmethod
    def print_latency_stats(cls, device_stats):
        if not device_stats:
            return

        lines = ["\n\nlatency result:", "(Avg:time, Cnt:count, Pct:share%)\n"]
        for (major, minor), stats in device_stats.items():
            lines.append(LatencyAsciiRenderer.render(stats, major, minor))

        print("\n".join(lines))

    @classmethod
    def print_io_size_histogram(cls, bpf):
        """打印IO大小分布直方图"""
        size_map = bpf.get_table("io_size_hist")

        lines = ["\n\nIO Size Distribution:"]
        print("\n".join(lines))
        size_map.print_log2_hist("size", strip_leading_zero=True)

    @classmethod
    def print_latency_histogram(cls, bpf):
        lat_map = bpf.get_table("io_lat_hist")
        # defaultdict需要工厂函数，lambda用于为每个新键创建独立的字典对象
        groups = defaultdict(lambda: {"total": 0, "slots": defaultdict(int)})

        for key, leaf in lat_map.items():
            count = leaf.value
            if count == 0:
                continue
            entry = groups[key.size]
            entry["total"] += count
            entry["slots"][key.slot] += count

        if not groups:
            return

        lines = ["\nLatency Distribution (by IO size):", SEPARATOR]

        for size_key in sorted(groups):
            entry = groups[size_key]
            total_count = entry["total"]

            latency_total = sum(
                slot_count * (1 << slot_key)
                for slot_key, slot_count in entry["slots"].items()
            )

            size_bytes = 1 << size_key
            size_str = format_size(size_bytes)
            lines.append(f"\nIO Size {size_str:>6}:")

            max_slot_count = max(entry["slots"].values())
            bar_width = 40

            for slot_key in sorted(entry["slots"]):
                slot_count = entry["slots"][slot_key]
                latency_us = 1 << slot_key
                latency_ns = latency_us * 1000

                count_pct = (slot_count * 100.0 / total_count) if total_count else 0
                lat_pct = (slot_count * latency_us * 100.0 / latency_total) if latency_total else 0

                bar = _stars(slot_count, max_slot_count, bar_width)
                time_str = format_time(latency_ns, precision=0)

                lines.append(f"  <{time_str:>8} |{bar:<40}| {slot_count:>5} "
                             f"({count_pct:>5.2f}% cnt, {lat_pct:>5.2f}% lat)")

        print("\n".join(lines))

    @classmethod
    def print_process_stats(cls, process_stats, global_total_bytes):
        if not process_stats:
            return

        header = f"{'PID':<8}{'COMM':<16}{'TOTAL':>12}{'READ':>12}{'WRITE':>12}{'AVG_BIO':>13}{'PERCENT':>10}"
        separator = "-" * len(header)

        lines = ["\n\nprocess IO statistics:\n", header, separator]
        for pid, stats in process_stats.items():
            total_str = format_size(stats.total_bytes)
            read_str = format_size(stats.read_bytes)
            write_str = format_size(stats.write_bytes)
            avg_bio_bytes = (stats.bio_size_sum / stats.bio_count) if stats.bio_count > 0 else 0
            avg_bio_str = format_size(avg_bio_bytes)
            percent = (stats.total_bytes * 100.0 / global_total_bytes) if global_total_bytes > 0 else 0.0
            lines.append(
                f"{pid:<8}{stats.comm[:15]:<16}{total_str:>12}{read_str:>12}"
                f"{write_str:>12}{avg_bio_str:>13}{percent:>9.2f}%"
            )

        print("\n".join(lines))

    @staticmethod
    def print_tracing_status(start_time, event_counts, lost_events):
        time_str = f"{int(time.time() - start_time)}s"
        total_events = sum(event_counts.values())

        with open('/proc/self/status', 'r') as f:
            mem_str = "N/A"
            for line in f:
                if line.startswith('VmRSS:'):
                    mem_kb = int(line.split()[1])
                    mem_str = format_size(mem_kb * 1024, precision=1)
                    break

        status = f"\rTracing: {time_str} | Events: {total_events:,} | Lost: {lost_events:,} | Mem: {mem_str}"
        print(status, end='', flush=True)

    @staticmethod
    def print_raw_event(debug_enabled, event_type, event):
        """输出原始tracepoint事件 (类似ftrace格式)"""
        if not debug_enabled:
            return

        comm = event.comm.decode('utf-8', 'replace').strip('\x00')
        bytes_sectors = f"{event.bytes}/{event.bytes // SECTOR_SIZE}"
        line = (f"[{event_type:6s}] {event.cpu:<3d} {comm:16s} "
                f"{event.pid:<7d} {event.sector:<12d} "
                f"{bytes_sectors:<15} {event.timestamp}")
        print(line)

    @staticmethod
    def print_bpf_event_counters(bpf):
        """打印BPF内核侧的事件计数器统计"""
        event_counters = bpf.get_table("event_counters")

        lines = ["\n\n# Event Counters (from kernel):", "#" + SEPARATOR]
        total_events = 0
        for idx in sorted(EVENT_COUNTER_NAMES.keys()):
            key = ct.c_int(idx)
            count = int(event_counters[key].value)
            total_events += count
            name = EVENT_COUNTER_NAMES[idx]
            lines.append(f"  {name:<15}: {count:>10,}")

        lines.extend([SEPARATOR, f"  {'TOTAL':<15}: {total_events:>10,}",
                      "#" + SEPARATOR + "\n"])

        print("\n".join(lines))

class EventMatcher(object):
    """事件匹配逻辑类"""

    def __init__(self, time_thresholds):
        self.Q2I_MAX = time_thresholds.get('Q2I_MAX', 1000 * NS_PER_MS)
        self.Q2D_MAX = self.Q2I_MAX  # none调度器: Q→D的时延上限等同Q→I
        self.I2D_MAX = time_thresholds.get('I2D_MAX', 1000 * NS_PER_MS)
        self.D2C_MAX = time_thresholds.get('D2C_MAX', 5000 * NS_PER_MS)

    def match_insert_to_queue(self, q_info, event, stats_collector):
        """匹配Insert事件到Queue记录"""
        if event.bytes > q_info['bytes']:
            return None

        if q_info['bio_type'] == 'ORIGINAL':
            time_diff = event.timestamp - q_info['ts']
            if time_diff > self.Q2I_MAX:
                return None

            q_info['matched'] = True
            stats_collector.record_latency('q2i', time_diff)
            return q_info

        elif q_info['bio_type'] == 'SPLIT':
            split_ts = q_info.get('split_ts', q_info['ts'])
            time_diff = event.timestamp - split_ts
            if time_diff > self.Q2I_MAX:
                return None

            q_info['matched'] = True
            stats_collector.record_latency('s2i', time_diff)
            return q_info

        return None

    def match_issue_to_queue(self, q_info, event, stats_collector):
        """匹配Issue事件到Queue记录（none调度器模式，跳过Insert）"""
        if event.bytes > q_info['bytes']:
            return None

        if q_info['bio_type'] == 'ORIGINAL':
            time_diff = event.timestamp - q_info['ts']
            if time_diff > self.Q2D_MAX:
                return None

            q_info['matched'] = True
            stats_collector.record_latency('q2d', time_diff)
            return q_info

        elif q_info['bio_type'] == 'SPLIT':
            # Split场景下，从split时间戳计算
            split_ts = q_info.get('split_ts', q_info['ts'])
            time_diff = event.timestamp - split_ts
            if time_diff > self.Q2D_MAX:
                return None

            q_info['matched'] = True
            stats_collector.record_latency('q2d', time_diff)
            return q_info

        return None

    def match_issue_to_insert(self, i_info, event, stats_collector):
        """匹配Issue事件到Insert记录"""
        time_diff = event.timestamp - i_info['ts']
        if time_diff > self.I2D_MAX:
            return None

        i_info['matched'] = True
        stats_collector.record_latency('i2d', time_diff)
        return i_info

    def match_complete_to_issue(self, d_info, event):
        """匹配Complete事件到Issue记录"""
        time_diff = event.timestamp - d_info['ts']
        if time_diff > self.D2C_MAX:
            return None
        return (d_info, time_diff)


class StatisticsCollector(object):
    """统计信息收集类"""

    def __init__(self, device_stats, process_stats, show_process_stats):
        self.device_stats = device_stats
        self.process_stats = process_stats
        self.show_process_stats = show_process_stats
        self.current_stats = None

    def set_current_device(self, dev_key):
        """设置当前处理的设备"""
        if dev_key not in self.device_stats:
            self.device_stats[dev_key] = LatencyStats()
        self.current_stats = self.device_stats[dev_key]

    def _get_or_create_process_stats(self, pid, comm):
        """获取或创建进程统计对象"""
        if pid not in self.process_stats:
            self.process_stats[pid] = ProcessStats()
            self.process_stats[pid].pid = pid
            self.process_stats[pid].comm = comm
        return self.process_stats[pid]

    def record_latency(self, metric_name, time_diff):
        """通用延迟统计: q2i/s2i/i2d/d2c/q2c/q2s"""
        getattr(self.current_stats, metric_name).add_sample(time_diff)

    def add_split_stats(self, after_bytes, original_bytes=0, original_bio_count=0):
        self.current_stats.split_count += 1
        self.current_stats.split_after_bytes += after_bytes
        if original_bytes > 0:
            self.current_stats.split_original_bytes += original_bytes
            self.current_stats.split_original_bio_count += original_bio_count

    def update_process_stats(self, pid, comm, bytes_val, is_write):
        """更新进程统计"""
        if not self.show_process_stats:
            return

        proc_stats = self._get_or_create_process_stats(pid, comm)
        proc_stats.update(bytes_val, is_write=is_write)

    def update_process_bio_size(self, pid, comm, bio_size):
        """记录进程的原始bio大小"""
        if not self.show_process_stats:
            return

        proc_stats = self._get_or_create_process_stats(pid, comm)
        proc_stats.add_bio(bio_size)

    def update_split_stats(self, original_q, event_timestamp, event_bytes):
        """更新Split统计信息并返回队列信息"""
        if not original_q:
            return None

        queue_ts = original_q['ts']
        original_pid = original_q['pid']
        original_comm = original_q['comm']

        original_q['split_consumed'] = original_q.get('split_consumed', 0) + event_bytes

        if original_q['split_consumed'] >= original_q['bytes']:
            original_q['matched'] = True

        q2s = event_timestamp - queue_ts
        self.record_latency('q2s', q2s)

        # 统计split信息
        is_first_split = (original_q['split_consumed'] == event_bytes)
        self.add_split_stats(
            event_bytes,
            original_q['bytes'] if is_first_split else 0,
            1 if is_first_split else 0
        )

        return queue_ts, original_pid, original_comm


class RequestLifecycleTracker(object):
    """IO请求生命周期追踪器 - 核心事件协调"""

    def __init__(self):
        self.device_stats = {}
        self.process_stats = {}
        self.devices = {}  # dev_key -> {'Q_tree': {}, 'I_tree': {}, 'D_tree': {}}

        self.debug_enabled = False
        self.total_io_bytes = 0
        self.global_total_bytes = 0
        self.lost_events = 0

        # 事件缓存
        self._events = []
        self.event_counts = {label: 0 for label in EVENT_TYPE_LABELS.values()}

        # 辅助类 - 初始化时占位,set_mode时创建
        self.matcher = None
        self.stats_collector = None

        # 事件处理表
        self._handlers = {
            ACTION_QUEUE: self._handle_queue,
            ACTION_INSERT: self._handle_insert,
            ACTION_ISSUE: self._handle_issue,
            ACTION_COMPLETE: self._handle_complete,
            ACTION_SPLIT: self._handle_split,
        }

    def set_mode(self, show_process_stats, debug_enabled):
        self.debug_enabled = debug_enabled

        # 初始化匹配器
        self.matcher = EventMatcher(TIME_THRESHOLDS)

        # 初始化统计收集器
        self.stats_collector = StatisticsCollector(
            self.device_stats,
            self.process_stats,
            show_process_stats,
        )

    def _prepare_device(self, dev_key):
        """准备设备环境：确保设备树存在并设置为当前设备"""
        if dev_key not in self.devices:
            self.devices[dev_key] = {
                'Q_tree': defaultdict(list),
                'I_tree': defaultdict(list),
                'D_tree': defaultdict(list),
            }
        if dev_key not in self.device_stats:
            self.device_stats[dev_key] = LatencyStats()
        self.stats_collector.set_current_device(dev_key)

    # ==================== 事件处理核心逻辑 ====================

    def _handle_queue(self, event, dev_key):
        """Queue事件: 创建原始BIO记录"""
        ResultPrinter.print_raw_event(self.debug_enabled, "Q", event)
        self._prepare_device(dev_key)

        # 统计Q2Q和bio字节
        self.stats_collector.current_stats.add_q2q(event.timestamp, event.disk_name)
        self.stats_collector.current_stats.add_bio_bytes(event.bytes)

        # 创建Q事件记录
        comm = event.comm.decode('utf-8', 'replace').strip('\x00')
        q_info = {
            'ts': event.timestamp,
            'bytes': event.bytes,
            'matched': False,
            'bio_type': 'ORIGINAL',
            'pid': event.pid,
            'comm': comm,
            'split_consumed': 0,
        }

        self.devices[dev_key]['Q_tree'][event.sector].append(q_info)

        # 统计进程IO: bio计数和大小(AVG_BIO)
        self.stats_collector.update_process_bio_size(event.pid, comm, event.bytes)
        # 统计进程IO: 读写字节数(TOTAL/READ/WRITE)
        is_write = bool(event.rwflag)
        self.stats_collector.update_process_stats(event.pid, comm, event.bytes, is_write)

    def _handle_split(self, event, dev_key):
        """Split事件: 内核产生的Split BIO"""
        ResultPrinter.print_raw_event(self.debug_enabled, "SPLIT", event)
        self._prepare_device(dev_key)

        Q_tree = self.devices[dev_key]['Q_tree']

        self.stats_collector.current_stats.add_request_bytes(event.bytes)

        original_q = find_original_bio(Q_tree, event)
        result = self.stats_collector.update_split_stats(original_q, event.timestamp, event.bytes)

        if result:
            queue_ts, original_pid, original_comm = result
        else:
            queue_ts = event.timestamp
            original_pid = event.pid
            original_comm = event.comm.decode('utf-8', 'replace').strip('\x00')

        split_info = {
            'ts': queue_ts,
            'split_ts': event.timestamp,
            'bytes': event.bytes,
            'matched': False,
            'bio_type': 'SPLIT',
            'pid': original_pid,
            'comm': original_comm,
        }

        Q_tree[event.sector].append(split_info)

    def _handle_insert(self, event, dev_key):
        """Insert事件: 匹配Q树,统计Q2I或S2I"""
        ResultPrinter.print_raw_event(self.debug_enabled, "I", event)
        self._prepare_device(dev_key)

        Q_tree = self.devices[dev_key]['Q_tree']
        I_tree = self.devices[dev_key]['I_tree']

        # 统计Insert事件数，用于判断调度器模式
        self.stats_collector.current_stats.insert_count += 1

        # lambda用于将额外的上下文(event, stats_collector)注入验证函数
        matched_q, _ = find_matching_entry(
            Q_tree, event.sector,
            lambda q_info: self.matcher.match_insert_to_queue(q_info, event, self.stats_collector)
        )

        I_tree[event.sector].append({
            'ts': event.timestamp,
            'bytes': event.bytes,
            'matched': False,
            'pid': event.pid,
            'comm': event.comm.decode('utf-8', 'replace').strip('\x00'),
            'q_info': matched_q,
        })

    def _handle_issue(self, event, dev_key):
        """Issue事件: 根据调度器模式选择匹配路径（I树或Q树）"""
        ResultPrinter.print_raw_event(self.debug_enabled, "D", event)
        self._prepare_device(dev_key)

        Q_tree = self.devices[dev_key]['Q_tree']
        I_tree = self.devices[dev_key]['I_tree']
        D_tree = self.devices[dev_key]['D_tree']

        stats = self.stats_collector.current_stats

        # 根据insert_count判断调度器模式
        if stats.insert_count > 0:
            # 传统调度器路径: I → D (有Insert事件)
            matched_i, _ = find_matching_entry(
                I_tree, event.sector,
                lambda i_info: self.matcher.match_issue_to_insert(i_info, event, self.stats_collector)
            )

            if matched_i:
                D_tree[event.sector].append({
                    'ts': event.timestamp,
                    'bytes': event.bytes,
                    'matched': False,
                    'pid': matched_i['pid'],
                    'comm': matched_i['comm'],
                    'q_info': matched_i['q_info'],
                })
        else:
            # none调度器路径: Q → D (无Insert事件)
            matched_q, _ = find_matching_entry(
                Q_tree, event.sector,
                lambda q_info: self.matcher.match_issue_to_queue(q_info, event, self.stats_collector)
            )

            if matched_q:
                D_tree[event.sector].append({
                    'ts': event.timestamp,
                    'bytes': event.bytes,
                    'matched': False,
                    'pid': matched_q['pid'],
                    'comm': matched_q['comm'],
                    'q_info': matched_q,
                })

    def _handle_complete(self, event, dev_key):
        """Complete事件: 统计D2C和Q2C,清理内存"""
        ResultPrinter.print_raw_event(self.debug_enabled, "C", event)
        self._prepare_device(dev_key)

        D_tree = self.devices[dev_key]['D_tree']

        # lambda用于将额外的上下文(event)注入验证函数
        result, d_index = find_matching_entry(
            D_tree, event.sector,
            lambda d_info: self.matcher.match_complete_to_issue(d_info, event)
        )

        if result:
            matched_d, time_diff = result

            # 统计D2C
            self.stats_collector.record_latency('d2c', time_diff)

            # 统计Q2C
            q_info = matched_d.get('q_info')
            if q_info and q_info['bio_type'] in ['ORIGINAL', 'SPLIT']:
                q2c = event.timestamp - q_info['ts']
                self.stats_collector.record_latency('q2c', q2c)

            # 从D_tree中删除已匹配的条目
            stack = D_tree.get(event.sector, [])
            if stack and d_index is not None:
                stack.pop(d_index)
                if not stack:
                    del D_tree[event.sector]

        self.total_io_bytes += event.bytes
        self.global_total_bytes += event.bytes

    def process_deferred_events(self):
        """在采集结束后处理所有已缓存的事件"""
        if not self._events:
            return

        # 按时间戳排序
        self._events.sort(key=attrgetter('timestamp'))
        total_events = len(self._events)
        handlers = self._handlers

        for index, event in enumerate(self._events):
            # 显示进度
            if index > 0 and index % 100000 == 0:
                progress = (index * 100.0) / total_events
                print(f"\rProcessing... {progress:.1f}% ({index:,}/{total_events:,})",
                      end='', flush=True)

            major = event.device >> DEVICE_MINOR_SHIFT
            minor = event.device & DEVICE_MINOR_MASK
            dev_key = (major, minor)

            handler = handlers.get(event.action)
            if handler:
                handler(event, dev_key)
            elif event.action in EVENT_ACTION_NAMES:
                if dev_key not in self.device_stats:
                    self.device_stats[dev_key] = LatencyStats()
                self.stats_collector.set_current_device(dev_key)
                self.stats_collector.current_stats.add_merge(event.action)
                if self.debug_enabled:
                    action_name = EVENT_ACTION_NAMES[event.action]
                    line = (f"{action_name}: sec={event.sector} "
                            f"bytes={event.bytes}")
                    print(line)

        print(f"\rProcessing... 100.0% ({total_events:,}/{total_events:,})")
        self._events.clear()

    def store_event(self, event):
        """存储事件副本，结束时统一处理"""
        event_copy = LightBlktraceEvent.from_buffer_copy(event)
        self._events.append(event_copy)

        event_type = EVENT_TYPE_LABELS.get(event_copy.action)
        if event_type:
            self.event_counts[event_type] += 1



class ArgumentParser(object):
    """命令行参数解析和验证"""

    @staticmethod
    def parse():
        parser = argparse.ArgumentParser(description="Trace block IO latency")
        parser.add_argument(
            "-s", "--device",
            dest="device",
            help="block device path to filter",
        )
        parser.add_argument(
            "-p", "--pid",
            dest="pid",
            default=None,
            help="show process IO stats (-p 0), trace specific PIDs (-p 123,456), or search by name (-p fio)",
        )
        parser.add_argument(
            "-t", "--time",
            dest="time",
            type=ArgumentParser.positive_int,
            default=0,
            help="stop after N seconds",
        )
        parser.add_argument(
            "-d", "--debug",
            dest="debug",
            action="store_true",
            help="enable debug output",
        )
        parser.add_argument(
            "-v", "--detail",
            dest="detail",
            type=ArgumentParser.positive_int,
            default=0,
            help=(
                "detail level "
                f"({DETAIL_LEVEL_LATENCY_FIRST}: io size, "
                f"{DETAIL_LEVEL_LATENCY_SECOND}: latency histogram)"
            ),
        )
        return parser.parse_args()

    @staticmethod
    def positive_int(value):
        """验证非负整数参数"""
        ivalue = int(value)
        if ivalue < 0:
            raise argparse.ArgumentTypeError("must be non-negative")
        return ivalue

    @staticmethod
    def find_pids_by_name(process_name):
        """根据进程名搜索PID列表"""
        proc_dirs = os.listdir('/proc')
        result = []
        for pid_dir in proc_dirs:
            if not pid_dir.isdigit():
                continue
            with open(f'/proc/{pid_dir}/comm', 'r') as f:
                comm = f.read().strip()
                if process_name.lower() in comm.lower():
                    result.append(int(pid_dir))
        return result

    @staticmethod
    def parse_pid_argument(pid_arg):
        """
        解析PID参数，支持多种输入格式
        """
        # 无参数
        if pid_arg is None:
            return None, []

        # 尝试解析为数字（单个或多个PID）
        try:
            pids = [int(p.strip()) for p in str(pid_arg).split(',')]
            # -p 0 显示所有进程
            if len(pids) == 1 and pids[0] == 0:
                return "all", []
            pid_type = "pids" if len(pids) > 1 else "pid"
            return pid_type, pids
        except ValueError:
            # 解析失败，当作进程名处理
            pids = ArgumentParser.find_pids_by_name(str(pid_arg))
            if not pids:
                raise SystemExit(f"No processes found matching '{pid_arg}'")
            return "name", pids

    @staticmethod
    def setup_tracing_mode(args, pid_type, target_pids):
        """设置追踪模式和描述信息"""
        show_process_stats = (pid_type is not None)

        if pid_type is None:
            mode_desc = "latency analysis"
        elif pid_type == "all":
            mode_desc = "process IO statistics"
        elif pid_type == "name":
            mode_desc = f"process '{args.pid}' analysis (PIDs: {target_pids})"
        elif pid_type == "pids":
            mode_desc = f"PIDs {target_pids} analysis"
        else:  # pid_type == "pid"
            mode_desc = f"PID {target_pids[0]} analysis"

        return show_process_stats, mode_desc

    @staticmethod
    def setup_device_filter(args):
        def resolve_device(path):
            st = os.stat(path)
            if not stat.S_ISBLK(st.st_mode):
                raise ValueError(f"{path} is not a block device")
            return os.major(st.st_rdev), os.minor(st.st_rdev)

        """解析设备号并验证参数"""
        return resolve_device(args.device) if args.device else (0, 0)

# ==================== Main函数 ====================
def configure_filters(bpf, major, minor, target_pids=None, debug=False, detail=0):
    """
    配置BPF程序的过滤器和运行参数
    """
    # 配置全局参数数组（与内核侧cfg数组对应）
    cfg = bpf.get_table("cfg")
    cfg[ct.c_int(CFG_IDX_MAJOR)] = ct.c_uint(major)     # 设备主设备号过滤
    cfg[ct.c_int(CFG_IDX_MINOR)] = ct.c_uint(minor)     # 设备次设备号过滤
    cfg[ct.c_int(CFG_IDX_DETAIL)] = ct.c_uint(detail)   # 详细输出级别
    cfg[ct.c_int(CFG_IDX_PID_FILTER)] = ct.c_uint(1 if target_pids else 0)  # 是否启用PID过滤
    cfg[ct.c_int(CFG_IDX_DEBUG)] = ct.c_uint(1 if debug else 0)             # 是否启用调试输出

    # 初始化全局IO统计计数器（用于计算总IO量）
    total_bytes_map = bpf.get_table("total_bytes")
    total_bytes_map[ct.c_int(TOTAL_BYTES_INDEX)] = ct.c_uint64(0)

    # 如果指定了目标PID，将PID列表写入内核侧的hash map
    if target_pids:
        target_pids_table = bpf.get_table("target_pids")
        target_pids_table.clear()  # 清空旧数据
        for pid in target_pids:
            target_pids_table[ct.c_uint(pid)] = ct.c_uint(1)  # 值为1表示启用追踪


def load_bpf_program(major, minor, target_pids, pid_type, debug, detail):
    """加载并配置BPF程序"""
    script_dir = os.path.dirname(os.path.abspath(__file__))
    bpf_file = os.path.join(script_dir, "light_blktrace.c")

    with open(bpf_file) as f:
        bpf_text = f.read()

    bpf = BPF(text=bpf_text)
    pids_filter = target_pids if pid_type not in [None, "all"] else None
    configure_filters(bpf, major, minor, pids_filter, debug, detail)

    return bpf


def print_trace_header(args, mode_desc, major, minor):
    """打印追踪开始信息"""
    device_info = f"{args.device}'s data, dev num {major}-{minor}" if args.device else "all data"
    lines = [f"collecting {device_info} ({mode_desc})", "Tracing... Hit Ctrl-C to end."]

    if args.debug:
        lines.extend(["\n# Raw tracepoint events (kernel-level):",
                      "# FORMAT: [EVENT] CPU COMM PID START:SECTOR LENS:BYTES/SECTOR TIMESTAMP",
                      "#" + SEPARATOR])

    print("\n".join(lines))


def run_event_loop(bpf, tracker, args):
    """运行事件循环，采集数据"""
    start_time = time.time()
    last_status_print = start_time
    stop_flag = [False]

    def signal_handler(_signum, _frame):
        stop_flag[0] = True

    # 注册信号处理器
    signal.signal(signal.SIGINT, signal_handler)

    def handle_event(_cpu, data, _size):
        if not stop_flag[0]:
            event = ct.cast(data, ct.POINTER(LightBlktraceEvent)).contents
            tracker.store_event(event)

    # BCC框架的lost_cb需要一个接受count参数的回调函数
    bpf["events"].open_perf_buffer(handle_event, lost_cb=lambda cnt: setattr(tracker, 'lost_events', tracker.lost_events + cnt))

    while True:
        bpf.perf_buffer_poll(timeout=100)
        current_time = time.time()

        # 每5秒打印一次运行状态
        if current_time - last_status_print >= 5:
            ResultPrinter.print_tracing_status(start_time, tracker.event_counts, tracker.lost_events)
            last_status_print = current_time

        # 检查停止条件
        if (args.time and current_time - start_time >= args.time) or stop_flag[0]:
            if stop_flag[0]:
                print("\nExiting...")
            else:
                print()
            break


def finalize_and_print_results(bpf, tracker, detail, debug):
    """处理事件并打印结果"""
    tracker.process_deferred_events()

    # 在调试模式下输出BPF事件计数器统计
    if debug:
        ResultPrinter.print_bpf_event_counters(bpf)

    # 读取全局IO总量
    total_bytes_map = bpf.get_table("total_bytes")
    tracker.global_total_bytes = int(total_bytes_map[ct.c_int(TOTAL_BYTES_INDEX)].value)

    # 打印统计结果
    ResultPrinter.print_latency_stats(tracker.device_stats)
    ResultPrinter.print_process_stats(
        tracker.process_stats,
        tracker.global_total_bytes,
    )

    if detail >= DETAIL_LEVEL_LATENCY_FIRST:
        ResultPrinter.print_io_size_histogram(bpf)
    if detail >= DETAIL_LEVEL_LATENCY_SECOND:
        ResultPrinter.print_latency_histogram(bpf)


def main():
    """主函数：协调各阶段执行"""
    args = ArgumentParser.parse()

    # 阶段1: 解析参数并配置模式
    pid_type, target_pids = ArgumentParser.parse_pid_argument(args.pid)
    show_process_stats, mode_desc = ArgumentParser.setup_tracing_mode(args, pid_type, target_pids)
    major, minor = ArgumentParser.setup_device_filter(args)

    # 阶段2: 加载BPF程序并配置过滤器
    bpf = load_bpf_program(major, minor, target_pids, pid_type, args.debug, args.detail)

    # 阶段3: 创建追踪器
    tracker = RequestLifecycleTracker()
    tracker.set_mode(show_process_stats, args.debug)

    # 阶段4: 打印启动信息并运行事件循环
    print_trace_header(args, mode_desc, major, minor)
    run_event_loop(bpf, tracker, args)

    # 阶段5: 处理收集的事件并打印结果
    finalize_and_print_results(bpf, tracker, args.detail, args.debug)


if __name__ == "__main__":
    main()
