"""Поток картинок из TFRecord-шардов i1 (PNG 256×256 + 5 подписей) вместо parquet-потока FLUX.

Отличие от HFImageStream: запись TFRecord = [uint64 len][crc][payload][crc], внутри PNG-байты
и `caption` — пять длинных перефразировок одной картинки. Поэтому:
  * из пяти подписей берём случайно самую короткую или следующую по длине (бюджета на все нет,
    а короткая лучше влезает в --text-len);
  * картинки уже 256×256 — ресайз только если сторона не совпала с --size;
  * источник: локальная папка с шардами (быстро) или прямо репозиторий HF через range-запросы
    (block_size 16 МБ даёт ~90 записей/с, ~9.7 МБ/с — на уровне нашего темпа обучения).
"""
from __future__ import annotations

import io
import queue
import random
import struct
import threading
from pathlib import Path

import numpy as np
import torch
from PIL import Image

SHARD_GLOB = "*.tfrecord-*"


# ---------------------------------------------------------------- разбор формата
def _varint(buf: bytes, i: int) -> tuple[int, int]:
    shift = val = 0
    while True:
        b = buf[i]; i += 1
        val |= (b & 0x7F) << shift
        if not b & 0x80:
            return val, i
        shift += 7


def _fields(buf: bytes):
    i = 0
    while i < len(buf):
        key, i = _varint(buf, i)
        num, wire = key >> 3, key & 7
        if wire == 2:
            ln, i = _varint(buf, i)
            yield num, buf[i:i + ln]
            i += ln
        elif wire == 0:
            v, i = _varint(buf, i)
            yield num, v
        elif wire == 1:
            yield num, buf[i:i + 8]; i += 8
        elif wire == 5:
            yield num, buf[i:i + 4]; i += 4
        else:
            raise ValueError(f"wire type {wire}")


def _packed_i64(buf: bytes) -> list[int]:
    out, i = [], 0
    while i < len(buf):
        v, i = _varint(buf, i)
        out.append(v)
    return out


def parse_example(payload: bytes) -> dict:
    """tf.train.Example → {поле: список bytes | список int}"""
    out = {}
    for num, features in _fields(payload):
        if num != 1:
            continue
        for fn, entry in _fields(features):
            if fn != 1:
                continue
            key, val = None, {}
            for en, ev in _fields(entry):
                if en == 1:
                    key = ev.decode()
                elif en == 2:
                    for fnum, fval in _fields(ev):
                        if fnum == 1:
                            val["bytes"] = [v for n, v in _fields(fval) if n == 1]
                        elif fnum == 3:
                            val["int64"] = [x for n, v in _fields(fval) if n == 1 for x in _packed_i64(v)]
            out[key] = val.get("bytes", val.get("int64"))
    return out


def count_records(paths) -> int:
    """Сколько записей в шардах — читаем только заголовки (seek через payload), быстро."""
    total = 0
    for path in paths:
        with open(path, "rb") as fh:
            size = fh.seek(0, 2)
            pos = 0
            while pos + 12 <= size:
                fh.seek(pos)
                head = fh.read(8)
                if len(head) < 8:
                    break
                (n,) = struct.unpack("<Q", head)
                pos += 8 + 4 + n + 4
                total += 1
    return total


def count_shard_records(source: str) -> int | None:
    """Точное число записей для локальной папки; None, если источник удалённый."""
    if not Path(source).is_dir():
        return None
    paths = sorted(p.as_posix() for p in Path(source).glob(SHARD_GLOB))
    return count_records(paths) if paths else None


def read_records(fh):
    """файловый объект (локальный или HfFileSystem) → поток payload'ов"""
    while True:
        head = fh.read(8)
        if len(head) < 8:
            return
        (n,) = struct.unpack("<Q", head)
        fh.read(4)
        payload = fh.read(n)
        fh.read(4)
        yield payload


# ---------------------------------------------------------------- поток
class TFRecordStream:
    def __init__(self, source: str, size: int = 256, workers: int = 8, prefetch: int = 256,
                 block_mb: int = 16, seed: int = 0, caption_short2: bool = True):
        """source — локальная папка с шардами или id датасета на HF (repo_type=dataset)."""
        self.q: queue.Queue = queue.Queue(maxsize=prefetch)
        self.stop = False
        self.size, self.workers, self.seed = size, workers, seed
        self.caption_short2 = caption_short2
        self.local = Path(source).is_dir()
        self.source = str(source)
        self.block = block_mb * 2 ** 20
        threading.Thread(target=self._loop, daemon=True).start()

    # ---- список шардов
    def _shards(self) -> list[str]:
        if self.local:
            files = sorted(p.as_posix() for p in Path(self.source).glob(SHARD_GLOB))
        else:
            from huggingface_hub import HfFileSystem
            fs = HfFileSystem()
            files = [f"datasets/{self.source}/{p}" for p in sorted(fs.ls(f"datasets/{self.source}"))
                     if ".tfrecord-" in p]
        if not files:
            raise FileNotFoundError(f"нет шардов в {self.source}")
        random.Random(self.seed).shuffle(files)
        return files

    def _open(self, path: str):
        if self.local:
            return open(path, "rb")
        from huggingface_hub import HfFileSystem
        return HfFileSystem().open(path, "rb", block_size=self.block)

    # ---- подготовка одной записи
    def _prepare(self, payload: bytes):
        d = parse_example(payload)
        im = Image.open(io.BytesIO(d["image"][0])).convert("RGB")
        w, h = im.size
        if min(w, h) != self.size:
            sc = self.size / min(w, h)
            im = im.resize((max(self.size, round(w * sc)), max(self.size, round(h * sc))), Image.LANCZOS)
            w, h = im.size
            im = im.crop(((w - self.size) // 2, (h - self.size) // 2,
                          (w - self.size) // 2 + self.size, (h - self.size) // 2 + self.size))
        caps = [c.decode() if isinstance(c, bytes) else str(c) for c in d.get("caption", [])]
        caps = sorted(caps, key=len)
        if not caps:
            caps = [""]
        cap = random.choice(caps[:2]) if self.caption_short2 else random.choice(caps)
        arr = torch.from_numpy(np.asarray(im, dtype=np.float32)).permute(2, 0, 1) / 127.5 - 1.0
        return arr, cap

    def _loop(self):
        import collections
        from concurrent.futures import ThreadPoolExecutor
        with ThreadPoolExecutor(max_workers=self.workers) as pool:
            while not self.stop:
                try:
                    for shard in self._shards():
                        fh = self._open(shard)
                        try:
                            it = read_records(fh)
                            flying = collections.deque()
                            while not self.stop:
                                while len(flying) < self.workers:
                                    flying.append(pool.submit(self._prepare, next(it)))
                                self.q.put(flying.popleft().result())
                        finally:
                            fh.close()
                except Exception as e:
                    print(f"[tfrecord] сбой ({type(e).__name__}: {e}), переподключение через 5 с", flush=True)
                    import time
                    time.sleep(5)

    def batch(self, n: int):
        imgs, caps = [], []
        while len(imgs) < n:
            arr, cap = self.q.get()
            imgs.append(arr); caps.append(cap)
        return torch.stack(imgs), caps