from __future__ import annotations
import json
import logging
import time
import urllib.request
from pathlib import Path
from typing import Iterator
from riemann.paths import DATA_DIR, HF_CACHE_DIR

logger = logging.getLogger("riemann.data")
POWERCODING_MAX_EXAMPLES = 100000
POWERCODING_SHARDS = [
    "powercoding-python/powercoding-python-raw-00001-of-00006-chunk-000-of-005.jsonl",
    "powercoding-python/powercoding-python-raw-00001-of-00006-chunk-001-of-005.jsonl",
    "powercoding-python/powercoding-python-raw-00002-of-00006-chunk-000-of-005.jsonl",
    "powercoding-python/powercoding-python-raw-00003-of-00006-chunk-000-of-008.jsonl",
]


class ResilientWriter:

    def __init__(self, path: Path, mode: str = "a", attempts: int = 30):
        self.path = path
        self.attempts = attempts
        self._fh = None
        if mode == "w":
            self._open("w")
            self._fh.close()
            self._fh = None
        self._open("a")

    def _open(self, mode: str = "a") -> None:
        for attempt in range(self.attempts):
            try:
                self._fh = self.path.open(mode, encoding="utf-8")
                return
            except OSError as e:
                if attempt == self.attempts - 1:
                    raise
                wait = min(60, 2 * (attempt + 1))
                logger.warning(
                    "open %s failed (%d/%d), retrying in %ds: %s",
                    self.path.name,
                    attempt + 1,
                    self.attempts,
                    wait,
                    e,
                )
                time.sleep(wait)

    def write(self, text: str) -> None:
        for attempt in range(self.attempts):
            try:
                self._fh.write(text)
                return
            except (OSError, ValueError, AttributeError) as e:
                if attempt == self.attempts - 1:
                    raise
                wait = min(60, 2 * (attempt + 1))
                logger.warning(
                    "write to %s failed (%d/%d), reopening in %ds: %s",
                    self.path.name,
                    attempt + 1,
                    self.attempts,
                    wait,
                    e,
                )
                time.sleep(wait)
                try:
                    if self._fh is not None:
                        self._fh.close()
                except Exception:
                    pass
                self._fh = None
                self._open("a")

    def flush(self) -> None:
        try:
            if self._fh is not None:
                self._fh.flush()
        except OSError:
            pass

    def close(self) -> None:
        try:
            if self._fh is not None:
                self._fh.close()
        except OSError:
            pass
        self._fh = None

    def __enter__(self):
        return self

    def __exit__(self, *exc):
        self.close()


def _retry_write(fh, text: str, attempts: int = 3) -> None:
    if isinstance(fh, ResilientWriter):
        fh.write(text)
        return
    for attempt in range(attempts):
        try:
            fh.write(text)
            return
        except OSError as e:
            if attempt == attempts - 1:
                raise
            logger.warning("write failed (%s), retrying in 2s: %s", attempt + 1, e)
            time.sleep(2)


def _write_pretrain_doc(fh, text: str) -> None:
    text = text.strip()
    if not text:
        return
    _retry_write(fh, text.replace("\n", " \u2029 ") + "\n")


def _write_ntk_record(fh, prompt: str, completion: str) -> None:
    prompt, completion = (prompt.strip(), completion.strip())
    if not prompt or not completion:
        return
    _retry_write(
        fh,
        json.dumps({"prompt": prompt, "completion": completion}, ensure_ascii=False)
        + "\n",
    )


def _format_fable_message(m: dict) -> str:
    role = str(m.get("role", "") or "")
    content = m.get("content")
    content = "" if content is None else str(content)
    name = m.get("name")
    tool_calls = m.get("tool_calls")
    header = role
    if role == "tool" and name:
        header = f"tool[{name}]"
    pieces = []
    if content.strip():
        pieces.append(content)
    if tool_calls:
        for call in tool_calls:
            call = dict(call) if not isinstance(call, dict) else call
            fn = call.get("function") or {}
            fn = dict(fn) if not isinstance(fn, dict) else fn
            fn_name = fn.get("name") or call.get("name") or "tool"
            fn_args = fn.get("arguments") or call.get("arguments") or ""
            if not isinstance(fn_args, str):
                fn_args = json.dumps(fn_args, ensure_ascii=False)
            pieces.append(f"<tool_call name={fn_name}>{fn_args}</tool_call>")
    if not pieces:
        return ""
    return f"{header}: " + "\n".join(pieces)


def _retry_call(label: str, fn, attempts: int = 12):
    for attempt in range(attempts):
        try:
            return fn()
        except Exception as e:
            if attempt == attempts - 1:
                raise
            wait = min(120, 5 * (attempt + 1))
            logger.warning(
                "%s failed (%d/%d), retrying in %ds: %s",
                label,
                attempt + 1,
                attempts,
                wait,
                e,
            )
            time.sleep(wait)


def load_fable(pretrain_fh, ntk_fh, split_dir: str = "10k") -> int:
    from huggingface_hub import hf_hub_download

    repo = "MoreThought/Fable-5.1-Max-Reasoning-Filtered-10000x"
    rel = f"{split_dir}/train.jsonl"
    try:
        local_path = _retry_call(
            f"fable download {rel}",
            lambda: hf_hub_download(
                repo, rel, repo_type="dataset", cache_dir=HF_CACHE_DIR
            ),
        )
    except Exception as e:
        logger.error("fable: could not fetch %s: %s", rel, e)
        return 0
    n = 0
    with open(local_path, "r", encoding="utf-8") as fh:
        for line in fh:
            line = line.strip()
            if not line:
                continue
            try:
                rec = json.loads(line)
            except json.JSONDecodeError as e:
                logger.warning("fable: skipping malformed line: %s", e)
                continue
            messages = rec.get("messages") or []
            if not messages:
                continue
            parts = []
            user_turn, assistant_turn = (None, None)
            for m in messages:
                m = dict(m) if not isinstance(m, dict) else m
                formatted = _format_fable_message(m)
                if formatted:
                    parts.append(formatted)
                role = str(m.get("role", "") or "")
                content = m.get("content")
                content = "" if content is None else str(content)
                if role == "user" and user_turn is None:
                    user_turn = content
                elif role == "assistant" and content.strip():
                    assistant_turn = content
            if parts:
                _write_pretrain_doc(pretrain_fh, "\n".join(parts))
            if user_turn and assistant_turn:
                _write_ntk_record(ntk_fh, user_turn, assistant_turn)
            n += 1
            if n % 1000 == 0:
                logger.info("fable: %d examples so far...", n)
    logger.info("fable: wrote %d examples", n)
    return n


def load_jabarti_finetune(pretrain_fh, ntk_fh) -> int:
    from datasets import load_dataset

    n = 0
    for split in ("train", "eval"):
        ds = load_dataset(
            "bakrianoo/jabarti-llm-dataset",
            "finetune",
            split=split,
            streaming=True,
            cache_dir=HF_CACHE_DIR,
        )
        for ex in ds:
            question = str(ex.get("question", ""))
            answer = str(ex.get("answer", ""))
            if not question or not answer:
                continue
            _write_ntk_record(ntk_fh, question, answer)
            _write_pretrain_doc(pretrain_fh, f"Q: {question}\nA: {answer}")
            n += 1
    logger.info("jabarti_finetune: wrote %d examples", n)
    return n


def load_jabarti_pretrain(pretrain_fh, skip_first: int = 0) -> int:
    from datasets import load_dataset

    n = skip_first
    attempts = 12
    for attempt in range(attempts):
        try:
            ds = load_dataset(
                "bakrianoo/jabarti-llm-dataset",
                "pretrain",
                split="train",
                streaming=True,
                cache_dir=HF_CACHE_DIR,
            )
            if n > 0:
                logger.info("jabarti_pretrain: skipping first %d examples", n)
                ds = ds.skip(n)
            for ex in ds:
                text = str(ex.get("text", ""))
                if not text:
                    continue
                _write_pretrain_doc(pretrain_fh, text)
                n += 1
                if n % 100000 == 0:
                    logger.info("jabarti_pretrain: %d examples so far...", n)
            break
        except Exception as e:
            if attempt == attempts - 1:
                raise
            wait = min(120, 5 * (attempt + 1))
            logger.warning(
                "jabarti_pretrain stream failed at %d (%d/%d), resuming in %ds: %s",
                n,
                attempt + 1,
                attempts,
                wait,
                e,
            )
            time.sleep(wait)
    logger.info("jabarti_pretrain: wrote %d examples", n)
    return n


def load_exeworldbench(pretrain_fh) -> int:
    from datasets import load_dataset

    ds = load_dataset("HelenPeng/ExeWorldBench", split="test", cache_dir=HF_CACHE_DIR)
    n = 0
    for ex in ds:
        text = str(ex.get("text", ""))
        if not text:
            continue
        _write_pretrain_doc(pretrain_fh, text)
        n += 1
    logger.info("exeworldbench: wrote %d examples", n)
    return n


def load_codevideobench_captions(pretrain_fh) -> int:
    from huggingface_hub import HfApi, hf_hub_download

    api = HfApi()
    files = api.list_repo_files("code-video-model/CodeVideoBench", repo_type="dataset")
    prompt_files = sorted((f for f in files if f.endswith("prompt.txt")))
    n = 0
    for rel_path in prompt_files:
        try:
            local_path = hf_hub_download(
                "code-video-model/CodeVideoBench",
                rel_path,
                repo_type="dataset",
                cache_dir=HF_CACHE_DIR,
            )
        except Exception as e:
            logger.warning("codevideobench: failed to fetch %s: %s", rel_path, e)
            continue
        text = Path(local_path).read_text(encoding="utf-8", errors="replace")
        _write_pretrain_doc(pretrain_fh, f"[video caption: {rel_path}]\n{text}")
        n += 1
    logger.info("codevideobench_captions: wrote %d examples", n)
    return n


def _iter_jsonl_byte_range(url: str, max_bytes: int) -> Iterator[dict]:
    req = urllib.request.Request(url, headers={"Range": f"bytes=0-{max_bytes}"})
    with urllib.request.urlopen(req) as r:
        raw = r.read()
    lines = raw.decode("utf-8", errors="replace").split("\n")[:-1]
    for line in lines:
        line = line.strip()
        if not line:
            continue
        try:
            yield json.loads(line)
        except json.JSONDecodeError:
            continue


def load_powercoding(pretrain_fh, cap: int = POWERCODING_MAX_EXAMPLES) -> int:
    n = 0
    per_shard_cap = max(1, cap // len(POWERCODING_SHARDS))
    for shard in POWERCODING_SHARDS:
        if n >= cap:
            break
        url = f"https://huggingface.co/datasets/Tiiny/PowerCoding/resolve/main/{shard}"
        max_bytes = per_shard_cap * 8000
        shard_n = 0
        for rec in _iter_jsonl_byte_range(url, max_bytes):
            if n >= cap or shard_n >= per_shard_cap:
                break
            text = str(rec.get("text", ""))
            if not text:
                continue
            _write_pretrain_doc(pretrain_fh, text)
            n += 1
            shard_n += 1
        logger.info(
            "powercoding: %s -> %d examples (running total %d)", shard, shard_n, n
        )
    logger.info(
        "powercoding: wrote %d examples total (capped at %d of a multi-TB dataset)",
        n,
        cap,
    )
    return n


LOADERS = [
    ("fable", load_fable),
    ("jabarti_finetune", load_jabarti_finetune),
    ("jabarti_pretrain", lambda p, n, **kw: load_jabarti_pretrain(p, **kw)),
    ("exeworldbench", lambda p, n, **kw: load_exeworldbench(p)),
    ("codevideobench_captions", lambda p, n, **kw: load_codevideobench_captions(p)),
    ("powercoding", lambda p, n, **kw: load_powercoding(p)),
]


def prepare_all(
    data_dir: Path = DATA_DIR,
    only: list[str] | None = None,
    append: bool = False,
    jabarti_skip: int = 0,
) -> dict:
    data_dir.mkdir(parents=True, exist_ok=True)
    pretrain_path = data_dir / "pretrain_corpus.txt"
    ntk_path = data_dir / "ntk_finetune.jsonl"
    loaders = [(n, fn) for n, fn in LOADERS if only is None or n in only]
    mode = "a" if append else "w"
    counts = {}
    with ResilientWriter(pretrain_path, mode) as pfh, ResilientWriter(
        ntk_path, mode
    ) as nfh:
        for name, fn in loaders:
            try:
                if name == "jabarti_pretrain":
                    counts[name] = fn(pfh, nfh, skip_first=jabarti_skip)
                else:
                    counts[name] = fn(pfh, nfh)
            except Exception as e:
                logger.exception("Failed to load dataset %s: %s", name, e)
                counts[name] = 0
            pfh.flush()
            nfh.flush()
    counts["_pretrain_corpus_path"] = str(pretrain_path)
    counts["_ntk_finetune_path"] = str(ntk_path)
    return counts


if __name__ == "__main__":
    import argparse

    logging.basicConfig(
        level=logging.INFO, format="%(asctime)s %(levelname)s %(message)s"
    )
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument(
        "--only",
        nargs="*",
        default=None,
        help="only run these loader names, e.g. --only fable",
    )
    parser.add_argument(
        "--append",
        action="store_true",
        help="append to existing corpus files instead of overwriting them",
    )
    parser.add_argument(
        "--jabarti-skip",
        type=int,
        default=0,
        help="resume jabarti_pretrain by skipping this many already-ingested rows",
    )
    args = parser.parse_args()
    result = prepare_all(
        only=args.only, append=args.append, jabarti_skip=args.jabarti_skip
    )
    print(json.dumps(result, indent=2))