# ==============================================================================
# JiRack 1B - Safe Safetensors Version (model_prod.safetensors)
# ==============================================================================
import os
import sys
import time
import torch
import warnings
import torch.nn.functional as F
from transformers import AutoTokenizer

# ====================== ANTI-CRASH SETTINGS ======================
os.environ["TORCH_BLAS_BACKEND"] = "hipblas"
os.environ["PYTORCH_JIT"] = "0"
os.environ["OMP_NUM_THREADS"] = "4"
os.environ["MKL_DISABLE_FAST_MM"] = "1"
os.environ["ATEN_CPU_CAPABILITY"] = "avx"
os.environ["TORCH_CPU_CAPABILITY"] = "avx"

warnings.filterwarnings("ignore")

print("🛡️  Safe mode activated (AVX limited)")

# ====================== Device ======================
if torch.cuda.is_available():
    device = torch.device("cuda")
    dtype = torch.float16
    print(f"✅ Using GPU: {torch.cuda.get_device_name(0)}")
else:
    device = torch.device("cpu")
    dtype = torch.float32
    print("⚠️  Running on CPU (float32)")

# ====================== Import Model ======================
try:
    from JiRackTernaryPyTorch_1b_inf import TernaryTransformer1B, TernaryConfig
except ImportError as e:
    print(f"❌ Import error: {e}")
    sys.exit(1)


@torch.no_grad()
def chat_generate(model, tokenizer, prompt, device, max_new_tokens=256):
    formatted = f"<|begin_of_text|><|start_header_id|>user<|end_header_id|>\n\n{prompt}<|eot_id|><|start_header_id|>assistant<|end_header_id|>\n\n"
    ids = tokenizer.encode(formatted, return_tensors="pt").to(device)

    print("\nJiRack: ", end="", flush=True)

    temperature = 0.7
    top_p = 0.9
    repetition_penalty = 1.25
    generated = []

    for _ in range(max_new_tokens):
        try:
            logits, _ = model(ids[:, -1024:])

            next_token_logits = logits[:, -1, :].clone()

            # Repetition penalty
            for token_id in set(generated[-20:]):
                if token_id < next_token_logits.shape[-1]:
                    next_token_logits[0, token_id] /= repetition_penalty

            next_token_logits = next_token_logits / temperature
            probs = F.softmax(next_token_logits, dim=-1)

            # Top-P
            sorted_probs, sorted_indices = torch.sort(probs, descending=True)
            cumulative_probs = torch.cumsum(sorted_probs, dim=-1)
            sorted_indices_to_remove = cumulative_probs > top_p
            sorted_indices_to_remove[..., 1:] = sorted_indices_to_remove[..., :-1].clone()
            sorted_indices_to_remove[..., 0] = False

            probs[0, sorted_indices[sorted_indices_to_remove]] = 0.0
            if probs.sum() > 0:
                probs /= probs.sum(dim=-1, keepdim=True)

            next_token = torch.multinomial(probs, 1)
            next_id = next_token.item()

            ids = torch.cat([ids, next_token], dim=-1)
            generated.append(next_id)

            token_str = tokenizer.decode([next_id], skip_special_tokens=True)
            print(token_str, end="", flush=True)

            if next_id in [tokenizer.eos_token_id, 128001, 128009]:
                break

        except Exception as e:
            print(f"\n❌ Generation error: {e}")
            break

    print("\n" + "-" * 60)


def main():
    checkpoint = "model_prod.safetensors"

    if not os.path.exists(checkpoint):
        print(f"❌ File not found: {checkpoint}")
        print("Files in current folder:")
        for f in sorted(os.listdir(".")):
            if any(f.endswith(ext) for ext in [".safetensors", ".pt", ".bin"]):
                print(f"   • {f}")
        return

    print(f"📦 Found: {checkpoint}")

    try:
        tokenizer = AutoTokenizer.from_pretrained("meta-llama/Llama-3.2-1B-Instruct")

        config = TernaryConfig()
        model = TernaryTransformer1B(config).to(device=device, dtype=dtype)

        print(f"🚀 Loading weights: {checkpoint} ...")
        model.load_prod_weights(checkpoint, device)

        print("✅ JiRack 1B PROD Ready (model_prod.safetensors)")
        print("=" * 75)

        while True:
            try:
                u = input("\nUser: ").strip()
                if u.lower() in ["exit", "quit", "выход"]:
                    print("👋 Bye!")
                    break
                if not u:
                    continue

                start = time.time()
                chat_generate(model, tokenizer, u, device)
                print(f"⏱️  {time.time() - start:.2f}s\n")

            except KeyboardInterrupt:
                print("\n🛑 Stopped by user.")
                break
            except Exception as e:
                print(f"Error: {e}")

    except Exception as e:
        print(f"❌ Failed to start model: {e}")
        print("Make sure load_prod_weights() supports .safetensors files.")


if __name__ == "__main__":
    main()