"""
Nova 3.5 Nemotron — NVIDIA 70B Foundation Engine (nvidia/Llama-3.1-Nemotron-70B-Instruct-HF)
Runs NVIDIA Nemotron-70B-Instruct on AMD ROCm GPU (96 GB VRAM)
with Hierarchical Chain-of-Thought (CoT) and Autonomous Python Tool Execution.
"""

import os
import sys
import time
import torch
from transformers import AutoTokenizer, AutoModelForCausalLM, BitsAndBytesConfig

# Add project root to sys.path
sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), "..")))
from src.inference.tool_engine import parse_and_execute_tools


def run_nvidia_nemotron_engine(prompt: str, model_name: str = "nvidia/Llama-3.1-Nemotron-70B-Instruct-HF"):
    token = os.environ.get("HF_TOKEN")
    print(f"Loading NVIDIA Nemotron Model '{model_name}' onto AMD ROCm GPU (96 GB VRAM)...", flush=True)

    t0 = time.time()
    tokenizer = AutoTokenizer.from_pretrained(model_name, token=token)
    
    # 4-bit Quantization Config to comfortably fit 70B parameters inside 96 GB VRAM
    bnb_config = BitsAndBytesConfig(
        load_in_4bit=True,
        bnb_4bit_compute_dtype=torch.bfloat16,
        bnb_4bit_quant_type="nf4",
        bnb_4bit_use_double_quant=True
    )

    model = AutoModelForCausalLM.from_pretrained(
        model_name,
        quantization_config=bnb_config,
        device_map="auto",
        token=token
    )
    print(f"🚀 Loaded NVIDIA Nemotron 70B Model onto GPU in {time.time()-t0:.2f}s!", flush=True)

    messages = [
        {"role": "system", "content": "You are Nova 3.5 Nemotron, an advanced AI reasoning assistant powered by NVIDIA Nemotron-70B. Provide clear, detailed, accurate step-by-step answers. If python code execution is needed, output ```python code here ```."},
        {"role": "user", "content": prompt}
    ]

    text_input = tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
    model_inputs = tokenizer([text_input], return_tensors="pt").to("cuda")

    start_gen = time.time()
    generated_ids = model.generate(
        **model_inputs,
        max_new_tokens=350,
        temperature=0.6,
        top_p=0.9,
        do_sample=True
    )

    generated_ids = [
        output_ids[len(input_ids):] for input_ids, output_ids in zip(model_inputs.input_ids, generated_ids)
    ]

    response = tokenizer.batch_decode(generated_ids, skip_special_tokens=True)[0]
    gen_time = time.time() - start_gen

    print(f"\n=======================================================")
    print(f"🌌 NOVA 3.5 NEMOTRON GENERATION OUTPUT ({gen_time:.2f}s)")
    print(f"=======================================================")
    print(response)

    # Execute Autonomous Tool Engine if code block present
    clean_text, tool_results = parse_and_execute_tools(response)
    if tool_results:
        print("\n⚡ [Autonomous Tool Output]:")
        for res in tool_results:
            print(res)

    print("=======================================================\n")


if __name__ == "__main__":
    test_prompt = "Hello how are you? Can you tell me how many r's are in the word strawberry and write a python script to calculate 12**2 + 15**2?"
    run_nvidia_nemotron_engine(test_prompt)