#!/usr/bin/python3
"""
@file ptpanalyzer.py
@copyright ESO - European Southern Observatory
@defgroup meltcs	
@brief IFTEST - Subsystem Interface Testing
Tool to inspect PTP PCAP files
"""
import pyshark
import csv
import matplotlib.pyplot as plt
import argparse
from datetime import datetime, timedelta, tzinfo

class UTC(tzinfo):
    def utcoffset(self, dt): return timedelta(0)
    def tzname(self, dt): return "UTC"
    def dst(self, dt): return timedelta(0)

utc = UTC()

def parse_ptp_timestamp(seconds_str, nanoseconds_str):
    try:
        #print(f"  >> Parsing PTP timestamp: seconds={seconds_str}, nanos={nanoseconds_str}")
        if seconds_str is None or nanoseconds_str is None:
            raise ValueError("Missing seconds or nanoseconds")
        seconds = int(seconds_str)
        nanoseconds = int(nanoseconds_str)
        timestamp = seconds + nanoseconds / 1_000_000_000
        #print(f"    → Full timestamp (float): {timestamp}")
        return datetime.fromtimestamp(timestamp, tz=utc)
    except Exception as e:
        print(f"  !! Failed to parse PTP timestamp: {e}")
        return None


class PTPSessionAnalyzer:
    def __init__(self):
        self.sync_packets = []
        self.follow_up_packets = []
        self.delay_req_packets = []
        self.delay_resp_packets = []
        self.grandmaster_ids = set()
        self.sync_seq_ids = set()
        self.followup_seq_ids = set()
        self.follow_up_map = {}

    def analyze_packet(self, pkt):
        try:
            if not hasattr(pkt, 'ptp'):
                return

            ptp = pkt.ptp
            hw_timestamp = datetime.fromtimestamp(float(pkt.sniff_timestamp), tz=utc)
            msg_type = int(getattr(ptp, 'v2_messagetype', '0'), 16)
            seq_id = int(getattr(ptp, 'v2_sequenceid', '0'))
            grandmaster_id = getattr(ptp, 'v2_grandmasteridentity', 'unknown')

            if grandmaster_id != 'unknown':
                self.grandmaster_ids.add(grandmaster_id)

            if msg_type == 0x00:
                self.sync_packets.append((seq_id, hw_timestamp, int(pkt.number)))
                self.sync_seq_ids.add(seq_id)

            elif msg_type == 0x01:
                self.delay_req_packets.append((hw_timestamp, None, int(pkt.number)))

            elif msg_type == 0x08:
                seconds = getattr(ptp, 'v2_fu.preciseorigintimestamp.seconds', None)
                nanos = getattr(ptp, 'v2_fu.preciseorigintimestamp.nanoseconds', None)
                precise_ts = parse_ptp_timestamp(seconds, nanos)
                correction = float(getattr(ptp, 'v2_correction', '0'))
                self.follow_up_packets.append((seq_id, hw_timestamp, precise_ts, correction, int(pkt.number)))
                self.follow_up_map[seq_id] = (precise_ts, int(pkt.number))
                self.followup_seq_ids.add(seq_id)

            elif msg_type == 0x09:
                seconds = getattr(ptp, 'v2_dr_receivetimestamp_seconds', None)
                nanos = getattr(ptp, 'v2_dr_receivetimestamp_nanoseconds', None)
                recv_ts = parse_ptp_timestamp(seconds, nanos)
                #print(f"[Debug] Added Delay_Resp: Pkt {pkt.number} | HW: {hw_timestamp} | PTP: {recv_ts}")
                self.delay_resp_packets.append((None, recv_ts, int(pkt.number)))

        except Exception:
            pass

    def check_missing_seq(self):
        if not self.sync_seq_ids or not self.followup_seq_ids:
            return []
        combined = self.sync_seq_ids.union(self.followup_seq_ids)
        return [seq for seq in range(min(combined), max(combined) + 1) if seq not in self.sync_seq_ids or seq not in self.followup_seq_ids]

    def check_time_alignment(self):
        for seq_id, sync_hw, pkt_num in self.sync_packets:
            followup = self.follow_up_map.get(seq_id)
            if followup:
                followup_ts, _ = followup
                if followup_ts is None:
                    continue
                delta = abs((sync_hw - followup_ts).total_seconds())
                micro_delta = delta * 1_000_000
                status = "OK" if micro_delta <= 14 else "FAIL"
                print(f"Seq {seq_id} | Sync HW: {sync_hw} | FollowUp: {followup_ts} | Δ: {micro_delta:.2f}µs | {status}")

        for (dr_hw, _, dr_num), (_, dp_ptp, dp_num) in zip(self.delay_req_packets, self.delay_resp_packets):
            if dr_hw and dp_ptp:
                delta = abs((dp_ptp - dr_hw).total_seconds())
                micro_delta = delta * 1_000_000
                status = "OK" if micro_delta <= 7 else "FAIL"
                print(f"Delay_Req/Resp Pkt #{dr_num}->{dp_num} | Req HW: {dr_hw} | Resp PTP: {dp_ptp} | Δ: {micro_delta:.2f}µs | {status}")

    def export_to_csv(self, filename):
        with open(filename, 'w', newline='') as csvfile:
            writer = csv.writer(csvfile)
            writer.writerow(['Type', 'PacketNumber', 'SequenceID', 'HW Timestamp', 'PTP Timestamp', 'Delta (s)', 'Correction (ns)'])

            for seq_id, hw, num in self.sync_packets:
                writer.writerow(['Sync', num, seq_id, hw.isoformat(), '', '', ''])

            for seq_id, hw, precise, corr, num in self.follow_up_packets:
                if precise is not None:
                    delta = (hw - precise).total_seconds()
                    writer.writerow(['Follow_Up', num, seq_id, hw.isoformat(), precise.isoformat(), delta, corr])

            for (dr_hw, _, dr_num), (_, dp_ptp, dp_num) in zip(self.delay_req_packets, self.delay_resp_packets):
                if dr_hw and dp_ptp:
                    delta = (dp_ptp - dr_hw).total_seconds()
                    writer.writerow(['Delay_Req->Resp', f"{dr_num}->{dp_num}", '', dr_hw.isoformat(), dp_ptp.isoformat(), delta, ''])

    def plot_deltas(self):
        x_fu = [(num, (hw - precise).total_seconds()) for _, hw, precise, _, num in self.follow_up_packets if precise is not None]
        x_sync = [(num, hw.timestamp()) for _, hw, num in self.sync_packets]
        x_req_resp = [(dr_num, (dp_ptp - dr_hw).total_seconds()) for (dr_hw, _, dr_num), (_, dp_ptp, _) in zip(self.delay_req_packets, self.delay_resp_packets) if dr_hw and dp_ptp]

        plt.figure(figsize=(10, 6))
        if x_fu:
            x, y = zip(*x_fu)
            plt.plot(x, y, label='Follow_Up')
        if x_req_resp:
            x, y = zip(*x_req_resp)
            plt.plot(x, y, label='Delay_Req->Resp')
        if x_sync:
            x, y = zip(*x_sync)
            plt.plot(x, y, label='Sync Timestamps')

        plt.title('PTP Timing Data')
        plt.xlabel('Packet Number')
        plt.ylabel('Time / Offset (seconds)')
        plt.legend()
        plt.grid(True)
        plt.tight_layout()
        plt.show()

    def summary(self):
        print("=== PTP Session Summary ===")
        print(f"Total Sync packets: {len(self.sync_packets)}")
        print(f"Total Follow_Up packets: {len(self.follow_up_packets)}")
        print(f"Total Delay_Req packets: {len(self.delay_req_packets)}")
        print(f"Total Delay_Resp packets: {len(self.delay_resp_packets)}")
        print(f"Detected Grandmaster IDs: {', '.join(self.grandmaster_ids)}")

        missing_seq = self.check_missing_seq()
        
        self.check_time_alignment()
        if missing_seq:
            print(f"\nMissing Sequence Numbers (Sync <-> Follow_Up): {missing_seq}")
        else:
            print("\nNo missing sequence numbers between Sync and Follow_Up.")


def analyze_pcap(pcap_path, csv_output_path=None, plot=False):
    cap = pyshark.FileCapture(pcap_path)
    analyzer = PTPSessionAnalyzer()

    for pkt in cap:
        analyzer.analyze_packet(pkt)

    cap.close()
    analyzer.summary()

    print("\n=== Queue Sizes ===")
    print(f"Sync packets: {len(analyzer.sync_packets)}")
    print(f"Follow_Up packets: {len(analyzer.follow_up_packets)}")
    print(f"Delay_Req packets: {len(analyzer.delay_req_packets)}")
    print(f"Delay_Resp packets: {len(analyzer.delay_resp_packets)}")

    if csv_output_path:
        analyzer.export_to_csv(csv_output_path)

    if plot:
        analyzer.plot_deltas()

def main():
    parser = argparse.ArgumentParser(description="Stateful PTP Packet Analyzer")
    parser.add_argument("pcap_file", help="Path to the PCAPNG file")
    parser.add_argument("--csv", help="Output CSV file path", default=None)
    parser.add_argument("--plot", help="Plot timing offsets", action="store_true")
    args = parser.parse_args()
    analyze_pcap(args.pcap_file, csv_output_path=args.csv, plot=args.plot)

if __name__ == "__main__":
    main()
