"""
Nova All-In-One Ultimate — NVIDIA Nemotron-8B Foundation Engine
(nvidia/Mistral-NeMo-Minitron-8B-Instruct)
Runs NVIDIA 8B Nemotron on AMD ROCm GPU (96 GB VRAM) with
Autonomous Python Tool Execution and Hierarchical Reasoning.
"""

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

# 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_nova_nemotron_8b_engine(prompt: str, model_name: str = "nvidia/Mistral-NeMo-Minitron-8B-Instruct"):
    token = os.environ.get("HF_TOKEN")
    print(f"Loading NVIDIA Nemotron-8B Model '{model_name}' onto AMD ROCm GPU (96 GB VRAM)...", flush=True)

    t0 = time.time()
    tokenizer = AutoTokenizer.from_pretrained(model_name, token=token)
    model = AutoModelForCausalLM.from_pretrained(
        model_name,
        dtype=torch.bfloat16,
        device_map="cuda",
        token=token
    )
    print(f"🚀 Loaded NVIDIA Nemotron 8B Model onto GPU in {time.time()-t0:.2f}s!", flush=True)

    messages = [
        {"role": "system", "content": "You are Nova All-In-One Ultimate, an advanced AI reasoning assistant powered by NVIDIA Nemotron-8B. Provide clear, detailed, step-by-step reasoning. 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=300,
        temperature=0.7,
        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 ALL-IN-ONE (NVIDIA NEMOTRON 8B) 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 Execution Result]:", flush=True)
        for res in tool_results:
            print(res, flush=True)

    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 code snippet to calculate 12**2 + 15**2?"
    run_nova_nemotron_8b_engine(test_prompt)