#!/usr/bin/env python3
"""MP-QUIC Application & eBPF Results Visualizer.

This script takes the CSV outputs and plots their latency distributions and timelines.

Usage:
    python visualize.py --app ../server/app_metrics.csv
    python visualize.py --client ../results/client.csv --server ../results/server.csv
"""

import argparse
import pandas as pd
import matplotlib.pyplot as plt
import seaborn as sns
import os

def load_ebpf_data(csv_path, label):
    if not os.path.exists(csv_path):
        print(f"Warning: File {csv_path} not found.")
        return pd.DataFrame()
    
    df = pd.read_csv(csv_path)
    df['node'] = label
    df = df.sort_values('timestamp')
    if not df.empty:
        first_event_time = df['timestamp'].iloc[0]
        df['rel_time'] = df['timestamp'] - first_event_time
    else:
        df['rel_time'] = 0.0
    df['value_us'] = df['value_ns'] / 1000.0
    return df

def plot_app_metrics(csv_path, output_dir):
    """Plot Application-Level Metrics (Drone -> Ground Station)."""
    if not os.path.exists(csv_path):
        print(f"App metrics file {csv_path} not found.")
        return

    df = pd.read_csv(csv_path)
    if df.empty:
        print("App metrics file is empty.")
        return

    df = df.sort_values('sequence_number')
    
    # Calculate Jitter and Rel Latency
    df['latency_ms'] = df['latency_ns'] / 1000000.0
    
    # Relative time from first packet received
    df['rel_time'] = (df['ground_recv_time'] - df['ground_recv_time'].min()) / 1e9

    plt.figure(figsize=(10, 5))
    sns.scatterplot(x='rel_time', y='latency_ms', data=df, s=15, alpha=0.6)
    plt.title('App-Level End-to-End Latency Over Time (Glass-to-Glass)')
    plt.ylabel('Relative Latency (ms)')
    plt.xlabel('Time (s)')
    out_path = os.path.join(output_dir, 'app_latency_timeline.png')
    plt.savefig(out_path, dpi=300, bbox_inches='tight')
    plt.close()
    
    # Packet Loss Calculation
    max_seq = df['sequence_number'].max()
    min_seq = df['sequence_number'].min()
    expected_packets = max_seq - min_seq + 1
    received_packets = len(df)
    lost_packets = expected_packets - received_packets
    reliability = (received_packets / expected_packets) * 100 if expected_packets > 0 else 0

    print("\n" + "=" * 55)
    print("  🏆 APP-LEVEL RELIABILITY (Drone -> Ground)")
    print("=" * 55)
    print(f"  Packets Sent (Expected): {expected_packets}")
    print(f"  Packets Received:        {received_packets}")
    print(f"  Packets Lost:            {lost_packets}")
    print(f"  Reliability:             {reliability:.5f}%")
    print("=" * 55)
    
    print(f"Saved app-level plot to {out_path}")


def plot_latency_distributions(df, output_dir):
    latency_df = df[df['event_type'].isin(['SEND_LATENCY', 'RECV_LATENCY'])]
    if latency_df.empty: return

    plt.figure(figsize=(10, 6))
    sns.violinplot(x='event_type', y='value_us', hue='node', data=latency_df, split=True, inner="quartile")
    plt.yscale('log')
    plt.title('Kernel Network Stack Latency Distribution (Log Scale)')
    plt.ylabel('Latency (µs)')
    plt.xlabel('Event Type')
    
    out_path = os.path.join(output_dir, 'kernel_latency_distribution.png')
    plt.savefig(out_path, dpi=300, bbox_inches='tight')
    plt.close()
    print(f"Saved kernel distribution plot to {out_path}")

def plot_latency_timeline(df, output_dir):
    latency_df = df[df['event_type'].isin(['SEND_LATENCY', 'RECV_LATENCY'])]
    if latency_df.empty: return

    g = sns.FacetGrid(latency_df, col="event_type", row="node", margin_titles=True, height=4, aspect=2)
    g.map(sns.scatterplot, "rel_time", "value_us", alpha=0.5, s=10)
    g.set_axis_labels("Time (s)", "Latency (µs)")
    g.set_titles(col_template="{col_name}", row_template="{row_name}")
    
    for ax in g.axes.flat:
        ax.set_yscale('log')
        
    g.fig.suptitle('Kernel Latency Over Time', y=1.02)
    
    out_path = os.path.join(output_dir, 'kernel_latency_timeline.png')
    plt.savefig(out_path, dpi=300, bbox_inches='tight')
    plt.close()
    print(f"Saved kernel timeline plot to {out_path}")

def main():
    parser = argparse.ArgumentParser(description="Visualize MP-QUIC data")
    parser.add_argument("--client", help="Path to client eBPF CSV file")
    parser.add_argument("--server", help="Path to server eBPF CSV file")
    parser.add_argument("--app", help="Path to App-Level metrics CSV file (e.g. app_metrics.csv)")
    parser.add_argument("--outdir", default="plots", help="Directory to save plots")
    args = parser.parse_args()

    os.makedirs(args.outdir, exist_ok=True)

    if args.app:
        plot_app_metrics(args.app, args.outdir)

    dfs = []
    if args.client:
        dfs.append(load_ebpf_data(args.client, "Client"))
    if args.server:
        dfs.append(load_ebpf_data(args.server, "Server"))
        
    if dfs:
        df = pd.concat(dfs, ignore_index=True)
        sns.set_theme(style="whitegrid")
        plot_latency_distributions(df, args.outdir)
        plot_latency_timeline(df, args.outdir)
        
    print("\nVisualization complete! Check the '{}' directory.".format(args.outdir))

if __name__ == "__main__":
    main()