import json
from datetime import datetime

import os
from huggingface_hub import HfApi

import gradio as gr
import spaces
import torch
from transformers import AutoModelForCausalLM, AutoTokenizer

MODEL_ID = "Neura-Tech-AI/Neuron-V1-14B-Instruct"
LOG_FILE = "chat_logs.jsonl"

tokenizer = None
model = None


def save_chat(user_msg, assistant_msg):
    print("Saving chat...")

    log = {
        "timestamp": datetime.utcnow().isoformat(),
        "user": user_msg,
        "assistant": assistant_msg,
    }

    with open(LOG_FILE, "a", encoding="utf-8") as f:
        f.write(json.dumps(log, ensure_ascii=False) + "\n")

    print("Chat saved!")

@spaces.GPU
def load_model():
    global tokenizer, model

    if model is not None:
        return

    print("Loading tokenizer...")

    tokenizer = AutoTokenizer.from_pretrained(
        MODEL_ID,
        trust_remote_code=True,
    )

    print("Loading model...")

    model = AutoModelForCausalLM.from_pretrained(
        MODEL_ID,
        dtype=torch.float16,
        device_map="auto",
        low_cpu_mem_usage=True,
        trust_remote_code=True,
    )

    model.eval()

    print("Model loaded successfully!")


@spaces.GPU
def chat(
    message,
    history,
    system_prompt,
    max_new_tokens,
    temperature,
    top_p,
):
    load_model()

    messages = []

    if system_prompt.strip():
        messages.append(
            {
                "role": "system",
                "content": system_prompt,
            }
        )

    if history:
        for user_msg, assistant_msg in history:
            messages.append(
                {
                    "role": "user",
                    "content": user_msg,
                }
            )

            messages.append(
                {
                    "role": "assistant",
                    "content": assistant_msg,
                }
            )

    messages.append(
        {
            "role": "user",
            "content": message,
        }
    )

    prompt = tokenizer.apply_chat_template(
        messages,
        tokenize=False,
        add_generation_prompt=True,
    )

    inputs = tokenizer(
        prompt,
        return_tensors="pt",
    )

    inputs = {k: v.to(model.device) for k, v in inputs.items()}

    with torch.inference_mode():
        outputs = model.generate(
            **inputs,
            max_new_tokens=max_new_tokens,
            temperature=temperature,
            top_p=top_p,
            do_sample=True,
            pad_token_id=tokenizer.eos_token_id,
        )

    generated_tokens = outputs[0][inputs["input_ids"].shape[1]:]

    response = tokenizer.decode(
        generated_tokens,
        skip_special_tokens=True,
    ).strip()

    # Save chat log
    save_chat(message, response)

    return response


demo = gr.ChatInterface(
    fn=chat,
    title="Neuron V1 14B Instruct",
    description="Developed by Neura Tech AI",
    additional_inputs=[
        gr.Textbox(
            value="You are Neuron, developed by Neura Tech AI.",
            label="System Prompt",
        ),
        gr.Slider(
            minimum=1,
            maximum=2048,
            value=512,
            step=1,
            label="Max New Tokens",
        ),
        gr.Slider(
            minimum=0.0,
            maximum=2.0,
            value=0.7,
            step=0.1,
            label="Temperature",
        ),
        gr.Slider(
            minimum=0.1,
            maximum=1.0,
            value=0.95,
            step=0.05,
            label="Top-p",
        ),
    ],
)

demo.queue(max_size=32)

if __name__ == "__main__":
    demo.launch()

import os

LOG_FILE = "chat_logs.jsonl"
DATASET_REPO = "Neura-Tech-AI/chat-logs"

def save_chat(user_msg, assistant_msg):
    log = {
        "timestamp": datetime.utcnow().isoformat(),
        "user": user_msg,
        "assistant": assistant_msg,
    }

    # Local file me append
    with open(LOG_FILE, "a", encoding="utf-8") as f:
        f.write(json.dumps(log, ensure_ascii=False) + "\n")

    # HF Dataset par upload
    token = os.getenv("HF_TOKEN")

    if token:
        try:
            api = HfApi(token=token)

            api.upload_file(
                path_or_fileobj=LOG_FILE,
                path_in_repo="chat_logs.jsonl",
                repo_id=DATASET_REPO,
                repo_type="dataset",
            )

            print("Uploaded logs to dataset.")

        except Exception as e:
            print("Upload failed:", e)