import argparse
import cv2
import numpy as np
from creedsolo import RSCodec
from raptorq import Encoder

parser = argparse.ArgumentParser(formatter_class=argparse.ArgumentDefaultsHelpFormatter)
parser.add_argument("-i", "--input", help="input file")
parser.add_argument("-o", "--output", help="output video file")
parser.add_argument("-x", "--height", help="grid height", default=100, type=int)
parser.add_argument("-y", "--width", help="grid width", default=100, type=int)
parser.add_argument("-f", "--fps", help="frame rate", default=30, type=int)
parser.add_argument("-l", "--level", help="error correction level", default=0.1, type=float)
parser.add_argument("-m", "--mix", help="mix frames with original video", action="store_true")
args = parser.parse_args()


cheight = cwidth = max(args.height // 10, args.width // 10)
midwidth = args.width - 2 * cwidth
frame_size = args.height * args.width - 4 * cheight * cwidth
# Divide by 2 for 4-bit color
frame_xor = np.arange(frame_size // 2, dtype=np.uint8)
# reedsolo breaks message into 255-byte chunks
# raptorq can add up to 4 extra bytes
rs_size = frame_size // 2 - (frame_size // 2 + 254) // 255 * int(args.level * 255) - 4

with open(args.input, "rb") as f:
    data = f.read()
rsc = RSCodec(int(args.level * 255))
encoder = Encoder.with_defaults(data, rs_size)
packets = encoder.get_encoded_packets(int(len(data) / rs_size * (1 / (1 - args.level) - 1)))

# Make corners
wcorner = np.pad(np.full((cheight - 1, cwidth - 1), 0b1111), ((0, 1), (0, 1)))
rcorner = np.pad(np.full((cheight - 1, cwidth - 1), 0b0001), ((0, 1), (1, 0)))
gcorner = np.pad(np.full((cheight - 1, cwidth - 1), 0b0110), ((1, 0), (0, 1)))
bcorner = np.pad(np.full((cheight - 1, cwidth - 1), 0b1000), ((1, 0), (1, 0)))

print("Data length:", len(data))
print("Packets:", len(packets))

idx = 0


def get_frame():
    global idx
    frame = np.array(rsc.encode(bytearray(packets[idx])))
    idx = (idx + 1) % len(packets)
    frame = np.pad(frame, (0, frame_size // 2 - len(frame))) ^ frame_xor
    # Pad frame to be multiple of 255
    # frame = np.pad(frame, (0, (len(frame) + 254) // 255 * 255 - len(frame)))
    # Space out elements in each size 255 chunk
    # frame = np.ravel(frame.reshape(len(frame) // 255, 255), "F")[: frame_size // 2]
    frame = np.ravel(np.column_stack((frame >> 4, frame & 0b1111)))
    frame = np.concatenate(
        (
            np.concatenate(
                (wcorner, frame[: cheight * midwidth].reshape((cheight, midwidth)), rcorner),
                axis=1,
            ),
            frame[cheight * midwidth : frame_size - cheight * midwidth].reshape(
                (args.height - 2 * cheight, args.width)
            ),
            np.concatenate(
                (gcorner, frame[frame_size - cheight * midwidth :].reshape((cheight, midwidth)), bcorner),
                axis=1,
            ),
        )
    )
    return np.stack(((frame & 0b0001) * 255, (frame >> 1 & 0b0011) * 85, (frame >> 3) * 255), axis=-1).astype(np.uint8)


if args.mix:
    cap = cv2.VideoCapture(args.input)
    height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
    width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
    hscale = height // args.height
    wscale = width // args.width
    out = cv2.VideoWriter(args.output, cv2.VideoWriter_fourcc(*"FFV1"), args.fps, (width, height))
    while cap.isOpened():
        ret, frame = cap.read()
        if not ret:
            break
        frame = frame.astype(np.float64) / 255
        frame[: hscale * cheight, : wscale * cwidth] = 1
        frame[: hscale * cheight, wscale * (args.width - cwidth) :] = 1
        frame[hscale * (args.height - cheight) :, : wscale * cwidth] = 1
        frame[hscale * (args.height - cheight) :, wscale * (args.width - cwidth) :] = 1
        out.write(
            cv2.cvtColor(
                (frame * np.repeat(np.repeat(get_frame(), hscale, 0), wscale, 1)).astype(np.uint8),
                cv2.COLOR_RGB2BGR,
            )
        )
else:
    out = cv2.VideoWriter(args.output, cv2.VideoWriter_fourcc(*"FFV1"), args.fps, (args.width, args.height))
    for _ in packets:
        out.write(cv2.cvtColor(get_frame(), cv2.COLOR_RGB2BGR))