#!/usr/bin/env python3
"""
UE Count Plotter
----------------
Parses log lines like:
[2025-09-08 21:33:40.789676] UE RNTI=44512 ...
Groups entries into 1-second slots (by wall-clock time) and plots the number of
unique UEs (by RNTI) per slot for the entire file.

Usage:
  python ue_count_plot.py /path/to/log.txt [--save output.png] [--no-show]

Notes:
- Uses matplotlib for a single-line plot.
- Does not set any colors or styles.
- For very large logs, this streams line-by-line with O(N) memory in the number
  of distinct seconds encountered (each holds a set of RNTIs).
"""
import argparse
import datetime as dt
import re
import sys
from collections import defaultdict, OrderedDict
from typing import Dict, Set, Iterable, Tuple

import matplotlib.pyplot as plt
import matplotlib.dates as mdates


TIMESTAMP_RE = re.compile(r"\[(?P<ts>\d{4}-\d{2}-\d{2} \d{2}:\d{2}:\d{2}(?:\.\d{1,6})?)\]")
RNTI_RE = re.compile(r"\bRNTI=(?P<rnti>\d+)\b")


def parse_lines(lines: Iterable[str]) -> Dict[dt.datetime, Set[int]]:
    """
    Parse lines and bucket unique RNTIs per second.

    Returns:
        dict mapping second (datetime truncated to seconds) -> set of RNTIs
    """
    buckets: Dict[dt.datetime, Set[int]] = defaultdict(set)

    for line in lines:
        ts_m = TIMESTAMP_RE.search(line)
        rnti_m = RNTI_RE.search(line)
        if not ts_m or not rnti_m:
            # Skip lines that don't fit the schema
            continue

        ts_str = ts_m.group("ts")
        # Parse timestamp with or without fractional seconds
        # Accept both "%Y-%m-%d %H:%M:%S" and with "%f"
        if "." in ts_str:
            ts = dt.datetime.strptime(ts_str, "%Y-%m-%d %H:%M:%S.%f")
        else:
            ts = dt.datetime.strptime(ts_str, "%Y-%m-%d %H:%M:%S")

        # Truncate to whole-second wall time
        ts_minute = ts.replace(second=0, microsecond=0)
        rnti = int(rnti_m.group("rnti"))
        buckets[ts_minute].add(rnti)
    return buckets


def sort_and_fill(buckets: Dict[dt.datetime, Set[int]]) -> Tuple[list, list]:
    """
    Sorts buckets by time and fills missing seconds with zero counts.

    Returns:
        times (list[datetime]), counts (list[int])
    """
    if not buckets:
        return [], []

    # Sort keys
    sorted_times = sorted(buckets.keys())
    start = sorted_times[0]
    end = sorted_times[-1]

    all_times = []
    counts = []

    cur = start
    one_sec = dt.timedelta(minutes=1)
    while cur <= end:
        all_times.append(cur)
        counts.append(len(buckets.get(cur, set())))
        cur += one_sec

    return all_times, counts


def plot_counts(times: list, counts: list, title: str = "Unique UEs per minute"):
    """
    Render a simple matplotlib line plot (no specific colors/styles set).
    """
    fig, ax = plt.subplots(figsize=(10, 4))
    ax.plot(times, counts, linestyle="-")
    ax.set_title(title)
    ax.set_xlabel("Time (s)")
    ax.set_ylabel("Unique UEs")
    ax.xaxis.set_major_formatter(mdates.DateFormatter("%H:%M:%S"))
    fig.autofmt_xdate()
    plt.tight_layout()
    return fig, ax


def main():
    parser = argparse.ArgumentParser(description="Plot unique UEs (by RNTI) per minute from a log file.")
    parser.add_argument("--save", help="Optional path to save the plot image (e.g., output.png)")
    parser.add_argument("--no-show", action="store_true", help="Do not display the plot window")
    args = parser.parse_args()

    # Read lines
    try:
        lines = open("artifacts/mac_stats_20250909_231855.log", "r", encoding="utf-8")
    except OSError as e:
        print(f"Error opening {args.logfile}: {e}", file=sys.stderr)
        sys.exit(1)

    buckets = parse_lines(lines)
    if lines is not sys.stdin:
        lines.close()

    times, counts = sort_and_fill(buckets)
    if not times:
        print("No valid entries found to plot.", file=sys.stderr)
        sys.exit(2)

    fig, _ = plot_counts(times, counts)

    if args.save:
        try:
            fig.savefig(args.save, dpi=150, bbox_inches="tight")
            print(f"Saved plot to: {args.save}")
        except Exception as e:
            print(f"Failed to save figure: {e}", file=sys.stderr)

    if not args.no_show:
        plt.show()


if __name__ == "__main__":
    main()