#!/usr/bin/env python3
# -*- coding: utf-8 -*-

from __future__ import annotations

import argparse
import json
import os
import random
import sys
from dataclasses import dataclass
from typing import List, Tuple, Optional

from tokenizers import Tokenizer

@dataclass
class BuildConfig:
    corpus_path: str
    tokenizer_path: str
    output_dir: str
    max_seq_len: int = 128
    mlm_prob: float = 0.15
    task: str = "nsp"  # nsp | sop | none
    seed: int = 42
    max_samples: Optional[int] = None

def read_lines(path: str) -> List[str]:
    with open(path, "r", encoding="utf-8") as f:
        return [line.rstrip("\n") for line in f if line.strip()]

def ensure_dir(path: str) -> None:
    os.makedirs(path, exist_ok=True)

def load_tokenizer(path: str) -> Tokenizer:
    if not os.path.isfile(path):
        print(f"❌ Tokenizer file not found: {path}", file=sys.stderr)
        sys.exit(1)
    return Tokenizer.from_file(path)

def get_special_ids(tok: Tokenizer) -> dict:
    def _id(t: str) -> int:
        tid = tok.token_to_id(t)
        if tid is None:
            raise RuntimeError(f"Special token {t} not in tokenizer vocab")
        return tid
    return {
        "cls": _id("[CLS]"),
        "sep": _id("[SEP]"),
        "pad": _id("[PAD]"),
        "unk": _id("[UNK]"),
        "mask": _id("[MASK]"),
    }

def encode_text(tok: Tokenizer, text: str) -> List[int]:
    return tok.encode(text).ids  # without special tokens

def truncate_pair(a: List[int], b: List[int], max_len: int) -> Tuple[List[int], List[int]]:
    # max_len here is the space for a+b tokens (excluding [CLS] and two [SEP])
    while len(a) + len(b) > max_len:
        if len(a) > len(b):
            a.pop()
        else:
            b.pop()
    return a, b

def build_pair_samples(lines: List[str], task: str, rng: random.Random) -> List[Tuple[str, str, int]]:
    pairs: List[Tuple[str, str, int]] = []
    if task == "nsp":
        # Adjacent lines -> positive; random line -> negative
        for i in range(len(lines) - 1):
            a = lines[i]
            if rng.random() < 0.5:
                # positive
                b = lines[i + 1]
                label = 1
            else:
                # negative
                j = rng.randrange(len(lines))
                # avoid trivial same index
                if j == i:
                    j = (j + 1) % len(lines)
                b = lines[j]
                label = 0
            pairs.append((a, b, label))
    elif task == "sop":
        # Adjacent lines: sometimes swap order
        for i in range(len(lines) - 1):
            a = lines[i]
            b = lines[i + 1]
            if rng.random() < 0.5:
                pairs.append((a, b, 1))  # correct order
            else:
                pairs.append((b, a, 0))  # swapped
    else:  # none -> single-sentence MLM only
        for i in range(len(lines)):
            pairs.append((lines[i], "", -1))
    return pairs

def apply_mlm(
    ids: List[int],
    special: dict,
    vocab_size: int,
    mlm_prob: float,
    rng: random.Random,
) -> Tuple[List[int], List[int]]:
    input_ids = list(ids)
    labels = [-100] * len(ids)

    candidate_positions = [i for i, tid in enumerate(ids) if tid not in (special["cls"], special["sep"], special["pad"]) ]
    num_to_mask = max(1, int(round(len(candidate_positions) * mlm_prob)))
    rng.shuffle(candidate_positions)
    mask_positions = set(candidate_positions[:num_to_mask])

    for pos in mask_positions:
        original_id = ids[pos]
        labels[pos] = original_id
        p = rng.random()
        if p < 0.8:
            input_ids[pos] = special["mask"]
        elif p < 0.9:
            # random token (not special preferred)
            rid = rng.randrange(vocab_size)
            input_ids[pos] = rid
        else:
            # keep original
            input_ids[pos] = original_id

    return input_ids, labels

def pad_to_length(ids: List[int], pad_id: int, length: int, pad_value: int = 0) -> List[int]:
    if len(ids) >= length:
        return ids[:length]
    return ids + [pad_id if pad_value == 0 else pad_value] * (length - len(ids))

def build_record(
    tok: Tokenizer,
    special: dict,
    a_text: str,
    b_text: str,
    max_seq_len: int,
    mlm_prob: float,
    task: str,
    rng: random.Random,
) -> dict:
    a_ids = encode_text(tok, a_text)
    b_ids = encode_text(tok, b_text) if b_text else []

    # reserve space for [CLS] A [SEP] [B [SEP]]
    max_pair_len = max_seq_len - 3 if b_ids else max_seq_len - 2
    a_ids, b_ids = truncate_pair(a_ids, b_ids, max_pair_len)

    # build sequence with specials
    if b_ids:
        input_ids = [special["cls"], *a_ids, special["sep"], *b_ids, special["sep"]]
        token_type_ids = [0] * (len(a_ids) + 2) + [1] * (len(b_ids) + 1)
        is_pair = True
    else:
        input_ids = [special["cls"], *a_ids, special["sep"]]
        token_type_ids = [0] * len(input_ids)
        is_pair = False

    # attention mask before padding
    attention_mask = [1] * len(input_ids)

    # Apply MLM on non-padded portion
    mlm_input_ids, mlm_labels = apply_mlm(
        input_ids, special, tok.get_vocab_size(), mlm_prob, rng
    )

    # Pad
    mlm_input_ids = pad_to_length(mlm_input_ids, special["pad"], max_seq_len)
    attention_mask = pad_to_length(attention_mask, 0, max_seq_len, pad_value=0)
    token_type_ids = pad_to_length(token_type_ids, 0, max_seq_len, pad_value=0)
    mlm_labels = pad_to_length(mlm_labels, -100, max_seq_len, pad_value=-100)

    rec = {
        "input_ids": mlm_input_ids,
        "attention_mask": attention_mask,
        "token_type_ids": token_type_ids,
        "mlm_labels": mlm_labels,
    }

    if task == "nsp" and is_pair:
        # label will be filled by caller
        pass
    elif task == "sop" and is_pair:
        pass
    return rec

def main() -> None:
    parser = argparse.ArgumentParser(description="Build MLM + NSP/SOP pretraining samples (JSONL)")
    parser.add_argument("--corpus", required=True, help="Path to corpus txt (one document per line)")
    parser.add_argument("--tokenizer", required=True, help="Path to tokenizer.json")
    parser.add_argument("--output_dir", required=True, help="Directory to write dataset files")
    parser.add_argument("--max_seq_len", type=int, default=128)
    parser.add_argument("--mlm_prob", type=float, default=0.15)
    parser.add_argument("--task", choices=["nsp", "sop", "none"], default="nsp")
    parser.add_argument("--seed", type=int, default=42)
    parser.add_argument("--max_samples", type=int, default=0, help="Max records to write (0 = all)")
    args = parser.parse_args()

    cfg = BuildConfig(
        corpus_path=args.corpus,
        tokenizer_path=args.tokenizer,
        output_dir=args.output_dir,
        max_seq_len=args.max_seq_len,
        mlm_prob=args.mlm_prob,
        task=args.task,
        seed=args.seed,
        max_samples=None if args.max_samples == 0 else args.max_samples,
    )

    if not os.path.isfile(cfg.corpus_path):
        print(f"❌ Corpus file not found: {cfg.corpus_path}", file=sys.stderr)
        sys.exit(1)

    ensure_dir(cfg.output_dir)

    tok = load_tokenizer(cfg.tokenizer_path)
    special = get_special_ids(tok)

    lines = read_lines(cfg.corpus_path)
    rng = random.Random(cfg.seed)

    pairs = build_pair_samples(lines, cfg.task, rng)

    out_path = os.path.join(cfg.output_dir, "pretraining_dataset.jsonl")
    stats_path = os.path.join(cfg.output_dir, "pretraining_dataset_stats.json")

    written = 0
    unk_id = special["unk"]
    unk_token_count = 0
    total_token_count = 0

    with open(out_path, "w", encoding="utf-8") as fout:
        for (a, b, label) in pairs:
            rec = build_record(
                tok=tok,
                special=special,
                a_text=a,
                b_text=b,
                max_seq_len=cfg.max_seq_len,
                mlm_prob=cfg.mlm_prob,
                task=cfg.task,
                rng=rng,
            )
            # Simple UNK count on input_ids
            total_token_count += sum(1 for _ in rec["input_ids"])  # all positions
            unk_token_count += sum(1 for t in rec["input_ids"] if t == unk_id)

            if cfg.task == "nsp" and b:
                rec["nsp_label"] = 1 if label == 1 else 0
            if cfg.task == "sop" and b:
                rec["sop_label"] = 1 if label == 1 else 0

            fout.write(json.dumps(rec, ensure_ascii=False) + "\n")
            written += 1
            if cfg.max_samples is not None and written >= cfg.max_samples:
                break

    coverage = 1.0 - (unk_token_count / max(1, total_token_count))
    stats = {
        "num_records": written,
        "task": cfg.task,
        "max_seq_len": cfg.max_seq_len,
        "mlm_prob": cfg.mlm_prob,
        "coverage_rate_on_inputs": coverage,
        "output_files": {
            "dataset_jsonl": os.path.abspath(out_path),
        },
    }
    with open(stats_path, "w", encoding="utf-8") as f:
        json.dump(stats, f, ensure_ascii=False, indent=2)

    print("✅ Dataset built.")
    print(f"   Records: {written}")
    print(f"   Coverage (inputs): {coverage:.4f}")
    print(f"   Saved to: {os.path.abspath(cfg.output_dir)}")

if __name__ == "__main__":
    main()