import gradio as gr
import requests
import json
import os
import markdown
import gradio as gr
import re
import markdown
from pygments import highlight
from pygments.lexers import HtmlLexer, PythonLexer, TextLexer
from pygments.formatters import HtmlFormatter

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 render_avatars(userid):
    try:
        response = requests.post(
            'https://host.palple.polrambora.com/userexistence',
            json={
                'userid': userid
            },
            timeout=10
        )

        if response.status_code == 200:
            response_json = response.json()
            return response_json["avatar"]["link"]
    except Exception as e:
        return None
    except requests.exceptions.Timeout:
        return None


def authorize(user, api_key, system_message):
    test_data = {
        "user": user,
        "key": api_key
    }

    try:
        response = requests.post(
            "https://host.palple.polrambora.com/check_key_impv",
            json=test_data,
        )

        if response.status_code == 200:
            response_json = response.json()
            avatar = render_avatars(user) or USER_PIC_PATH
            
            if api_key not in sessions:
                sessions[api_key] = {
                    "history": [], 
                    "headers": {   
                        "authorization": api_key,
                        "Content-Type": 'application/json'
                    },
                    "avatar": avatar,
                    "system_message": system_message 
                }
            return True
        elif response.status_code == 403:
            return 403
        else:
            return False
    except Exception as 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", sessions[api_key]["avatar"], ASSISTANT_PIC_PATH))
        sessions[api_key]["history"] = history  
        return history, assistant_reply  
    else:
        return history, "Error: " + response.json().get("error", "Unknown error occurred.")


def detect_content_type(message):
    """
    Automatically detect whether the message contains Math, HTML, or plain Markdown.
    Returns a tuple: (is_math, is_html)
    """
    is_math = bool(re.search(r"(\$\$.*?\$\$|\$.*?\$)", message, re.DOTALL))
    
    is_html = bool(re.search(r"<\/?[a-z][\s\S]*?>", message))
    
    return is_math, is_html

def render_message(message):
    """
    Render a message dynamically, detecting if it's Math, HTML, or Markdown with code.
    """
    is_math, is_html = detect_content_type(message)

    if is_math:
        return gr.HTML(f"""
        <div style="font-size: 16px; font-family: Arial;">
            <script type="text/javascript" async src="https://cdn.jsdelivr.net/npm/mathjax@3/es5/tex-mml-chtml.js"></script>
            {message}
        </div>
        """)
    elif is_html:
        formatter = HtmlFormatter(style="default", noclasses=True)
        highlighted_html = highlight(message, HtmlLexer(), formatter)
        return gr.HTML(f"""
        <div style="background: #f5f5f5; border: 1px solid #ddd; padding: 10px; overflow-x: auto;">
            {highlighted_html}
        </div>
        """)
    else:
        def markdown_with_code_blocks(md):
            html = markdown.markdown(md, extensions=["fenced_code"])
            return html

        formatted_message = markdown_with_code_blocks(message)
        return gr.HTML(f"""
        <div style="font-size: 16px; font-family: Arial;">
            {formatted_message}
        </div>
        """)

def escape_html(unsafe_text):
    """Escape HTML special characters to prevent code injection."""
    return (
        unsafe_text.replace("&", "&amp;")
        .replace("<", "&lt;")
        .replace(">", "&gt;")
        .replace('"', "&quot;")
        .replace("'", "&#039;")
    )

js = """
function HTMLClean() {
    alert("This page is in-dev stage, expect lots of restarts, the interface is open-sourced (not the backend that is responsible for API Keys)")
 }
"""
css="""
.chatbox {height: 400px; overflow: auto; border: 1px solid #262626; padding: 10px; background-color: #171717; display: flex;  flex-direction: column-reverse;}
#chatbox-container .katex-display { display: inline; }
#chatbox-container .katex-display>.katex { display: inline; }
#chatbox-container .katex-display>.katex>.katex-html { display: inline; }
"""
with gr.Blocks(css=css, js=js) 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.")
        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-container")

        msg_input = gr.Text(
            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):
            loading_message = history + [(message, "Loading...", "You", "P-ALPLE", sessions[api_key]["avatar"], ASSISTANT_PIC_PATH)]
            yield render_message(loading_message), loading_message, ""

            history, assistant_reply = respond(message, api_key, max_tokens, top_p, temperature)
            yield render_message(history), history, ""

        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],
        )

        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],
        )

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

    with gr.Column(visible=False) as blacklist_view:
        gr.Markdown("## P-MSQ Authorization")
        gr.Markdown("Your linked ID appears to be blacklisted, and your API Key is pending on removal, if you believe this is a mistake, please try reaching us out.")
    
    
    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
            )
        elif authorize(user, api_key, PRIMARY_SYSTEM_INSTRUCTIONS) == 403: 
            return (
                gr.update(visible=False),  
                gr.update(visible=False), 
                gr.update(visible=True), 
            )
        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.launch(show_api=False)
if __name__ == "__main__":
    demo.queue = False