#!/usr/bin/env python3
import sqlite3
import csv
import argparse
from datetime import datetime, timedelta

def parse_args():
    parser = argparse.ArgumentParser()
    parser.add_argument("--csv", required=True, help="Original speed test CSV file")
    parser.add_argument("--kismet", required=True, help=".kismet SQLite log file")
    parser.add_argument("--output", required=True, help="Output enriched CSV file")
    return parser.parse_args()

def get_rf_metrics(cursor, bssid, channel, timestamp):
    # Match devices associated with the BSSID
    cursor.execute("""
        SELECT COUNT(*) FROM devicelink
        WHERE type = 'Wi-Fi Client' AND mac = ? AND last_time >= ? - 10 AND last_time <= ? + 10
    """, (bssid, timestamp, timestamp))
    clients_on_ap = cursor.fetchone()[0]

    # Match clients on the same channel
    cursor.execute("""
        SELECT COUNT(DISTINCT devicelink.remote_mac)
        FROM devicelink
        JOIN devices ON devicelink.mac = devices.base_mac
        WHERE devicelink.type = 'Wi-Fi Client'
          AND devices.channel = ?
          AND devicelink.last_time >= ? - 10 AND devicelink.last_time <= ? + 10
    """, (channel, timestamp, timestamp))
    clients_on_channel = cursor.fetchone()[0]

    # Count APs on the same channel
    cursor.execute("""
        SELECT COUNT(*) FROM devices
        WHERE type = 'Wi-Fi AP' AND channel = ?
          AND last_time >= ? - 10 AND last_time <= ? + 10
    """, (channel, timestamp, timestamp))
    aps_on_channel = cursor.fetchone()[0]

    congestion_score = (
        round(clients_on_channel / aps_on_channel, 2) if aps_on_channel > 0 else 0.0
    )

    return clients_on_ap, clients_on_channel, aps_on_channel, congestion_score

def main():
    args = parse_args()
    conn = sqlite3.connect(args.kismet)
    cursor = conn.cursor()

    with open(args.csv, newline="") as infile, open(args.output, "w", newline="") as outfile:
        reader = csv.reader(infile)
        writer = csv.writer(outfile)

        headers = next(reader)
        headers += ["ClientsOnAP", "ClientsOnChannel", "APsOnChannel", "CongestionScore"]
        writer.writerow(headers)

        for row in reader:
            try:
                timestamp_str = row[0]  # assuming first column is timestamp
                bssid = row[3].strip()
                channel = int(row[-4])  # assuming you store channel near the end
                timestamp = int(datetime.strptime(timestamp_str, "%Y-%m-%d_%H:%M:%S").timestamp())
                rf = get_rf_metrics(cursor, bssid, channel, timestamp)
                writer.writerow(row + list(rf))
            except Exception as e:
                print(f"[!] Failed to process row: {row}\n    Error: {e}")
                writer.writerow(row + ["ERR", "ERR", "ERR", "ERR"])

    conn.close()

if __name__ == "__main__":
    main()