import gradio as gr
import requests
import json
import os

API_URL = "https://host.palple.polrambora.com/pmsq"

sessions = {}
PRIMARY_SYSTEM_INSTRUCTIONS = "You are P-MSQ (Messaging Service Query), a friendly AI Chatbot that can help in any situations"
ASSISTANT_PIC_PATH = "https://huggingface.co/spaces/PLRMB/P-MSQ-API-PREVIEW/resolve/main/API.png"  
USER_PIC_PATH = "https://huggingface.co/spaces/PLRMB/P-MSQ-API-PREVIEW/resolve/main/usr.png"

def authorize(user, api_key, system_message):
    test_data = {
        "user": user,
        "key": api_key
    }
    test_headers = {
        'Accept': 'application/json'
    }

    print("Preparing to send the request...") 
    try:
        response = requests.post(
            "https://host.palple.polrambora.com/check_key_impv",
            json=test_data,
        )

        if response.status_code == 200:
            response_json = response.json()
            print(f"Response: {response_json}")
            
            if api_key not in sessions:
                sessions[api_key] = {
                    "history": [], 
                    "headers": {   
                        "authorization": api_key,
                        "Content-Type": 'application/json'
                    },
                    "system_message": system_message 
                }
                
            return True
        else:
            print(f"Failed request - Status code: {response.status_code}, Response: {response.text}")
            return False
    except requests.exceptions.Timeout:
        print("Request timed out")
        return False
    except requests.exceptions.ConnectionError:
        print("Network problem occurred (DNS failure, refused connection, etc.)")
        return False
    except Exception as e:
        print(f"Exception occurred: {str(e)}")
        return False


def respond(message, api_key, max_tokens, top_p, temperature):
    session = sessions.get(api_key, {})
    history = session.get("history", [])
    headers = session.get("headers", {})
    system_message = session.get("system_message", PRIMARY_SYSTEM_INSTRUCTIONS)
    messages = []
    for user_message, assistant_message, user_profile, assistant_profile, user_pic, assistant_pic in history:
        if user_message:
            messages.append({
                "role": "user",
                "content": user_message,
                "profile": user_profile,
                "picture": user_pic
            })
        if assistant_message:
            messages.append({
                "role": "assistant",
                "content": assistant_message,
                "profile": assistant_profile,
                "picture": assistant_pic
            })

    data = {
        "preferences": {
            "max_char": max_tokens,
            "temperature": temperature,
            "top_p": top_p,
            "system_message": system_message
        },
        "conversation_history": messages,
        "input": message
    }
    
    response = requests.post(API_URL, headers=headers, data=json.dumps(data))
    
    if response.status_code == 200:
        response_json = response.json()
        assistant_reply = response_json["msq"]["message"][0]
        

        history.append((message, assistant_reply, "You", "P-ALPLE", USER_PIC_PATH, ASSISTANT_PIC_PATH))
        sessions[api_key]["history"] = history  
        
        return history, assistant_reply  
    else:
        return history, "Error: " + response.json().get("error", "Unknown error occurred.")



def render_message(history):
    messages_html = """
    <div id="chatbox-container" class="chatbox" style="height: 400px; overflow-y: auto; border: 1px solid #262626; padding: 10px; background-color: #171717;">
        <div id="messages">"""

    for user_message, assistant_message, user_profile, assistant_profile, user_pic, assistant_pic in history:
        if user_message:
            messages_html += f"<div style='display: flex; align-items: center; margin-bottom: 10px;'>"
            if user_pic:
                messages_html += f"<img src='{user_pic}' style='width: 40px; height: 40px; border-radius: 50%; margin-right: 10px;'>"
            messages_html += f"<span style='color: white;'>{user_message}</span></div><br>"

        if assistant_message:
            messages_html += f"<div style='display: flex; align-items: center; margin-bottom: 10px;'>"
            if assistant_pic:
                messages_html += f"<img src='{assistant_pic}' style='width: 40px; height: 40px; border-radius: 50%; margin-right: 10px;'>"
            messages_html += f"<span style='color: white;'>{assistant_message}</span></div><br>"

    messages_html += "</div></div>" 

    scroll_script = """
    <script>
    const container = document.getElementById('chatbox-container');

    const observer = new MutationObserver(() => {
        container.scrollTop = container.scrollHeight;  // Force scroll to bottom after re-render
    });

    observer.observe(container, { childList: true, subtree: true });

    window.onload = function() {
        container.scrollTop = container.scrollHeight;
    };
    </script>
    """

    return messages_html + scroll_script

with gr.Blocks(css=".chatbox {height: 400px; overflow-y: auto; border: 1px solid #262626; padding: 10px; background-color: #171717;}") as demo:

    with gr.Column(visible=True) as auth_view:
        gr.Markdown("## P-MSQ Authorization")
        gr.Markdown("P-MSQ is in closed alpha test! The model, api and more are subject to change. Even though the space is public, P-MSQ is restricted for public use, available for testers.")
        api_user_input = gr.Textbox(placeholder="snowflake", label="UserID", type='email')
        api_key_input = gr.Textbox(placeholder="Enter your API key", label="Token", type='password')
        auth_button = gr.Button("Authorize")
        auth_status = gr.Textbox(label="Authorization Status", interactive=False)

    with gr.Column(visible=False) as chat_view:
        gr.Markdown("## P-MSQ Chat Interface")

        chatbot_output = gr.HTML(elem_id="chatbox")

        msg_input = gr.Textbox(
            show_label=False, 
            placeholder="Type your message and press Shift+Enter...", 
            lines=2,
            elem_id="input-text"
        )

        send_btn = gr.Button("Send")
        regen_btn = gr.Button("Clear")

        system_instructions_input = gr.Textbox(placeholder="Enter custom instructions (optional)", 
                                               label="Custom System Instructions", 
                                               lines=2)
        save_instructions_btn = gr.Button("Save Instructions")        
        gr.Markdown("### Settings")
        
        max_tokens = gr.Slider(minimum=1, maximum=2048, value=1024, step=1, label="Max new tokens")
        top_p = gr.Slider(minimum=0, maximum=2, value=0.8, step=0.1, label="Top P")
        temperature = gr.Slider(minimum=0.1, maximum=1, value=0.7, step=0.1, label="Temperature")

        history_state = gr.State([])  
        last_message_state = gr.State("")  

        def user_interaction(message, history, api_key, max_tokens, top_p, temperature):
           history, assistant_reply = respond(message, api_key, max_tokens, top_p, temperature)

           full_history_html = render_message(history)

           return full_history_html, history, "", message



        def regenerate_response(history, last_message, max_tokens, top_p, temperature):
            return "", []

        def clear_history(api_key):
            if api_key in sessions:
                sessions[api_key]["history"] = []  
            return "", []

        def load_conversation(api_key):
            session = sessions.get(api_key, {})
            history = session.get("history", [])
            return render_message(history), history

        msg_input.submit(user_interaction, 
                 inputs=[msg_input, history_state, api_key_input, max_tokens, top_p, temperature], 
                 outputs=[chatbot_output, history_state, msg_input, last_message_state])

        send_btn.click(user_interaction, 
                inputs=[msg_input, history_state, api_key_input, max_tokens, top_p, temperature], 
                outputs=[chatbot_output, history_state, msg_input, last_message_state])

        regen_btn.click(clear_history, 
                        inputs=[api_key_input], 
                        outputs=[chatbot_output, history_state])

    def authorize_and_proceed(user, api_key):
        if authorize(user, api_key, PRIMARY_SYSTEM_INSTRUCTIONS):
            gr.Info("Loading, please wait.")
            messages_html, history = load_conversation(api_key)
            return gr.update(visible=False), gr.update(visible=True), messages_html, history
        else:
            gr.Warning("Incorrect userid/token")
            return gr.update(visible=True), gr.update(visible=False), auth_status.update(value="Invalid userid/token")

    def save_custom_instructions(api_key, custom_instructions):
        if api_key in sessions:
            gr.Info("Instructions updated, we recommend to start the new conversation to make it more efficient.")
            sessions[api_key]["system_message"] = custom_instructions  
            return "Instructions updated!", gr.update(value="")
        else:
            gr.Warning("Your session has been expired, please refresh the page and login again.")
            return "Session not found.", gr.update(value="")

            
    auth_button.click(authorize_and_proceed, inputs=[api_user_input, api_key_input], outputs=[auth_view, chat_view, chatbot_output, history_state])
    save_instructions_btn.click(save_custom_instructions, inputs=[api_key_input, system_instructions_input], outputs=auth_status)
    demo.load(lambda: None, [], None)
    demo.launch(
        js="""
    function scrollToBottom() {
        const container = document.getElementById('chatbox-container');
        if (container) {
            container.scrollTop = container.scrollHeight;
        }
    }

    window.onload = scrollToBottom;

    const observer = new MutationObserver(scrollToBottom);
    observer.observe(document.getElementById('chatbox-container'), { childList: true, subtree: true });
        """
    )
if __name__ == "__main__":
    demo.queue = False