import argparse
import sys
import numpy as np
from creedsolo import RSCodec
from PyQt6.QtWidgets import QApplication, QWidget, QLabel, QVBoxLayout
from PyQt6.QtGui import QPixmap
from PyQt6.QtCore import QTimer
from PIL import Image, ImageQt
from raptorq import Encoder

parser = argparse.ArgumentParser(formatter_class=argparse.ArgumentDefaultsHelpFormatter)
parser.add_argument("file", help="output file for decoded data")
parser.add_argument("--height", help="grid height", default=100, type=int)
parser.add_argument("--width", help="grid width", default=100, type=int)
parser.add_argument("--fps", help="framerate", default=30, type=int)
parser.add_argument("--level", help="error correction level", default=0.1, type=float)
args = parser.parse_args()


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

with open(args.file, "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 * args.level))
print("Data length:", len(data))
print("Packets:", len(packets))
input("Seizure warning!")


class EncoderWidget(QWidget):
    def __init__(self):
        super().__init__()
        self.timer = QTimer(self)
        self.timer.timeout.connect(self.update)
        self.timer.start(1000 // args.fps)
        self.idx = 0
        self.label = QLabel(self)
        layout = QVBoxLayout(self)
        layout.addWidget(self.label)
        layout.setContentsMargins(0, 0, 0, 0)
        self.setLayout(layout)
        self.showFullScreen()

    def update(self):
        frame_data = np.array(rsc.encode(packets[self.idx]))
        # Pad frame to fit frame_size since raptorq might not add 4 bytes
        frame_data = np.pad(frame_data, (0, frame_size - len(frame_data))) ^ frame_xor
        self.idx = (self.idx + 1) % len(packets)
        frame = np.concatenate(
            (
                np.concatenate(
                    (
                        np.zeros((cheight, cwidth), dtype=np.uint8),  # TODO
                        frame_data[: cheight * midwidth].reshape((cheight, midwidth)),
                        np.zeros((cheight, cwidth), dtype=np.uint8),  # TODO
                    ),
                    axis=1,
                ),
                frame_data[
                    cheight * midwidth : frame_size - cheight * midwidth
                ].reshape((args.height - 2 * cheight, args.width)),
                np.concatenate(
                    (
                        np.zeros((cheight, cwidth), dtype=np.uint8),  # TODO
                        frame_data[frame_size - cheight * midwidth :].reshape(
                            (cheight, midwidth)
                        ),
                        np.zeros((cheight, cwidth), dtype=np.uint8),  # TODO
                    ),
                    axis=1,
                ),
            )
        )
        color_frame = np.stack(
            (frame << 5 & 0b11100000, frame << 2 & 0b11100000, frame & 0b11000000),
            axis=-1,
        )
        img = Image.fromarray(color_frame)
        qt_img = ImageQt.ImageQt(img)
        pixmap = QPixmap.fromImage(qt_img).scaled(self.size())
        self.label.setPixmap(pixmap)


app = QApplication([])
widget = EncoderWidget()
sys.exit(app.exec())