import sys, os, re, glob, json, shutil
from pyspark.sql import SparkSession
from pyspark.sql import functions as f, types as t
import argparse

def read_args():
    parser = argparse.ArgumentParser(
        description="Aggregate Tstat logs into hourly user-level summaries."
    )

    parser.add_argument(
        "input_glob",
        help="Input directory glob (e.g. file:///data/2025_01_*.out/)"
    )

    parser.add_argument(
        "services_json",
        help="Path to services.json used for hostname → service classification"
    )

    parser.add_argument(
        "-o", "--output",
        default="output",
        help="Output directory for aggregated CSV files (default: output)"
    )

    args = parser.parse_args()

    return args.input_glob.rstrip("/"), args.services_json, args.output

def uri_to_local(p): 
    return p[7:] if p.startswith("file://") else p

def local_to_uri(p):
    return p if p.startswith("file://") else "file://" + p

def list_out_dirs(glob_uri):
    pat = uri_to_local(glob_uri)
    rx = re.compile(r"\d{4}_\d{2}_\d{2}_\d{2}_\d{2}\.out$")
    out = [d for d in glob.glob(pat) if os.path.isdir(d) and rx.search(os.path.basename(d))]
    out.sort()
    return [local_to_uri(d if d.endswith("/") else d + "/") for d in out]

def paths_for(glob_uri, filename):
    return [d + filename for d in list_out_dirs(glob_uri)]

def build_service_udf(spark, services_json_path):
    with open(services_json_path, "r", encoding="utf-8") as fh:
        raw = json.load(fh)

    import re as _re

    rules = []
    for cat, services in raw.items():
        for svc, pats in services.items():
            regs = [_re.compile(p, _re.IGNORECASE) for p in pats]
            rules.append((cat, svc, regs))

    @f.udf(t.StringType())
    def udf(c_tls_SNI, http_hostname, fqdn, quic_SNI):
        for host in (c_tls_SNI, http_hostname, fqdn, quic_SNI):
            h = (host or "").strip()
            if not h:
                continue
            for cat, svc, regs in rules:
                for rgx in regs:
                    if rgx.search(h):
                        return f"{cat}:{svc}"
        return "Unknown"

    return udf

def read_tstat_csv(spark, paths):
    df = spark.read.csv(paths, header=True, inferSchema=True, sep=" ")
    return df.toDF(*[c.split("#")[-1].split(":")[0] for c in df.columns])

def get_or_null(df, name):
    return f.col(name) if name in df.columns else f.lit(None)

def prepare_tcp_complete(spark, paths, udf_service):
    df = read_tstat_csv(spark, paths)
    day = f.regexp_extract(f.input_file_name(), r"(\d{4}_\d{2}_\d{2})_\d{2}_\d{2}\.out", 1)
    return (df
        .withColumn("ts", f.to_timestamp(f.from_unixtime(f.col("first").cast("double")/1000.0)))
        .withColumn("protocol", f.lit("TCP"))
        .withColumn("service", udf_service(get_or_null(df,"c_tls_SNI"), get_or_null(df,"http_hostname"), get_or_null(df,"fqdn"), f.lit(None)))
        .withColumn("day", day)
        .select("c_ip","s_bytes_uniq","c_bytes_uniq","ts","protocol","service","day"))

def prepare_udp_complete(spark, paths, udf_service):
    df = read_tstat_csv(spark, paths)
    udp_ts = f.coalesce(get_or_null(df,"c_first_abs"), get_or_null(df,"s_first_abs"))
    day = f.regexp_extract(f.input_file_name(), r"(\d{4}_\d{2}_\d{2})_\d{2}_\d{2}\.out", 1)
    return (df
        .withColumn("ts", f.to_timestamp(f.from_unixtime(udp_ts.cast("double")/1000.0)))
        .withColumn("protocol", f.lit("UDP"))
        .withColumn("service", udf_service(f.lit(None), f.lit(None), get_or_null(df,"fqdn"), get_or_null(df,"quic_SNI")))
        .withColumn("day", day)
        .withColumnRenamed("s_bytes_all","s_bytes_uniq")
        .withColumnRenamed("c_bytes_all","c_bytes_uniq")
        .select("c_ip","s_bytes_uniq","c_bytes_uniq","ts","protocol","service","day"))

def prepare_tcp_nocomplete(spark, paths):
    df = read_tstat_csv(spark, paths)
    day = f.regexp_extract(f.input_file_name(), r"(\d{4}_\d{2}_\d{2})_\d{2}_\d{2}\.out", 1)
    return (df
        .withColumn("ts", f.to_timestamp(f.from_unixtime(f.col("first").cast("double")/1000.0)))
        .withColumn("protocol", f.lit("TCP_NC"))
        .withColumn("service", f.lit("Unknown"))
        .withColumn("day", day)
        .select("c_ip","s_bytes_uniq","c_bytes_uniq","ts","protocol","service","day"))

def add_hour_bin(df):
    return (df.withColumn("hour_start", f.date_trunc("hour", f.col("ts")))
              .withColumn("hour_bin",
                          f.concat_ws("-", f.date_format(f.col("hour_start"), "HH"),
                                      f.date_format(f.expr("hour_start + INTERVAL 1 HOUR"), "HH")))
           )

def aggregate(df):
    return (df.groupBy("day","c_ip","hour_bin","service","protocol")
              .agg(
                  f.sum(f.coalesce(f.col("s_bytes_uniq"), f.lit(0))).alias("bytes_down"),
                  f.sum(f.coalesce(f.col("c_bytes_uniq"), f.lit(0))).alias("bytes_up"),
                  f.count(f.lit(1)).alias("flows")
              )
           )

def export_csv_per_day(df, output_root):
    days = [r["day"] for r in df.select("day").distinct().collect()]
    for day in days:
        out_day = df.filter(f.col("day") == day)
        base_dir = os.path.join(output_root, day)
        os.makedirs(base_dir, exist_ok=True)
        tmp_dir = os.path.join(base_dir, "_tmp")
        final_csv = os.path.join(base_dir, "aggregation.csv")

        (out_day.coalesce(1)
            .write
            .mode("overwrite")
            .option("header", True)
            .csv(tmp_dir))

        part = glob.glob(os.path.join(tmp_dir, "part-*.csv"))
        if not part:
            shutil.rmtree(tmp_dir, ignore_errors=True)
            raise FileNotFoundError(f"No CSV generated for day {day}")

        if os.path.exists(final_csv):
            os.remove(final_csv)

        shutil.move(part[0], final_csv)
        shutil.rmtree(tmp_dir, ignore_errors=True)
        print(final_csv)

def main():
    glob_uri, services_json, output_root = read_args()
    os.environ["PYSPARK_PYTHON"] = sys.executable
    os.environ["PYSPARK_DRIVER_PYTHON"] = sys.executable
    spark = (SparkSession.builder
             .appName("TSTAT-Day-Aggregation")
             .master("local[*]")
             .config("spark.sql.execution.arrow.pyspark.enabled", "true")
             .config("spark.sql.shuffle.partitions", "128")
             .config("spark.driver.memory", "12g")
             .config("spark.executor.memory", "12g")
             .config("spark.local.dir", "/tmp/spark") 
             .getOrCreate())
    udf_service = build_service_udf(spark, services_json)
    tcp_paths  = paths_for(glob_uri, "log_tcp_complete.gz")
    udp_paths  = paths_for(glob_uri, "log_udp_complete.gz")
    nc_paths   = paths_for(glob_uri, "log_tcp_nocomplete.gz")
    if not (tcp_paths or udp_paths or nc_paths):
        print(f"No input files under {glob_uri}")
        sys.exit(1)
    dfs = []
    if tcp_paths:
        dfs.append(prepare_tcp_complete(spark, tcp_paths, udf_service))
    if udp_paths:
        dfs.append(prepare_udp_complete(spark, udp_paths, udf_service))
    if nc_paths:
        dfs.append(prepare_tcp_nocomplete(spark, nc_paths))
    unified = dfs[0]
    for d in dfs[1:]:
        unified = unified.unionByName(d, allowMissingColumns=True)
    binned = add_hour_bin(unified)
    daily_agg = aggregate(binned)
    export_csv_per_day(daily_agg, output_root)
    
if __name__ == "__main__":
    main()