#!/usr/bin/env python3
import argparse
import os
import re
import shutil
import sys
import time

ANSI_RE = re.compile(r"\x1b\[[0-9;]*m")


def read_exact(stream, size: int) -> bytes | None:
    data = bytearray()
    while len(data) < size:
        chunk = stream.read(size - len(data))
        if not chunk:
            return None if not data else bytes(data)
        data.extend(chunk)
    return bytes(data)


def build_lut(charset: str, invert: bool) -> list[str]:
    if not charset:
        raise ValueError("Charset cannot be empty")
    lut: list[str] = []
    last = len(charset) - 1
    for value in range(256):
        base = 255 - value if invert else value
        lut.append(charset[(base * last) // 255])
    return lut


def rgb_to_ansi256(r: int, g: int, b: int) -> int:
    r6 = int(round((r / 255) * 5))
    g6 = int(round((g / 255) * 5))
    b6 = int(round((b / 255) * 5))
    return 16 + (36 * r6) + (6 * g6) + b6


def ansi_fg(r: int, g: int, b: int, color_depth: str) -> str:
    if color_depth == "256":
        return f"\x1b[38;5;{rgb_to_ansi256(r, g, b)}m"
    return f"\x1b[38;2;{r};{g};{b}m"


def ansi_bg(r: int, g: int, b: int, color_depth: str) -> str:
    if color_depth == "256":
        return f"\x1b[48;5;{rgb_to_ansi256(r, g, b)}m"
    return f"\x1b[48;2;{r};{g};{b}m"


def setup_console() -> bool:
    if os.name != "nt":
        return True

    # First try native VT mode support.
    try:
        import ctypes

        kernel32 = ctypes.windll.kernel32
        handle = kernel32.GetStdHandle(-11)  # STD_OUTPUT_HANDLE
        mode = ctypes.c_uint32()
        if kernel32.GetConsoleMode(handle, ctypes.byref(mode)):
            if kernel32.SetConsoleMode(handle, mode.value | 0x0004):
                return True
    except Exception:
        pass

    # Fallback: colorama can translate ANSI to Windows API calls.
    try:
        import colorama

        colorama.just_fix_windows_console()
        return True
    except Exception:
        return False


def strip_ansi(text: str) -> str:
    return ANSI_RE.sub("", text)


def truncate_ansi_visual(text: str, max_cols: int) -> str:
    if max_cols <= 0:
        return ""

    out: list[str] = []
    cols = 0
    i = 0
    n = len(text)
    while i < n and cols < max_cols:
        if text[i] == "\x1b":
            m = ANSI_RE.match(text, i)
            if m:
                out.append(m.group(0))
                i = m.end()
                continue
        out.append(text[i])
        cols += 1
        i += 1
    out.append("\x1b[0m")
    return "".join(out)


def build_index_map(src_len: int, dst_len: int) -> list[int]:
    if src_len <= 0 or dst_len <= 0:
        return []
    return [min(src_len - 1, (i * src_len) // dst_len) for i in range(dst_len)]


def compute_fit(src_w: int, src_h: int, max_w: int, max_h: int) -> tuple[int, int]:
    max_w = max(1, max_w)
    max_h = max(1, max_h)
    scale = min(max_w / src_w, max_h / src_h)
    out_w = max(1, int(src_w * scale))
    out_h = max(1, int(src_h * scale))
    return out_w, out_h


def parse_ansi_cells(line: str, expected_cells: int) -> list[tuple[str, str]]:
    cells: list[tuple[str, str]] = []
    active_fg = ""
    active_bg = ""
    i = 0
    n = len(line)
    while i < n:
        if line[i] == "\x1b":
            m = ANSI_RE.match(line, i)
            if m:
                code = m.group(0)
                if code == "\x1b[0m":
                    active_fg = ""
                    active_bg = ""
                elif code.startswith("\x1b[38;"):
                    active_fg = code
                elif code.startswith("\x1b[48;"):
                    active_bg = code
                i = m.end()
                continue
        cells.append((f"{active_fg}{active_bg}", line[i]))
        i += 1

    if len(cells) < expected_cells:
        cells.extend([("", " ")] * (expected_cells - len(cells)))
    elif len(cells) > expected_cells:
        cells = cells[:expected_cells]
    return cells


def sample_mono_line(line: str, col_map: list[int], src_w: int) -> str:
    raw = strip_ansi(line)
    if len(raw) < src_w:
        raw = raw + (" " * (src_w - len(raw)))
    elif len(raw) > src_w:
        raw = raw[:src_w]
    return "".join(raw[c] for c in col_map)


def sample_color_line(line: str, col_map: list[int], src_w: int) -> str:
    # Fast path: stream-parse ANSI and only sample required columns.
    # This avoids building a full per-cell table for very wide frames.
    if not col_map:
        return "\x1b[0m"

    targets = col_map
    target_i = 0

    active_fg = ""
    active_bg = ""
    last_style = None

    parts: list[str] = []
    col = 0
    i = 0
    n = len(line)
    while i < n and col < src_w and target_i < len(targets):
        ch = line[i]
        if ch == "\x1b":
            m = ANSI_RE.match(line, i)
            if m:
                code = m.group(0)
                if code == "\x1b[0m":
                    active_fg = ""
                    active_bg = ""
                elif code.startswith("\x1b[38;"):
                    active_fg = code
                elif code.startswith("\x1b[48;"):
                    active_bg = code
                i = m.end()
                continue

        if col == targets[target_i]:
            style = f"{active_fg}{active_bg}"
            if style != last_style:
                parts.append(style if style else "\x1b[0m")
                last_style = style
            parts.append(ch)
            target_i += 1
            while target_i < len(targets) and targets[target_i] == col:
                parts.append(ch)
                target_i += 1

        col += 1
        i += 1

    if target_i < len(targets):
        if last_style != "":
            parts.append("\x1b[0m")
        parts.append(" " * (len(targets) - target_i))

    parts.append("\x1b[0m")
    return "".join(parts)


def parse_header(file_obj) -> dict[str, str]:
    magic = file_obj.readline().rstrip("\n")
    if magic != "ASV1":
        raise ValueError("Unsupported file format (expected ASV1)")

    meta: dict[str, str] = {}
    while True:
        line = file_obj.readline()
        if not line:
            raise ValueError("Unexpected EOF while reading header")
        line = line.rstrip("\n")
        if line == "ENDHDR":
            break
        if "=" not in line:
            raise ValueError(f"Invalid header line: {line}")
        key, value = line.split("=", 1)
        meta[key] = value
    return meta


def parse_asv2_header(file_obj) -> dict[str, str]:
    # file_obj is opened in binary mode.
    magic = file_obj.readline().decode("utf-8", errors="strict").rstrip("\n")
    if magic != "ASV2":
        raise ValueError("Unsupported file format (expected ASV2)")

    meta: dict[str, str] = {}
    while True:
        line_b = file_obj.readline()
        if not line_b:
            raise ValueError("Unexpected EOF while reading header")
        line = line_b.decode("utf-8", errors="strict").rstrip("\n")
        if line == "ENDHDR":
            break
        if "=" not in line:
            raise ValueError(f"Invalid header line: {line}")
        key, value = line.split("=", 1)
        meta[key] = value
    return meta


def play_once_text(
    path: str,
    clear_screen: bool,
    no_color: bool,
    ansi_ok: bool,
    cols: int,
    rows: int,
    max_skip: int,
    speed: float,
) -> None:
    with open(path, "r", encoding="utf-8", newline="") as f:
        meta = parse_header(f)
        fps = float(meta.get("fps", "12"))
        width = int(meta["width"])
        height = int(meta["height"])
        has_color = meta.get("color", "0") == "1"
        render_color = has_color and not no_color and ansi_ok

        effective_fps = fps * speed
        frame_duration = 1.0 / effective_fps if effective_fps > 0 else 0.0

        if clear_screen:
            sys.stdout.write("\x1b[2J")
        sys.stdout.write("\x1b[H\x1b[?7l")
        sys.stdout.write("\x1b[?25l")
        sys.stdout.flush()

        next_frame_time = time.perf_counter()
        last_drawn_h = 0
        cached_size: tuple[int, int] | None = None
        row_map: list[int] = []
        col_map: list[int] = []
        try:
            while True:
                now = time.perf_counter()
                if frame_duration > 0 and max_skip > 0:
                    skipped = 0
                    while now - next_frame_time > frame_duration and skipped < max_skip:
                        for _ in range(height):
                            if not f.readline():
                                return
                        next_frame_time += frame_duration
                        skipped += 1
                        now = time.perf_counter()

                term = shutil.get_terminal_size(fallback=(120, 40))
                max_w = cols if cols > 0 else term.columns
                max_h = rows if rows > 0 else term.lines
                view_w, view_h = compute_fit(width, height, max_w, max_h)

                if cached_size != (view_w, view_h):
                    cached_size = (view_w, view_h)
                    row_map = build_index_map(height, view_h)
                    col_map = build_index_map(width, view_w)
                    sys.stdout.write("\x1b[2J")

                needed_rows = set(row_map)
                needed_lines: dict[int, str] = {}
                for src_row in range(height):
                    line = f.readline()
                    if not line:
                        return
                    if src_row in needed_rows:
                        needed_lines[src_row] = line.rstrip("\n")

                rendered_lines: list[str] = []
                if render_color:
                    for r in row_map:
                        rendered_lines.append(sample_color_line(needed_lines.get(r, ""), col_map, width))
                else:
                    for r in row_map:
                        rendered_lines.append(sample_mono_line(needed_lines.get(r, ""), col_map, width))

                out: list[str] = ["\x1b[H"]
                for line in rendered_lines:
                    out.append(line)
                    out.append("\x1b[K\n")
                if len(rendered_lines) < last_drawn_h:
                    out.append("\x1b[J")
                sys.stdout.write("".join(out))
                sys.stdout.flush()
                last_drawn_h = len(rendered_lines)

                next_frame_time += frame_duration
                sleep_for = next_frame_time - time.perf_counter()
                if sleep_for > 0:
                    time.sleep(sleep_for)
        finally:
            sys.stdout.write("\x1b[0m\x1b[?25h\x1b[?7h\n")
            sys.stdout.flush()


def render_halfblock_line(
    frame: bytes,
    src_w: int,
    src_scaled_h: int,
    src_cell_row: int,
    col_map: list[int],
    color_depth: str,
    no_color: bool,
) -> str:
    # src_scaled_h is 2 * cell_height for halfblock mode.
    row_stride = src_w * 3
    top_row = src_cell_row * 2
    bottom_row = top_row + 1
    if bottom_row >= src_scaled_h:
        bottom_row = src_scaled_h - 1

    top_off = top_row * row_stride
    bot_off = bottom_row * row_stride

    parts: list[str] = []
    last_code = None
    for c in col_map:
        idx = c * 3
        tr = frame[top_off + idx]
        tg = frame[top_off + idx + 1]
        tb = frame[top_off + idx + 2]
        br = frame[bot_off + idx]
        bg = frame[bot_off + idx + 1]
        bb = frame[bot_off + idx + 2]

        if no_color:
            code = None
        else:
            code = f"{ansi_fg(tr, tg, tb, color_depth)}{ansi_bg(br, bg, bb, color_depth)}"
        if code != last_code and code is not None:
            parts.append(code)
            last_code = code
        parts.append("▀")
    parts.append("\x1b[0m")
    return "".join(parts)


def render_chars_line(
    frame: bytes,
    src_w: int,
    src_row: int,
    col_map: list[int],
    lut: list[str],
    pix_fmt: str,
    color_depth: str,
    no_color: bool,
) -> str:
    parts: list[str] = []
    last_code = None
    if pix_fmt == "rgb24":
        row_stride = src_w * 3
        off = src_row * row_stride
        for c in col_map:
            idx = off + (c * 3)
            r = frame[idx]
            g = frame[idx + 1]
            b = frame[idx + 2]
            brightness = (77 * r + 150 * g + 29 * b) // 256
            ch = lut[brightness]
            if no_color:
                parts.append(ch)
            else:
                code = ansi_fg(r, g, b, color_depth)
                if code != last_code:
                    parts.append(code)
                    last_code = code
                parts.append(ch)
    else:
        # gray
        row_stride = src_w
        off = src_row * row_stride
        for c in col_map:
            parts.append(lut[frame[off + c]])
    parts.append("\x1b[0m")
    return "".join(parts)


def play_once_raw(
    path: str,
    clear_screen: bool,
    no_color: bool,
    ansi_ok: bool,
    cols: int,
    rows: int,
    max_skip: int,
    speed: float,
) -> None:
    with open(path, "rb") as f:
        meta = parse_asv2_header(f)
        fps = float(meta.get("fps", "12"))
        width = int(meta["width"])
        height = int(meta["height"])
        pixel_mode = meta.get("pixel_mode", "chars")
        pix_fmt = meta.get("pix_fmt", "gray")
        scaled_h = int(meta.get("scaled_h", str(height)))
        has_color = meta.get("color", "0") == "1"
        color_depth = meta.get("color_depth", "truecolor")
        charset = meta.get("charset", " .,:;irsXA253hMHGS#9B&@")
        invert = meta.get("invert", "0") == "1"

        render_color = has_color and not no_color and ansi_ok
        effective_fps = fps * speed
        frame_duration = 1.0 / effective_fps if effective_fps > 0 else 0.0
        frame_size = width * scaled_h * (3 if pix_fmt == "rgb24" else 1)

        lut = build_lut(charset, invert=invert)

        if clear_screen:
            sys.stdout.write("\x1b[2J")
        sys.stdout.write("\x1b[H\x1b[?7l")
        sys.stdout.write("\x1b[?25l")
        sys.stdout.flush()

        next_frame_time = time.perf_counter()
        last_drawn_h = 0
        cached_size: tuple[int, int] | None = None
        row_map: list[int] = []
        col_map: list[int] = []
        try:
            while True:
                now = time.perf_counter()
                if frame_duration > 0 and max_skip > 0:
                    skipped = 0
                    while now - next_frame_time > frame_duration and skipped < max_skip:
                        if f.seekable():
                            f.seek(frame_size, 1)
                        else:
                            if read_exact(f, frame_size) is None:
                                return
                        next_frame_time += frame_duration
                        skipped += 1
                        now = time.perf_counter()

                frame = read_exact(f, frame_size)
                if frame is None or len(frame) < frame_size:
                    return

                term = shutil.get_terminal_size(fallback=(120, 40))
                max_w = cols if cols > 0 else term.columns
                max_h = rows if rows > 0 else term.lines
                view_w, view_h = compute_fit(width, height, max_w, max_h)

                if cached_size != (view_w, view_h):
                    cached_size = (view_w, view_h)
                    row_map = build_index_map(height, view_h)
                    col_map = build_index_map(width, view_w)
                    sys.stdout.write("\x1b[2J")

                out: list[str] = ["\x1b[H"]
                if pixel_mode == "halfblock":
                    for src_cell_row in row_map:
                        out.append(
                            render_halfblock_line(
                                frame,
                                src_w=width,
                                src_scaled_h=scaled_h,
                                src_cell_row=src_cell_row,
                                col_map=col_map,
                                color_depth=color_depth,
                                no_color=not render_color,
                            )
                        )
                        out.append("\x1b[K\n")
                else:
                    # chars mode: scaled_h == height in the raw container.
                    for src_row in row_map:
                        out.append(
                            render_chars_line(
                                frame,
                                src_w=width,
                                src_row=src_row,
                                col_map=col_map,
                                lut=lut,
                                pix_fmt=pix_fmt,
                                color_depth=color_depth,
                                no_color=not render_color,
                            )
                        )
                        out.append("\x1b[K\n")

                if len(row_map) < last_drawn_h:
                    out.append("\x1b[J")
                sys.stdout.write("".join(out))
                sys.stdout.flush()
                last_drawn_h = len(row_map)

                next_frame_time += frame_duration
                sleep_for = next_frame_time - time.perf_counter()
                if sleep_for > 0:
                    time.sleep(sleep_for)
        finally:
            sys.stdout.write("\x1b[0m\x1b[?25h\x1b[?7h\n")
            sys.stdout.flush()


def play_once(
    path: str,
    clear_screen: bool,
    no_color: bool,
    ansi_ok: bool,
    cols: int,
    rows: int,
    max_skip: int,
    speed: float,
) -> None:
    # Probe magic in binary mode (works for both ASV1 and ASV2).
    with open(path, "rb") as f:
        magic = f.readline().decode("utf-8", errors="replace").rstrip("\n")

    if magic == "ASV2":
        return play_once_raw(path, clear_screen, no_color, ansi_ok, cols, rows, max_skip, speed)
    return play_once_text(path, clear_screen, no_color, ansi_ok, cols, rows, max_skip, speed)


def main() -> None:
    parser = argparse.ArgumentParser(description="Play a raw ASCII video (.asv) in a terminal.")
    parser.add_argument("input", help="Input .asv file")
    parser.add_argument("--loop", action="store_true", help="Loop playback")
    parser.add_argument("--no-clear", action="store_true", help="Do not clear screen before playback")
    parser.add_argument("--no-color", action="store_true", help="Ignore ANSI colors and render grayscale text only")
    parser.add_argument("--cols", type=int, default=0, help="Force render width in characters (0 = auto)")
    parser.add_argument("--rows", type=int, default=0, help="Force render height in rows (0 = auto)")
    parser.add_argument(
        "--max-skip",
        type=int,
        default=0,
        help="Max frames to skip when behind schedule (0 = never skip)",
    )
    parser.add_argument(
        "--speed",
        type=float,
        default=0.28,
        help="Playback speed multiplier (e.g. 0.5 = half speed, 2 = double speed)",
    )
    args = parser.parse_args()

    ansi_ok = setup_console()

    while True:
        play_once(
            args.input,
            clear_screen=not args.no_clear,
            no_color=args.no_color,
            ansi_ok=ansi_ok,
            cols=args.cols,
            rows=args.rows,
            max_skip=args.max_skip,
            speed=max(0.01, args.speed),
        )
        if not args.loop:
            break


if __name__ == "__main__":
    main()
