# ==============================================================================
# JiRack 1B PROD - ANTI-HANG + DEBUG VERSION
# ==============================================================================
import os
import sys
import time
import torch
import warnings
import torch.nn.functional as F
from transformers import AutoTokenizer

warnings.filterwarnings("ignore", message="Attempting to use hipBLASLt")
os.environ["TORCH_BLAS_BACKEND"] = "hipblas"

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.5
    top_p = 0.9
    repetition_penalty = 1.3
    generated = []

    for step in range(max_new_tokens):
        try:
            # Safety timeout per step (in case of hang in forward pass)
            start_step = time.time()

            logits, _ = model(ids[:, -1024:])

            if time.time() - start_step > 30:  # 30s per token = too slow
                print("\n⚠️  Token generation too slow - stopping")
                break

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

            # Repetition penalty
            for token_id in set(generated[-30:]):
                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
            probs /= probs.sum(dim=-1, keepdim=True)

            next_token = torch.multinomial(probs, num_samples=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❌ Error during generation step {step}: {e}")
            break

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


def main():
    print("🔍 Starting JiRack 1B PROD (Anti-Hang Mode)...")

    # Device & dtype
    if torch.cuda.is_available():
        device = torch.device("cuda")
        dtype = torch.float16
        print(f"✅ GPU detected: {torch.cuda.get_device_name(0)}")
    else:
        device = torch.device("cpu")
        dtype = torch.float32
        print("⚠️  Running on CPU (float32) - this may be slow")

    checkpoint = "jirack_pro_1b_prod.pt"

    try:
        print(f"📦 Loading model + weights from: {checkpoint}")
        start_load = time.time()

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

        # === LOAD WEIGHTS ===
        model.load_prod_weights(checkpoint, device)

        load_time = time.time() - start_load
        print(f"✅ Model loaded successfully in {load_time:.1f} seconds")

    except Exception as e:
        print(f"❌ CRITICAL LOAD ERROR: {e}")
        print("   → Check that the checkpoint file exists")
        print("   → Check that load_prod_weights method is correctly implemented")
        return

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

    print("\n💬 JiRack 1B PROD Ready (Anti-Hang Version)")
    print("=" * 75)

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

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

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


if __name__ == "__main__":
    main()