"""Prompt embedding cache: release text encoder (Qwen3.5) hidden states, one tensor per prompt.

    torchrun --nproc-per-node N -m rapid_moss.textenc --model DIR --prompts train.jsonl eval.jsonl --out CACHE
Each rank writes cache/shard-<rank>-<k>.safetensors; rank 0 merges cache/index.json at the end.
Key "" is the empty negative prompt.
"""
from __future__ import annotations

import argparse
import json
import os
import time
from pathlib import Path

import torch
import torch.distributed as dist
from safetensors import safe_open
from safetensors.torch import save_file


def read_prompts(paths):
    rows = []
    for p in paths:
        with open(p, encoding='utf-8') as f:
            rows += [json.loads(line) for line in f if line.strip()]
    return rows


class Encoder:
    def __init__(self, model_dir, device):
        from transformers import AutoModel, AutoTokenizer
        self.tok = AutoTokenizer.from_pretrained(Path(model_dir) / 'tokenizer', local_files_only=True)
        self.enc = AutoModel.from_pretrained(Path(model_dir) / 'text_encoder', dtype=torch.bfloat16,
                                             local_files_only=True).to(device).eval()
        self.device = device

    @torch.no_grad()
    def __call__(self, prompt, max_len=4096):
        # Same as MOSSImagePipeline.encode_prompt for one prompt.
        text = self.tok.apply_chat_template([{'role': 'user', 'content': prompt}], tokenize=False,
                                            add_generation_prompt=True, enable_thinking=True)
        inputs = self.tok(text, return_tensors='pt', truncation=True, max_length=max_len).to(self.device)
        return self.enc(**inputs, use_cache=False, return_dict=True).last_hidden_state[0]


class Cache:
    """Read side: id -> [L, 4096] bf16 (CPU, memory-mapped shards)."""

    def __init__(self, root):
        self.root = Path(root)
        self.index = json.loads((self.root / 'index.json').read_text())
        self._files = {}

    def __contains__(self, key):
        return key in self.index

    def get(self, key):
        name = self.index[key]
        if name not in self._files:
            self._files[name] = safe_open(self.root / name, 'pt')
        return self._files[name].get_tensor(key)


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument('--model', required=True)
    ap.add_argument('--prompts', nargs='+', required=True)
    ap.add_argument('--out', required=True)
    ap.add_argument('--shard-size', type=int, default=500)
    args = ap.parse_args()
    distributed = 'RANK' in os.environ
    if distributed:
        dist.init_process_group('nccl')
    rank, world = (dist.get_rank(), dist.get_world_size()) if distributed else (0, 1)
    device = torch.device('cuda', int(os.environ.get('LOCAL_RANK', 0)))
    torch.cuda.set_device(device)
    out = Path(args.out)
    out.mkdir(parents=True, exist_ok=True)
    rows = read_prompts(args.prompts)
    keys = [''] + [r['id'] for r in rows]
    texts = dict([('', '')] + [(r['id'], r['prompt']) for r in rows])
    if len(set(keys)) != len(keys):
        raise SystemExit('duplicate prompt ids')
    done = set()
    for p in out.glob(f'shard-{rank}-*.safetensors'):
        with safe_open(p, 'pt') as f:
            done.update(f.keys())
    mine = [k for i, k in enumerate(keys) if i % world == rank and k not in done]
    enc = Encoder(args.model, device)
    buf, shard, t0 = {}, len(list(out.glob(f'shard-{rank}-*.safetensors'))), time.time()

    def flush():
        nonlocal buf, shard
        if buf:
            save_file(buf, out / f'shard-{rank}-{shard:05d}.safetensors')
            shard += 1
            buf = {}

    for n, key in enumerate(mine):
        buf[key] = enc(texts[key]).to('cpu', torch.bfloat16).contiguous()
        if len(buf) >= args.shard_size:
            flush()
        if rank == 0 and n % 200 == 0:
            print(json.dumps({'encoded': n, 'of': len(mine), 's': round(time.time() - t0, 1)}), flush=True)
    flush()
    if distributed:
        dist.barrier()
    if rank == 0:
        index = {}
        for p in sorted(out.glob('shard-*.safetensors')):
            with safe_open(p, 'pt') as f:
                for k in f.keys():
                    index[k] = p.name
        missing = [k for k in keys if k not in index]
        (out / 'index.json').write_text(json.dumps(index))
        print(json.dumps({'index': len(index), 'missing': len(missing)}), flush=True)
    if distributed:
        dist.destroy_process_group()


if __name__ == '__main__':
    main()