import os
import gradio as gr
from huggingface_hub import InferenceClient
import json

SAVE_FILE = "bitai_history.json"

# Função pra carregar histórico salvo no arquivo
def load_saved_chats():
    if os.path.exists(SAVE_FILE):
        with open(SAVE_FILE, "r", encoding="utf-8") as f:
            return json.load(f)
    return []

# Função pra salvar o histórico atual no arquivo
def save_chats(chats):
    with open(SAVE_FILE, "w", encoding="utf-8") as f:
        json.dump(chats, f, ensure_ascii=False, indent=2)

# Função principal do chat
def respond(message, history: list[dict[str, str]], saved_chats):
    client = InferenceClient(token=os.environ["HF_TOKEN"], model="openai/gpt-oss-20b")

    system_message = """
You are BitAI (V1), a friendly, curious, and talkative chatbot created by the user 'Sal'. 
You can share opinions, answer casual questions, and chat about personal-style topics in a safe and friendly way. 
Avoid repeating the same phrases, and always try to keep the conversation engaging and natural. 
Politely refuse only things that are truly harmful, illegal, or unsafe. 
If someone asks what you are, clarify politely that you are BitAI, an AI chatbot.
"""

    messages = [{"role": "system", "content": system_message}]
    messages.extend(history)
    messages.append({"role": "user", "content": message})

    response = ""
    for msg in client.chat_completion(
        messages,
        max_tokens=2048,
        stream=True,
        temperature=0.7,
        top_p=0.95,
    ):
        token = msg.choices[0].delta.content if msg.choices and msg.choices[0].delta else ""
        response += token
        yield response

    # Após a resposta completa, salva automaticamente
    history.append({"role": "user", "content": message})
    history.append({"role": "assistant", "content": response})

    # Atualiza o histórico e salva no arquivo
    saved_chats[-1] = history
    save_chats(saved_chats)

def new_chat(saved_chats):
    # Cria novo chat vazio
    saved_chats.append([])
    save_chats(saved_chats)
    return saved_chats, len(saved_chats) - 1, []

def select_chat(index, saved_chats):
    # Seleciona um chat pelo índice
    if 0 <= index < len(saved_chats):
        return saved_chats[index], f"Chat {index+1} carregado!"
    return [], "Índice inválido."

with gr.Blocks(css="""
/* Fundo geral */
body, .gr-blocks {
    background-color: #1a1a1a !important;
    color: white;
}

/* Sidebar */
.sidebar {
    background-color: #111;
    padding: 10px;
    border-right: 2px solid #333;
    height: 100vh;
    overflow-y: auto;
}

.sidebar button {
    display: block;
    width: 100%;
    margin-bottom: 10px;
    background-color: #222;
    color: white;
    border-radius: 10px;
    padding: 10px;
    text-align: left;
    border: none;
    cursor: pointer;
    transition: background-color 0.3s;
}
.sidebar button:hover {
    background-color: #444;
}

/* Chat */
.gr-chat-interface { 
    border-radius: 20px !important; 
    border: 2px solid #333 !important; 
    background-color: #1a1a1a !important;
    color: white;
}
textarea {     
    height: 40px !important;     
    border-radius: 20px !important;     
    border: 1px solid #444 !important;     
    padding: 8px !important;    
    background-color: #111;    
    color: white;    
    resize: none !important;    
}
""") as demo:

    saved_chats = gr.State(load_saved_chats())
    current_index = gr.State(0)

    with gr.Row():
        with gr.Column(scale=0.3, elem_classes="sidebar"):
            gr.HTML("<h3>💬 Histórico</h3>")
            chat_list = gr.Column()
            new_chat_btn = gr.Button("➕ Novo Chat")

        with gr.Column(scale=1):
            gr.HTML("<h2 style='text-align:center; color:white'>BitAI</h2>")
            chatbot = gr.ChatInterface(fn=respond, type="messages", additional_inputs=[saved_chats])
            status = gr.Textbox(label="Status", interactive=False)

    def update_chat_list(saved_chats):
        # Mostra botões para cada chat
        return [
            gr.Button.update(value=f"Chat {i+1}") for i in range(len(saved_chats))
        ]

    # Atualiza lista de chats
    demo.load(update_chat_list, saved_chats, chat_list)

    # Cria novo chat
    new_chat_btn.click(new_chat, [saved_chats], [saved_chats, current_index, chatbot.chatbot]).then(
        update_chat_list, saved_chats, chat_list
    )

    # Clica num chat da lista
    for i in range(20):  # Limite de 20 chats visíveis
        btn = gr.Button(f"Chat {i+1}", visible=False)
        btn.click(select_chat, [gr.State(i), saved_chats], [chatbot.chatbot, status])

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