"""streaming_dataset.py — Carregador de dataset streaming para o CNN-BiGRU.

Adaptado de xavante_work/flexnet/streaming_datasets_v13_9.py e do repositório
'PowerMachine/CNN-BiGRU' no HuggingFace, com:
    - Modo streaming (IterableDataset, sem materialização completa)
    - Suporte ao repositório próprio 'PowerMachine/CNN-BiGRU' (streaming_datasets.py)
    - Fallback sintético quando datasets externos estão indisponíveis
    - Suporte multimodal: texto + imagem (placeholder) + áudio (placeholder)
    - Garantia de produzir N amostras para o teste (até 500 samples em batches de 100)

v3.0:
    - Prioriza o repositório 'PowerMachine/CNN-BiGRU' (conforme requisição do usuário)
    - Mantém fallback para os datasets V13.9.1 originais
    - Suporte a batches de 100 amostras (até 500 no total) para testes
"""
from __future__ import annotations

import logging
import os
import random
from dataclasses import dataclass, field
from typing import Any, Dict, Iterator, List, Optional, Sequence

import numpy as np
import torch

logger = logging.getLogger(__name__)


@dataclass
class MultimodalSample:
    """Amostra multimodal: texto + imagem (HxWxC float) + áudio (spec TxF)."""
    sample_id: int
    text_a: str            # Stream A (e.g. pergunta/título)
    text_b: str            # Stream B (e.g. contexto/corpo)
    image: Optional[np.ndarray] = None     # [H, W, C] float32 in [0,1]
    audio: Optional[np.ndarray] = None     # [T, F] spectrogram float32
    label: int = 0
    metadata: Dict[str, Any] = field(default_factory=dict)


# Dataset padrão V13.9.1 (mesmos do reference streaming_datasets_v13_9.py)
# NOVO v3.0: prioriza o repositório próprio 'PowerMachine/CNN-BiGRU'
# (conforme requisição do usuário: "usar do repositório 'PowerMachine/CNN-BiGRU' streaming_datasets.py")
DEFAULT_DATASETS = [
    # Prioridade 1: repositório próprio do projeto
    "PowerMachine/CNN-BiGRU",
    # Prioridade 2: datasets V13.9.1 originais
    "CEIA-POSITIVO/ultrachat_br_clustred_balanced_v1",
    "Madras1/corpus-ptbr-v2",
    "rhaymison/multmodal_175k_portuguese",
    "TucanoBR/GigaVerbo",
    "nvidia/OpenMathReasoning",
    "MathLLMs/MathVision",
    "nvidia/OpenMathInstruct-2",
    "dominguesm/restore-punctuation-ptbr-dataset",
    "carolina-c4ai/corpus-carolina",
]


def _extract_field(sample: Dict[str, Any], candidates: List[str]) -> Optional[str]:
    for f in candidates:
        if f in sample:
            v = sample[f]
            if isinstance(v, str) and v.strip():
                return v
            if isinstance(v, list):
                parts = []
                for m in v:
                    if isinstance(m, dict):
                        c = m.get("content", "")
                        if isinstance(c, str) and c.strip():
                            parts.append(c)
                    elif isinstance(m, str) and m.strip():
                        parts.append(m)
                if parts:
                    return "\n".join(parts)
    return None


def _try_load_hf_streaming(
    dataset_name: str,
    split: str = "train",
    hf_token: Optional[str] = None,
):
    """Tenta carregar dataset HF em streaming. Retorna None se falhar."""
    try:
        from datasets import load_dataset
        ds = load_dataset(dataset_name, split=split, streaming=True, token=hf_token)
        logger.info("Streaming OK: %s[%s]", dataset_name, split)
        return ds
    except Exception as e:
        logger.info("Streaming falhou para %s: %s", dataset_name, str(e)[:120])
        return None


def synthetic_multimodal_stream(
    n_samples: int,
    seed: int = 42,
    image_size: tuple = (28, 28, 1),
    audio_shape: tuple = (64, 40),
    vocab_texts: Optional[Sequence[str]] = None,
) -> Iterator[MultimodalSample]:
    """Gera amostras sintéticas multimodais determinísticas.

    Usado quando o dataset HF está indisponível (offline, rate-limit, etc.).
    Garante que o teste de 50 amostras sempre execute.
    """
    rng = random.Random(seed)
    np_rng = np.random.RandomState(seed)

    if vocab_texts is None:
        vocab_texts = [
            "o modelo aprende padrões locais com convoluções",
            "a ponte de cooperação troca informações entre fluxos",
            "grus capturam dependências temporais bidirecionais",
            "atenção cruzada reduz perplexidade em textos longos",
            "a fusão multimodal combina texto imagem e áudio",
            "penalidades evitam repetições viciosas na geração",
            "a camada anti-alucinação usa lógica fuzzy de lukasiewicz",
            "o verificador classifica passos com sigmoid binária",
            "ajustes dinâmicos de lr controlam explosão de gradiente",
            "normalização espectral estabiliza o treinamento",
        ]

    n_vocab = len(vocab_texts)
    for i in range(n_samples):
        text_a = vocab_texts[rng.randrange(n_vocab)]
        text_b = vocab_texts[rng.randrange(n_vocab)]
        # Imagem sintética: padrões estruturados simples
        img = np_rng.rand(*image_size).astype(np.float32)
        # Adiciona um padrão que depende do índice (para discriminação)
        img[i % image_size[0], :, :] = 1.0
        # Áudio sintético: espectrograma
        aud = np_rng.rand(*audio_shape).astype(np.float32) * 0.5
        # Espectro com picos determinísticos
        aud[i % audio_shape[0], :] += 0.5
        yield MultimodalSample(
            sample_id=i,
            text_a=text_a,
            text_b=text_b,
            image=img,
            audio=aud,
            label=i % 3,
            metadata={"source": "synthetic", "idx": i},
        )


class MultimodalStreamingDataset(torch.utils.data.IterableDataset):
    """Dataset streaming multimodal. Suporta HF + fallback sintético.

    Args:
        n_samples: número total de amostras a produzir.
        hf_datasets: lista de datasets HF para tentar (em ordem).
        hf_token: token HF (será limpo após uso).
        use_synthetic_fallback: se True, usa sintético quando HF falha.
        seed: seed para reprodutibilidade.
    """

    def __init__(
        self,
        n_samples: int = 50,
        hf_datasets: Optional[List[str]] = None,
        hf_token: Optional[str] = None,
        use_synthetic_fallback: bool = True,
        seed: int = 42,
        image_size: tuple = (28, 28, 1),
        audio_shape: tuple = (64, 40),
    ):
        super().__init__()
        self.n_samples = n_samples
        self.hf_datasets = hf_datasets or DEFAULT_DATASETS
        # NOTA: o hf_token NÃO é persistido como atributo de instância
        # para evitar que seja exposto em dumps/logs. Em vez disso, é
        # passado como parâmetro local durante a iteração.
        self._hf_token = hf_token  # private, limpo após iter
        self.use_synthetic_fallback = use_synthetic_fallback
        self.seed = seed
        self.image_size = image_size
        self.audio_shape = audio_shape

    @property
    def hf_token(self) -> Optional[str]:
        """Retorna o token HF atual (ou None se já limpo)."""
        return getattr(self, "_hf_token", None)

    def clear_hf_token(self) -> None:
        """Limpa o token HF da memória da instância (boa prática de segurança)."""
        self._hf_token = None

    def _try_hf(self) -> Iterator[Dict[str, Any]]:
        """Tenta carregar amostras HF. Retorna iterator vazio se falhar."""
        for ds_name in self.hf_datasets:
            ds = _try_load_hf_streaming(ds_name, split="train", hf_token=self._hf_token)
            if ds is None:
                continue
            count = 0
            text_candidates = ["text", "content", "question", "problem", "input",
                               "conversa", "description", "prompt", "instruction"]
            label_candidates = ["answer", "response", "output", "solution",
                                 "punctuated", "restored"]
            for raw in ds:
                if count >= self.n_samples:
                    break
                try:
                    text_a = _extract_field(raw, text_candidates) or ""
                    text_b = _extract_field(raw, label_candidates) or ""
                    if len(text_a) < 5:
                        text_a = "pergunta de exemplo sobre o tema"
                    if len(text_b) < 5:
                        text_b = "resposta de exemplo para contexto"
                    # Image/audio placeholder: geramos sintéticos para manter multimodal
                    img = np.random.rand(*self.image_size).astype(np.float32) * 0.5
                    aud = np.random.rand(*self.audio_shape).astype(np.float32) * 0.5
                    yield {
                        "sample_id": count,
                        "text_a": text_a[:200],  # truncado para teste
                        "text_b": text_b[:200],
                        "image": img,
                        "audio": aud,
                        "label": count % 3,
                        "metadata": {"source": ds_name, "idx": count},
                    }
                    count += 1
                except Exception as e:
                    logger.debug("skip sample %d from %s: %s", count, ds_name, e)
                    continue
            if count > 0:
                logger.info("HF yield %d samples de %s", count, ds_name)
                return

    def __iter__(self) -> Iterator[Dict[str, Any]]:
        """Itera sobre as amostras: tenta HF primeiro, depois fallback sintético.

        CORREÇÃO v2.0: O bug original retornava imediatamente após `yield from
        self._try_hf()` mesmo se HF não tivesse produzido nenhuma amostra,
        tornando o fallback sintético INACESSÍVEL quando HF falhava silenciosamente.
        Agora rastreamos o número de amostras produzidas e fazemos fallback se zero.
        """
        n_yielded = 0

        # Tentar HF primeiro
        if self.hf_datasets:
            try:
                for sample in self._try_hf():
                    yield sample
                    n_yielded += 1
                    if n_yielded >= self.n_samples:
                        return
            except Exception as e:
                logger.warning(f"HF streaming falhou ({e}) — usando fallback sintético")

        # Se HF não produziu amostras suficientes, usar fallback sintético
        if n_yielded < self.n_samples and self.use_synthetic_fallback:
            if n_yielded == 0:
                logger.info("Nenhuma amostra HF produzida — usando 100%% sintético")
            else:
                logger.info(f"HF produziu apenas {n_yielded}/{self.n_samples} — completando com sintético")
            remaining = self.n_samples - n_yielded
            for s in synthetic_multimodal_stream(
                n_samples=remaining,
                seed=self.seed + n_yielded,  # seed diferente para variar
                image_size=self.image_size,
                audio_shape=self.audio_shape,
            ):
                yield {
                    "sample_id": s.sample_id + n_yielded,  # offset para não colidir
                    "text_a": s.text_a,
                    "text_b": s.text_b,
                    "image": s.image,
                    "audio": s.audio,
                    "label": s.label,
                    "metadata": {**s.metadata, "source": "synthetic_fallback"},
                }
                n_yielded += 1
                if n_yielded >= self.n_samples:
                    # Limpa o token HF após uso completo (boa prática de segurança)
                    self.clear_hf_token()
                    return
        # Limpa o token HF após uso completo (boa prática de segurança)
        self.clear_hf_token()


def collate_multimodal(
    batch: List[Dict[str, Any]],
    tokenizer,
    max_len: int = 64,
) -> Dict[str, torch.Tensor]:
    """Cola um batch de amostras multimodais em tensores.

    Returns dict com:
        input_ids_a: [B, T]   (stream A)
        input_ids_b: [B, T]   (stream B)
        attn_mask_a: [B, T]
        attn_mask_b: [B, T]
        images:      [B, C, H, W]
        audios:      [B, 1, T, F]
        labels:      [B]
    """
    texts_a = [b["text_a"] for b in batch]
    texts_b = [b["text_b"] for b in batch]
    ids_a = tokenizer.encode_batch(texts_a, add_special=True)
    ids_b = tokenizer.encode_batch(texts_b, add_special=True)

    pad_id = tokenizer.pad_id

    def _pad(seqs, max_len):
        out = []
        masks = []
        for s in seqs:
            s = s[:max_len]
            n = len(s)
            padded = s + [pad_id] * (max_len - n)
            mask = [1] * n + [0] * (max_len - n)
            out.append(padded)
            masks.append(mask)
        return out, masks

    padded_a, masks_a = _pad(ids_a, max_len)
    padded_b, masks_b = _pad(ids_b, max_len)

    images = np.stack([b["image"] for b in batch])  # [B, H, W, C]
    # converte para [B, C, H, W]
    if images.ndim == 4:
        images = np.transpose(images, (0, 3, 1, 2))
    else:
        images = images[:, None, :, :]  # add channel dim

    audios = np.stack([b["audio"] for b in batch])  # [B, T, F]
    audios = audios[:, None, :, :]  # [B, 1, T, F]

    labels = np.array([b["label"] for b in batch], dtype=np.int64)

    return {
        "input_ids_a": torch.tensor(padded_a, dtype=torch.long),
        "input_ids_b": torch.tensor(padded_b, dtype=torch.long),
        "attn_mask_a": torch.tensor(masks_a, dtype=torch.float),
        "attn_mask_b": torch.tensor(masks_b, dtype=torch.float),
        "images": torch.tensor(images, dtype=torch.float32),
        "audios": torch.tensor(audios, dtype=torch.float32),
        "labels": torch.tensor(labels, dtype=torch.long),
        "sample_ids": [b["sample_id"] for b in batch],
    }


__all__ = [
    "MultimodalSample",
    "MultimodalStreamingDataset",
    "collate_multimodal",
    "DEFAULT_DATASETS",
    "synthetic_multimodal_stream",
]