#!/usr/bin/env python3
import argparse
from pathlib import Path
import os
import subprocess
from io import BytesIO
import numpy as np
from tqdm import tqdm
from datetime import datetime
import pandas as pd
import pcapng
from struct import unpack
from PIL import Image


# Create the parser
parser = argparse.ArgumentParser(description="Process a pcap file.")
parser.add_argument("--live", action="store_true", help="Process images live")

# Add an argument for the pcap file, with a default value
parser.add_argument('input_file', nargs='?', default='in.pcap', help='The pcap file to process')

# Parse the arguments
args = parser.parse_args()


# TODO - probably a better way to do this
def live_capture():
    while True:
        # TODO: This is wrong:
        # during scapy.all.sniff construction we lose about 27 packets.
        # we work around it by capturing 2* packets(frame), so we just lose
        # about half the frames.
        cap = scapy.all.sniff(iface="enp1s0f0", count=32 * 2)
        for pkt in cap:
            l = len(bytes(pkt))
            if l != 6972:
                continue
            data = bytes(pkt)
            with open('udp.bytes', 'ab+') as fd:
                fd.write(bytes(pkt.payload))
            yield data


if args.live:
    print('live stream, import scapy')
    import scapy.all
    print('open stream')
    stream = live_capture()
    start = datetime.now()
    timestamp = start.strftime('%Y%m%d_%H%M%S')
    basename = f'live_{timestamp}'
    blocks = tqdm(stream)
else:
    # Now use args.input_file as the file to process
    input_file = args.input_file
    basename = os.path.splitext(os.path.basename(input_file))[0]
    stream = open(input_file, 'rb')
    # Read packets from a pcap file
    scanner = pcapng.scanner.FileScanner(stream)
    blocks = tqdm(scanner)


def rightsize(it):
    for i, obj in enumerate(it):
        if isinstance(obj, bytes):
            l = len(obj)
            data = obj
        else:
            if not hasattr(obj, 'packet_len'):
                continue
            l = obj.packet_len
            data = obj.packet_data
        if l != 6972:
            continue
        yield data


def removestart(it):
    "Remove the UDP header from the packets"
    for x in it:
        yield x[0x2A:]


# Function to parse packet data
def parse(data):
    hdr = 4 + 2 * 7  # Header length
    # Unpack data into variables
    c1, c2, part, a, ffaa, b, c, d = unpack(">Lhhhhhhh", data[:hdr])
    ret = locals()
    del ret["data"]
    del ret["hdr"]
    ret["data"] = data[hdr:]
    return ret


def parsed(it):
    for x in it:
        yield parse(x)


# Function to group data into frames
def frames(it):
    current = []
    #otherdata = []
    for obj in it:
        if obj['part'] == 0:
            if len(current) > 0:
                yield b"".join(current)
            current = []
            #otherdata = []
        current.append(obj["data"])
        #otherdata.append(obj)
    if len(current) > 0:
        yield b"".join(current)

WIDTH = 384
HEIGHT = 288


def skip_bad_frames(it, width=WIDTH, height=HEIGHT):
    for frame in it:
        if len(frame) != width * height * 2:  # 16 bpp
            # Will be fixed when we stopped doing restarts
            #print(f'{len(frame)} != {width} * {height} * 2')
            continue
        yield frame


def iterimages(it, width=WIDTH, height=HEIGHT, pixelformat=">H"):
    for frame in it:
        yield Image.fromarray(np.frombuffer(frame, dtype=pixelformat).reshape(width, height))


# Get frames and convert them to images
frames = skip_bad_frames(frames(parsed(removestart(rightsize(blocks)))))

# Create the directory for frames if not exists
frame_dir = f"frames/{basename}"
if not os.path.exists(frame_dir):
    os.makedirs(frame_dir)

# Save each image as a PNG file
if not args.live:
    images = iterimages(it=frames)
    for i, img in enumerate(images):
        img.save(f'frames/{basename}/{basename}_{i:04}.png')
    ffmpeg_input = f"frames/{basename}/{basename}_%04d.png"
    command = [
        "ffmpeg", 
        "-y",                  # Overwrite output file without asking
        "-hide_banner",        # Hide banner
        "-loglevel", "info",   # Log level
        "-f", "image2",        # Input format
        "-framerate", "25",    # Framerate
        "-i", ffmpeg_input,     # Input file pattern
        "-vf", "transpose=1",  # Video filter for transposing
        "-s", "384x288",       # Size of one frame
        "-vcodec", "libopenh264",  # Video codec
        "-pix_fmt", "yuv420p", # Pixel format: YUV 4:2:0
        "thermal.mp4",         # Output file in MP4 container
    ]

    subprocess.run(command)
    print("to play: ffplay thermal.mp4")

else:
    # TODO: to video via ffmpeg; right now just a single png
    # of the last frame
    def todo_live_ffmpeg():
        output = 'to_ffmpeg'
        # live: write to named pipe
        if not Path(output).exists():
            print(f'making fifo at {output}')
            os.mkfifo(output)
        fd = open(output, 'wb')
        for frame in frames:
            fd.write(frame)

    for frame in frames:
        print('.')
        Image.fromarray(np.frombuffer(frame, dtype='>H').reshape(WIDTH, HEIGHT)).save(f'live.png')