import os
from datetime import datetime
import uuid
import gradio as gr
from transformers import AutoTokenizer, AutoModelForCausalLM
import torch
from huggingface_hub import login

from dotenv import load_dotenv
# Load environment variables
load_dotenv()

# Authenticate with Hugging Face
login(token=os.getenv("HUGGINGFACE_TOKEN"))

# Load model and tokenizer
model_name = "meta-llama/Meta-Llama-3.1-8B-Instruct"

tokenizer = AutoTokenizer.from_pretrained(model_name, token=True)
model = AutoModelForCausalLM.from_pretrained(model_name, torch_dtype=torch.bfloat16, device_map="auto", token=True)


# Set pad_token_id if it's not already set
if tokenizer.pad_token_id is None:
    tokenizer.pad_token_id = tokenizer.eos_token_id

def chat_with_model(messages):
    # Prepare the input
    input_ids = tokenizer.encode(str(messages), return_tensors="pt").to(model.device)
    attention_mask = torch.ones_like(input_ids)
    
    # Generate response
    with torch.no_grad():
        output = model.generate(
            input_ids,
            attention_mask=attention_mask,
            max_length=1000,
            num_return_sequences=1,
            temperature=0.7,
            pad_token_id=tokenizer.pad_token_id
        )
    
    response = tokenizer.decode(output[0], skip_special_tokens=True)
    return response

def chat_with_model_gradio(message, history, session_id):
    messages = [
        {"role": "system", "content": f"너의 이름은 ChatMBTI. 사람들의 MBTI유형에 알맞은 상담을 진행할 수 있어. 상대방의 MBTI 유형을 먼저 물어보고, 그 유형에 알맞게 상담을 진행해줘. 참고로 현재 시각은 {datetime.now().strftime('%Y-%m-%d %H:%M:%S')}이야."},
    ]
    messages.extend([{"role": "user" if i % 2 == 0 else "assistant", "content": m} for i, m in enumerate(history)])
    messages.append({"role": "user", "content": message})

    response = chat_with_model(messages)
    history.append((message, response))

    return "", history

def main():
    session_id = str(uuid.uuid4())
    with gr.Blocks() as demo:
        chatbot = gr.Chatbot(label="ChatMBTI")
        msg = gr.Textbox(label="메시지를 입력하세요")
        clear = gr.Button("대화 초기화")

        msg.submit(chat_with_model_gradio, [msg, chatbot, gr.State(session_id)], [msg, chatbot])
        clear.click(lambda: None, None, chatbot, queue=False)

    demo.launch()

if __name__ == "__main__":
    main()