from __future__ import annotations

import struct
import sys
from pathlib import Path

import cv2
import numpy as np
import soundfile as sf
import torch

from encodec import EncodecModel


# ============================================================
# CONFIGURATION
# ============================================================

RAW_OUTPUT = Path("received.raw")
WAV_OUTPUT = Path("reconstructed.wav")

BIT_DEPTH = 10
MAGIC = b"OCS1"

BW_TO_NQ = {
    1500: 2,
    3000: 4,
    6000: 8,
    12000: 16,
    24000: 32,
}


# ============================================================
# BIT UNPACKING
# ============================================================

def unpack_bits(
    data: bytes,
    bits_per_value: int = 10,
    expected_values: int | None = None,
) -> np.ndarray:

    values = []

    bitstream = 0
    bitcount = 0

    mask = (1 << bits_per_value) - 1

    for byte in data:

        bitstream |= byte << bitcount
        bitcount += 8

        while bitcount >= bits_per_value:

            values.append(
                bitstream & mask
            )

            bitstream >>= bits_per_value
            bitcount -= bits_per_value

            if (
                expected_values is not None
                and len(values) >= expected_values
            ):
                return np.array(
                    values,
                    dtype=np.uint16,
                )

    return np.array(
        values,
        dtype=np.uint16,
    )


# ============================================================
# HEADER PARSING
# ============================================================

def parse_header(blob: bytes):

    offset = 0

    if len(blob) < 20:
        raise ValueError(
            "File is too small to contain a valid header."
        )

    magic = blob[
        offset:
        offset + 4
    ]

    offset += 4

    if magic != MAGIC:
        raise ValueError(
            "Invalid music file: bad magic."
        )

    version = struct.unpack_from(
        "<H",
        blob,
        offset,
    )[0]

    offset += 2

    bandwidth_milli = struct.unpack_from(
        "<H",
        blob,
        offset,
    )[0]

    offset += 2

    sample_rate = struct.unpack_from(
        "<I",
        blob,
        offset,
    )[0]

    offset += 4

    n_q = struct.unpack_from(
        "<H",
        blob,
        offset,
    )[0]

    offset += 2

    num_frames = struct.unpack_from(
        "<I",
        blob,
        offset,
    )[0]

    offset += 4

    total_tokens = struct.unpack_from(
        "<I",
        blob,
        offset,
    )[0]

    offset += 4

    frame_lengths = []

    for _ in range(num_frames):

        frame_length = struct.unpack_from(
            "<I",
            blob,
            offset,
        )[0]

        offset += 4

        frame_lengths.append(
            int(frame_length)
        )

    return (
        version,
        bandwidth_milli,
        sample_rate,
        n_q,
        num_frames,
        total_tokens,
        frame_lengths,
        offset,
    )


# ============================================================
# ENCODEC DECODER
# ============================================================

def decode_music(
    input_path: Path,
    output_path: Path,
):

    print()
    print("=" * 60)
    print("DECODING MUSIC")
    print("=" * 60)

    blob = input_path.read_bytes()

    print(
        f"Input size: {len(blob):,} bytes"
    )

    (
        version,
        bandwidth_milli,
        sample_rate,
        n_q,
        num_frames,
        total_tokens,
        frame_lengths,
        payload_offset,
    ) = parse_header(blob)

    print(
        f"Version:     {version}"
    )

    print(
        f"Bandwidth:   {bandwidth_milli / 1000:.3f} kbps"
    )

    print(
        f"Sample rate: {sample_rate} Hz"
    )

    print(
        f"Codebooks:   {n_q}"
    )

    print(
        f"Frames:      {num_frames}"
    )

    print(
        f"Tokens:      {total_tokens}"
    )

    if bandwidth_milli not in BW_TO_NQ:
        raise ValueError(
            f"Unsupported bandwidth: "
            f"{bandwidth_milli / 1000:.3f} kbps"
        )

    expected_n_q = BW_TO_NQ[
        bandwidth_milli
    ]

    if n_q != expected_n_q:
        raise ValueError(
            "Number of codebooks does not "
            "match the bandwidth."
        )

    payload = blob[
        payload_offset:
    ]

    tokens = unpack_bits(
        payload,
        bits_per_value=BIT_DEPTH,
        expected_values=total_tokens,
    )

    if len(tokens) < total_tokens:
        raise ValueError(
            "Token stream is truncated."
        )

    tokens = tokens[
        :total_tokens
    ]

    # --------------------------------------------------------
    # Reconstruct EnCodec frames
    # --------------------------------------------------------

    encoded_frames = []

    index = 0

    for frame_length in frame_lengths:

        token_count = (
            n_q * frame_length
        )

        frame_tokens = tokens[
            index:
            index + token_count
        ]

        index += token_count

        if len(frame_tokens) != token_count:
            raise ValueError(
                "Not enough tokens to rebuild frame."
            )

        codes = torch.tensor(
            frame_tokens,
            dtype=torch.long,
        ).reshape(
            1,
            n_q,
            frame_length,
        )

        encoded_frames.append(
            (codes, None)
        )

    # --------------------------------------------------------
    # Decode
    # --------------------------------------------------------

    print(
        "Loading EnCodec..."
    )

    model = (
        EncodecModel
        .encodec_model_24khz()
    )

    model.set_target_bandwidth(
        bandwidth_milli / 1000.0
    )

    model.eval()

    print(
        "Reconstructing audio..."
    )

    with torch.inference_mode():

        decoded = model.decode(
            encoded_frames
        )

    decoded = (
        decoded
        .squeeze(0)
        .cpu()
        .numpy()
        .T
    )

    sf.write(
        output_path,
        decoded,
        sample_rate,
    )

    duration = (
        len(decoded) /
        sample_rate
    )

    print(
        f"Saved: {output_path}"
    )

    print(
        f"Duration: {duration:.2f} seconds"
    )


# ============================================================
# QR CODE DECODING
# ============================================================

def decode_qr_image(
    image_path: Path,
):

    image = cv2.imread(
        str(image_path)
    )

    if image is None:
        raise FileNotFoundError(
            f"Could not open image: {image_path}"
        )

    detector = cv2.QRCodeDetector()

    results = []

    # Newer OpenCV versions provide byte-oriented
    # multi-QR decoding, which is useful here because
    # our QR payload is arbitrary binary data.

    try:

        (
            ok,
            decoded_info,
            points,
            _,
        ) = detector.detectAndDecodeBytesMulti(
            image
        )

        if ok:

            for data in decoded_info:

                if data is None:
                    continue

                if isinstance(data, str):
                    data = data.encode(
                        "latin1"
                    )

                data = bytes(data)

                if len(data) < 1:
                    continue

                qr_id = data[0]

                if not 1 <= qr_id <= 8:
                    continue

                results.append(
                    (
                        qr_id,
                        data[1:],
                    )
                )

    except AttributeError:

        raise RuntimeError(
            "Your OpenCV version does not "
            "support byte-oriented multi-QR "
            "decoding. Install a recent "
            "opencv-python version."
        )

    return results


# ============================================================
# READ SIDE A + SIDE B
# ============================================================

def read_qr_cassette():

    print()
    print("=" * 60)
    print("QR MUSIC CASSETTE")
    print("=" * 60)

    print()
    print(
        "Enter the path to the Side A image."
    )

    side_a = Path(
        input(
            "Side A: "
        ).strip().strip('"')
    )

    print()
    print(
        "Enter the path to the Side B image."
    )

    side_b = Path(
        input(
            "Side B: "
        ).strip().strip('"')
    )

    if not side_a.exists():
        raise FileNotFoundError(
            f"Side A not found: {side_a}"
        )

    if not side_b.exists():
        raise FileNotFoundError(
            f"Side B not found: {side_b}"
        )

    print()
    print(
        "Scanning Side A..."
    )

    side_a_codes = decode_qr_image(
        side_a
    )

    print(
        f"Found {len(side_a_codes)} QR codes."
    )

    for qr_id, data in side_a_codes:
        print(
            f"  QR {qr_id}: "
            f"{len(data)} bytes"
        )

    print()
    print(
        "Scanning Side B..."
    )

    side_b_codes = decode_qr_image(
        side_b
    )

    print(
        f"Found {len(side_b_codes)} QR codes."
    )

    for qr_id, data in side_b_codes:
        print(
            f"  QR {qr_id}: "
            f"{len(data)} bytes"
        )

    # --------------------------------------------------------
    # Combine both sides.
    # --------------------------------------------------------

    all_codes = (
        side_a_codes +
        side_b_codes
    )

    chunks = {}

    for qr_id, data in all_codes:

        if qr_id in chunks:

            raise ValueError(
                f"Duplicate QR ID detected: "
                f"{qr_id}"
            )

        chunks[qr_id] = data

    # --------------------------------------------------------
    # Make sure every QR exists.
    # --------------------------------------------------------

    missing = [
        i
        for i in range(1, 9)
        if i not in chunks
    ]

    if missing:

        raise RuntimeError(
            "Missing QR codes: " +
            ", ".join(
                str(x)
                for x in missing
            )
        )

    # --------------------------------------------------------
    # Reconstruct original music.raw.
    # --------------------------------------------------------

    print()
    print(
        "Reconstructing music data..."
    )

    reconstructed = b"".join(
        chunks[i]
        for i in range(1, 9)
    )

    RAW_OUTPUT.write_bytes(
        reconstructed
    )

    print(
        f"Reconstructed: "
        f"{len(reconstructed):,} bytes"
    )

    print(
        f"Saved as: {RAW_OUTPUT}"
    )

    # --------------------------------------------------------
    # Decode it.
    # --------------------------------------------------------

    decode_music(
        RAW_OUTPUT,
        WAV_OUTPUT,
    )


# ============================================================
# LORA MODE
# ============================================================

def receive_lora_file():
    """
    The LoRa receiver ESP32 reconstructs received.raw
    and sends:

        BEGIN_RAW:<size>
        <binary data>

    over USB serial.

    This function is intentionally kept separate from the
    QR cassette reader so both input methods end up producing
    the same music.raw file.
    """

    import serial

    SERIAL_PORT = input(
        "ESP32 serial port (e.g. COM5): "
    ).strip()

    BAUD_RATE = 115200

    print()
    print(
        f"Opening {SERIAL_PORT}..."
    )

    try:

        ser = serial.Serial(
            SERIAL_PORT,
            BAUD_RATE,
            timeout=1,
        )

    except serial.SerialException as e:

        raise RuntimeError(
            f"Could not open serial port: {e}"
        )

    try:

        print()
        print(
            "Waiting for reconstructed file..."
        )

        while True:

            line = ser.readline()

            if not line:
                continue

            text = line.decode(
                "ascii",
                errors="ignore",
            ).strip()

            print(
                f"ESP32: {text}"
            )

            if not text.startswith(
                "BEGIN_RAW:"
            ):
                continue

            size_text = text[
                len("BEGIN_RAW:"):
            ]

            expected_size = int(
                size_text
            )

            print()
            print(
                f"Receiving "
                f"{expected_size:,} bytes..."
            )

            data = bytearray()

            while len(data) < expected_size:

                remaining = (
                    expected_size -
                    len(data)
                )

                chunk = ser.read(
                    min(
                        remaining,
                        4096,
                    )
                )

                if not chunk:
                    raise TimeoutError(
                        "Timed out while "
                        "receiving file."
                    )

                data.extend(
                    chunk
                )

                percent = (
                    len(data) /
                    expected_size *
                    100
                )

                print(
                    f"\r"
                    f"{len(data):,}/"
                    f"{expected_size:,} "
                    f"({percent:.1f}%)",
                    end="",
                    flush=True,
                )

            print()

            RAW_OUTPUT.write_bytes(
                data
            )

            print(
                f"Saved: {RAW_OUTPUT}"
            )

            break

    finally:

        ser.close()

    decode_music(
        RAW_OUTPUT,
        WAV_OUTPUT,
    )


# ============================================================
# MAIN MENU
# ============================================================

def main():

    while True:

        print()
        print("=" * 60)
        print(
            "       OFF-GRID MUSIC RECEIVER"
        )
        print("=" * 60)

        print(
            "1. Receive music over LoRa"
        )

        print(
            "2. Read QR music cassette"
        )

        print(
            "3. Exit"
        )

        print()

        choice = input(
            "Select: "
        ).strip()


        try:

            if choice == "1":

                receive_lora_file()

            elif choice == "2":

                read_qr_cassette()

            elif choice == "3":

                print(
                    "Goodbye."
                )

                break

            else:

                print(
                    "Invalid choice."
                )

        except Exception as e:

            print()
            print(
                "ERROR:"
            )

            print(e)


if __name__ == "__main__":
    main()