"""بارگذاری مدل آموزش‌دیده و تولید پاسخ."""
from pathlib import Path
import json, threading
from text_utils import build_prompt

BASE = Path(__file__).resolve().parent
WEIGHTS = BASE / 'model' / 'weights.pt'
VOCAB = BASE / 'model' / 'vocab.json'

TEMPERATURE = 0.0   # 0 = همیشه محتمل‌ترین نویسه (پایدارترین پاسخ)؛ برای تنوع مثلاً 0.5
MAX_NEW = 200       # حداکثر نویسه‌های پاسخ
_CACHE = {}
_LOCK = threading.Lock()


def is_trained() -> bool:
    return WEIGHTS.exists() and VOCAB.exists()


def _stamp():
    """امضای فایل‌های مدل؛ با هر آموزش جدید عوض می‌شود تا مدل بدون راه‌اندازی مجدد دوباره بارگذاری شود."""
    return (WEIGHTS.stat().st_mtime_ns, VOCAB.stat().st_mtime_ns)


def _load():
    if not is_trained():
        return None
    with _LOCK:
        stamp = _stamp()
        cached = _CACHE.get('model')
        if cached and cached[0] == stamp:
            return cached[1]
        import torch
        from model_model import TinyGPT
        ckpt = torch.load(WEIGHTS, map_location='cpu', weights_only=True)
        vocab = json.loads(VOCAB.read_text(encoding='utf-8'))
        model = TinyGPT(len(vocab), ckpt['block_size'], ckpt['n_embd'], ckpt['n_head'], ckpt['n_layer'])
        model.load_state_dict(ckpt['state_dict']); model.eval()
        _CACHE['model'] = (stamp, (model, vocab, ckpt['block_size']))
        return _CACHE['model'][1]


def reply(question: str):
    """پاسخ مدل به سؤال؛ اگر مدل آموزش ندیده یا پاسخی نساخت None."""
    loaded = _load()
    if not loaded:
        return None
    import torch
    from model_model import encode, decode
    model, vocab, block = loaded
    ids = encode(build_prompt(question), vocab)[-block:]
    if not ids:
        return None
    x = torch.tensor([ids], dtype=torch.long); n0 = x.shape[1]
    nl = vocab.get('\n', -1)
    with torch.no_grad():
        for _ in range(MAX_NEW):
            logits = model(x[:, -block:])[:, -1, :]
            if TEMPERATURE <= 0:
                nxt = logits.argmax(dim=-1, keepdim=True)
            else:
                nxt = torch.multinomial(torch.softmax(logits / TEMPERATURE, dim=-1), 1)
            if nxt.item() == nl:
                break
            x = torch.cat([x, nxt], dim=1)
    return decode(x[0, n0:].tolist(), vocab).strip() or None