"""Generic text decision task: the state is a string (document, entity, log, ...)
or empty. Used by general QA, sentiment and any JSONL decision dataset.

States are stored as ragged UTF-8 bytes and tokenised at batch time with the
model's tokenizer (`task.tokenizer`: byte-level or the pretrained encoder's), so
one stored dataset works for every text backend.
"""
from __future__ import annotations

from typing import Sequence

import numpy as np

from minijev.text_task_utils import normalise_log
from minijev.tokenizers import ByteTokenizer
from minijev.tasks import Task, ragged_pack, ragged_take


class TextDecisionTask(Task):
    name = "text"
    state_encoder = "text"
    default_question = ""
    default_choices: tuple[str, ...] = ()

    def __init__(self, cfg: dict | None = None):
        super().__init__(cfg)
        self.tokenizer = ByteTokenizer()

    @property
    def max_len(self) -> int:
        return int(self.cfg.get("max_len", 256))

    def preprocess(self, text: str) -> str:
        return normalise_log(text, int(self.cfg.get("max_lines", 20))) if self.state_encoder == "log" else text

    def states_to_arrays(self, states: Sequence[str | None]) -> dict[str, np.ndarray]:
        flat, offsets = ragged_pack([np.frombuffer(self.preprocess(s or "").encode("utf-8"), np.uint8)
                                     for s in states], np.uint8)
        return {"text": flat, "text_offsets": offsets}

    def texts(self, arrays, idx) -> list[str]:
        flat, lengths = ragged_take(arrays["text"], arrays["text_offsets"], np.asarray(idx))
        out, pos = [], 0
        for n in lengths.tolist():
            out.append(bytes(flat[pos:pos + n]).decode("utf-8", "replace"))
            pos += n
        return out

    def encode_arrays(self, arrays, idx):
        texts = self.texts(arrays, idx)
        ids = self.tokenizer.encode_batch(texts, self.max_len, dynamic=True)
        return {"state_ids": ids}

    def example_states(self, n=4, seed=0):
        base = ["hotdog", "whale", "Python", "", "The support was terrible and nobody answered me.",
                "Great update, everything works now!", "Redis", "balena"]
        rng = np.random.default_rng(seed)
        return [base[i] for i in rng.integers(0, len(base), n)]


class LogDecisionTask(TextDecisionTask):
    name = "log"
    state_encoder = "log"

    def encode_arrays(self, arrays, idx):
        return {"state_log_ids": super().encode_arrays(arrays, idx)["state_ids"]}


def make_task(cfg: dict) -> TextDecisionTask:
    return LogDecisionTask(cfg) if cfg.get("name") == "log" else TextDecisionTask(cfg)