# x-run: python3 % ~/downloads/moZtoMP7HAA.mp4 /tmp/video.cani
from typing import Literal
from dataclasses import dataclass
from struct import pack
from sys import argv
from subprocess import Popen, PIPE, run
from tqdm import tqdm
from tempfile import TemporaryDirectory
from glob import glob
from PIL import Image
from functools import lru_cache


PIX_BITS = [[1, 2], [4, 8], [16, 0]]


@dataclass
class VideoMetadata:
    framerate: Literal[20, 10, 5] = 10
    audio_channels: Literal[1, 2] = 2
    sample_rate: Literal[12000, 24000, 48000] = 48000
    screen_width: int = 164
    screen_height: int = 81

    @property
    def audio_samples_per_frame(self) -> int:
        return self.sample_rate // self.framerate

    def serialize(self) -> bytes:
        return bytes([ self.framerate, self.audio_channels ]) \
            + pack("<HHH", self.sample_rate, self.screen_width, self.screen_height)

@dataclass
class VideoFrame:
    audio: list[bytes]
    video: list[bytes]
    palette: list[int]

    def serialize(self) -> bytes:
        return bytes.join(b"", self.audio) \
            + bytes.join(b"", self.video) \
            + bytes.join(b"", [ bytes.fromhex("%06x" % color) for color in self.palette ])

@lru_cache
def _brightness(palette: tuple, i: int) -> float:
    r, g, b = palette[i * 3 : (i + 1) * 3]
    return (r + g + b) / 768

@lru_cache
def _distance(palette: tuple, a: int, b: int) -> float:
    r1, g1, b1 = palette[a * 3 : (a + 1) * 3]
    r2, g2, b2 = palette[b * 3 : (b + 1) * 3]
    rd, gd, bd = r1 - r2, g1 - g2, b1 - b2
    return (rd * rd + gd * gd + bd * bd) / 1966608

@lru_cache
def _get_colors(imgdata, palette: tuple, x: int, y: int) -> tuple[int, int]:
    brightest_i, brightest_l = 0, 0
    darkest_i, darkest_l = 0, 768
    for oy, line in enumerate(PIX_BITS):
        for ox in range(len(line)):
            pix = imgdata[x + ox, y + oy]
            assert pix < 16, f"{pix} is too big at {x+ox}:{y+oy}"
            brightness = _brightness(palette, pix)
            if brightness > brightest_l:
                brightest_l, brightest_i = brightness, pix
            if brightness < darkest_l:
                darkest_l, darkest_i = brightness, pix
    return darkest_i, brightest_i

@lru_cache()
def _is_darker(palette: tuple, bg: int, fg: int, c: int) -> bool:
    return _distance(palette, bg, c) < _distance(palette, fg, c)

def _get_block(imgdata, palette: tuple, x: int, y: int) -> tuple[int, int, int]:
    dark_i, bri_i = _get_colors(imgdata, palette, x, y)
    assert dark_i < 16, f"{dark_i} is too big"
    assert bri_i < 16, f"{bri_i} is too big"
    out: int = 0
    for oy, line in enumerate(PIX_BITS):
        for ox, bit in enumerate(line):
            if not _is_darker(
                palette, dark_i, bri_i, imgdata[x + ox, y + oy]
            ):
                out |= bit
    # bottom right pixel fix?
    if not _is_darker(palette, dark_i, bri_i, imgdata[x + 1, y + 2]):
        out ^= 31
        dark_i, bri_i = bri_i, dark_i
    return out, dark_i, bri_i

metadata = VideoMetadata(
    framerate=20,
    audio_channels=2,
    sample_rate=24000,
    screen_width=164,
    screen_height=81
)

input_video = argv[1]

with TemporaryDirectory() as tmpdir:
    run([
        "ffmpeg",
        "-i", input_video,
        "-f", "s8",
        "-ac", str(metadata.audio_channels),
        "-ar", str(metadata.sample_rate),
        f"{tmpdir}/audio.s8"
    ])

    run([
        "ffmpeg",
        "-i", input_video,
        "-an",
        "-r", str(metadata.framerate),
        "-vf", f"scale={metadata.screen_width * 2}:{metadata.screen_height * 3}",
        f"{tmpdir}/video%06d.jpg"
    ])

    with open(argv[2], "w") as fp_out, open(f"{tmpdir}/audio.s8", "rb") as fp_audio:
        print(metadata.serialize().hex(), file=fp_out)
        for i, frame_path in tqdm(enumerate(glob(f"{tmpdir}/video*.jpg"))):
            with Image.open(frame_path) as img_in:
                img_in = img_in.convert("P", palette=Image.Palette.ADAPTIVE, colors=16)
                img_data = img_in.load()
                img_palette = tuple(img_in.getpalette()) # type: ignore

                audio_samples = fp_audio.read(metadata.audio_samples_per_frame * metadata.audio_channels)

                frame = VideoFrame(
                    [
                        audio_samples[i::metadata.audio_channels]
                        for i in range(metadata.audio_channels)
                    ],
                    [],
                    [
                        (r << 16) | (g << 8) | b
                        for r, g, b
                        in zip(img_palette[0::3], img_palette[1::3], img_palette[2::3])  # type: ignore
                    ])

                for y in range(0, img_in.height - 2, 3):
                    line = bytearray()
                    for x in range(0, img_in.width - 1, 2):
                        ch, bg, fg = _get_block(img_data, img_palette, x, y)
                        line.extend([(ch + 0x80) & 0xFF, fg << 4 | bg])
                    frame.video.append(line)

                print(frame.serialize().hex(), file=fp_out)