from __future__ import annotations

import json
from pathlib import Path
from typing import Iterable, Iterator, List, Sequence


def read_text_records(paths: Sequence[str]) -> Iterator[str]:
    """Yield text records from txt/md/json/jsonl files or directories."""
    files: List[Path] = []
    for item in paths:
        p = Path(item)
        if p.is_dir():
            files.extend([x for x in sorted(p.rglob("*")) if x.is_file()])
        elif p.is_file():
            files.append(p)
        else:
            raise FileNotFoundError(item)

    for p in files:
        suffix = p.suffix.lower()
        if suffix in {".txt", ".md"}:
            text = p.read_text(encoding="utf-8", errors="ignore").strip()
            if text:
                yield text
        elif suffix == ".jsonl":
            with p.open("r", encoding="utf-8", errors="ignore") as f:
                for line in f:
                    line = line.strip()
                    if not line:
                        continue
                    obj = json.loads(line)
                    yield extract_text_from_json(obj)
        elif suffix == ".json":
            obj = json.loads(p.read_text(encoding="utf-8", errors="ignore"))
            if isinstance(obj, list):
                for row in obj:
                    yield extract_text_from_json(row)
            else:
                yield extract_text_from_json(obj)


def extract_text_from_json(obj) -> str:
    if isinstance(obj, str):
        return obj
    if isinstance(obj, dict):
        if "text" in obj:
            return str(obj["text"])
        if "content" in obj:
            return str(obj["content"])
        if "messages" in obj:
            return format_messages(obj["messages"])
        # Conservative fallback: concatenate scalar values.
        parts = []
        for value in obj.values():
            if isinstance(value, (str, int, float)):
                parts.append(str(value))
        return "\n".join(parts)
    return str(obj)


def format_messages(messages) -> str:
    parts = []
    for msg in messages:
        role = msg.get("role", "user")
        content = msg.get("content", "")
        parts.append(f"<|{role}|>\n{content}\n<|end|>")
    return "\n".join(parts)


class PackedTokenDataset:
    """Simple in-memory packed causal-LM dataset.

    For serious training, replace this with streaming shards / mmap arrays. This class is intentionally
    small and readable for the first Ares smoke tests.
    """

    def __init__(self, token_ids: Sequence[int], seq_len: int):
        if len(token_ids) < seq_len + 1:
            raise ValueError("Not enough tokens for one training example")
        self.token_ids = list(map(int, token_ids))
        self.seq_len = int(seq_len)
        self.n = (len(self.token_ids) - 1) // self.seq_len

    def __len__(self) -> int:
        return self.n

    def __getitem__(self, idx: int):
        import torch

        start = (idx % self.n) * self.seq_len
        chunk = self.token_ids[start : start + self.seq_len + 1]
        x = torch.tensor(chunk[:-1], dtype=torch.long)
        y = torch.tensor(chunk[1:], dtype=torch.long)
        return x, y


def encode_corpus(tokenizer_path: str, text_paths: Sequence[str]) -> List[int]:
    try:
        from tokenizers import Tokenizer
    except ImportError as exc:
        raise SystemExit("Install tokenizers first: pip install tokenizers") from exc

    tok = Tokenizer.from_file(tokenizer_path)
    ids: List[int] = []
    for record in read_text_records(text_paths):
        enc = tok.encode(record)
        ids.extend(enc.ids)
    return ids


def make_dataloader(tokenizer_path: str, text_paths: Sequence[str], seq_len: int, batch_size: int, shuffle: bool = True):
    import torch
    from torch.utils.data import DataLoader

    ids = encode_corpus(tokenizer_path, text_paths)
    dataset = PackedTokenDataset(ids, seq_len)
    return DataLoader(dataset, batch_size=batch_size, shuffle=shuffle, drop_last=True)