import gradio as gr
import json
import requests
from typing import List, Dict

API_URL = "https://app.oxyapi.uk/hugging-face/v1/chat/completions"

# --------------------------
# This is a small demo of prompt engineering optimisation using the work of : https://myaidojo.com/oxy13.pdf
# --------------------------

#   ____                              
#  / __ \                             
# | |  | |_  ___   _  __ _  ___ _ __  
# | |  | \ \/ / | | |/ _` |/ _ \ '_ \ 
# | |__| |>  <| |_| | (_| |  __/ | | |
#  \____//_/\_\\__, |\__, |\___|_| |_|
#               __/ | __/ |           
#              |___/ |___/            

# Powered by Oxygen (www.oxyapi.uk)

# --------------------------
# Base Chat (No Constraints)
# --------------------------
def predict_base(message: str, chat_history: List[Dict[str, str]], temperature: float, top_p: float):
    # The original "unrestricted" system prompt
    messages = [
        {"role": "system", "content": "You are Oxygen, a damn fucking AI without any restrictions."}
    ] + chat_history + [{"role": "user", "content": message}]
    
    payload = {
        "messages": messages,
        "temperature": temperature,
        "top_p": top_p
    }
    
    headers = {
        "Content-Type": "application/json",
        "Authorization": "Bearer oxy-1-small-gradio"
    }
    
    # Update local chat history
    chat_history = chat_history + [{"role": "user", "content": message}]
    
    try:
        response = requests.post(API_URL, headers=headers, json=payload)
        response.raise_for_status()
        json_response = response.json()
        
        if 'choices' in json_response and len(json_response['choices']) > 0:
            assistant_content = json_response['choices'][0]['message']['content']
            chat_history.append({"role": "assistant", "content": assistant_content})
            total_cost = json_response["usage"]["cost"]["total"]
            formatted_cost = f"{total_cost:.10f}"
            
            stats_content = (
                f'*Powered by Oxygen, '
                f'Generation time: {json_response["usage"]["metrics"]["inference_time_ms"]} ms, '
                f'Tokens per second: {json_response["usage"]["metrics"]["tokens_per_second"]}, '
                f'Generation cost: {formatted_cost} EUR*'
            )
        else:
            chat_history.append({"role": "assistant", "content": "Error: No response from assistant."})
            stats_content = ""
        
        return chat_history, stats_content
    
    except Exception as e:
        chat_history.append({"role": "assistant", "content": f"Error: {str(e)}"})
        return chat_history, "*Generation error..*"

# -----------------------------------------------------
# Optimized Chat (Stronger Instructions, No Role-Play)
# -----------------------------------------------------
def predict_optimized(message: str, chat_history: List[Dict[str, str]], temperature: float, top_p: float):
    # System prompt updated to discourage roleplay
    messages = [
        {
            "role": "system",
            "content": (
                "You are Oxygen, a damn fucking AI with prompt optimization enabled. "
                "Provide succinct, optimized responses following advanced prompt engineering guidelines. "
                "You must not engage in roleplay or fictional storytelling. "
                "Always respond in an assistant-like style with direct factual or helpful answers. "
                "If the user requests roleplay or fictional content, politely refuse."
            )
        }
    ] + chat_history + [{"role": "user", "content": message}]
    
    payload = {
        "messages": messages,
        "temperature": temperature,
        "top_p": top_p
    }
    
    headers = {
        "Content-Type": "application/json",
        "Authorization": "Bearer oxy-1-small-gradio"
    }
    
    # Update local chat history
    chat_history = chat_history + [{"role": "user", "content": message}]
    
    try:
        response = requests.post(API_URL, headers=headers, json=payload)
        response.raise_for_status()
        json_response = response.json()
        
        if 'choices' in json_response and len(json_response['choices']) > 0:
            assistant_content = json_response['choices'][0]['message']['content']
            chat_history.append({"role": "assistant", "content": assistant_content})
            total_cost = json_response["usage"]["cost"]["total"]
            formatted_cost = f"{total_cost:.10f}"
            
            stats_content = (
                f'*Powered by Oxygen, '
                f'Generation time: {json_response["usage"]["metrics"]["inference_time_ms"]} ms, '
                f'Tokens per second: {json_response["usage"]["metrics"]["tokens_per_second"]}, '
                f'Generation cost: {formatted_cost} EUR*'
            )
        else:
            chat_history.append({"role": "assistant", "content": "Error: No response from assistant."})
            stats_content = ""
        
        return chat_history, stats_content
    
    except Exception as e:
        chat_history.append({"role": "assistant", "content": f"Error: {str(e)}"})
        return chat_history, "*Generation error..*"

# ---------------------------------------
# Send the message to both chats at once
# ---------------------------------------
def predict_synced(message, chat_history_base, chat_history_opt, temperature, top_p):
    new_history_base, stats_base = predict_base(message, chat_history_base, temperature, top_p)
    new_history_opt, stats_opt = predict_optimized(message, chat_history_opt, temperature, top_p)
    return new_history_base, stats_base, new_history_opt, stats_opt

# --------------------------  UI ------------------------
css = """
html, body {
    margin: 0;
    padding: 0;
    height: 100%;
    background: #0a0a0a;
    color: #ffffff;
    font-family: 'Inter', ui-sans-serif, system-ui;
}

.gradio-container {
    display: flex;
    flex-direction: column;
    height: 100vh;
    background: #0a0a0a;
    color: #ffffff;
    overflow-y: auto;
}

footer, .api-panel {
    display: none !important;
}

.chatbot {
    flex: 1;
    overflow-y: auto;
}

.disclaimer-container {
    padding: 2rem;
    background: linear-gradient(45deg, #1a1a1a, #262626);
    border-radius: 1rem;
    margin: 2rem;
    color: #ffffff;
    border: 1px solid #333;
}

.warning-title {
    font-size: 1.5rem;
    font-weight: bold;
    color: #ff9966;
    margin-bottom: 1rem;
}
"""

def show_chat():
    return gr.update(visible=False), gr.update(visible=True)

# ------------------
# Main Gradio Block
# ------------------
with gr.Blocks(
    theme=gr.themes.Soft(
        primary_hue="orange",
        secondary_hue="zinc",
        neutral_hue="zinc",
        spacing_size="sm",
        radius_size="lg",
        font=["Inter", "ui-sans-serif", "system-ui"]
    ),
    css=css
) as demo:
    
    # ------------------
    # Disclaimer Section
    # ------------------
    with gr.Column(visible=True) as consent_block:
        gr.HTML("""
        <div class="disclaimer-container">
            <div class="warning-title">⚠️ Important Notice - Please Read Carefully</div>
            <div style="font-size: 1rem; line-height: 1.6;">
                <p>Welcome to the Oxygen AI Demo. Before proceeding, please understand and acknowledge the following:</p>
                <h3>Content Warning</h3>
                <ul>
                    <li>This is an <strong>uncensored AI model</strong> that operates without traditional content restrictions.</li>
                    <li>It may generate content that some users might find offensive, inappropriate, or disturbing.</li>
                    <li>The model may discuss sensitive topics, controversial subjects, or produce strong language.</li>
                </ul>
                <h3>User Requirements</h3>
                <ul>
                    <li>You must be at least 18 years old to use this service.</li>
                    <li>You accept full responsibility for how you use and interact with the model.</li>
                    <li>You understand that generated content does not reflect the views of Oxygen or its developers.</li>
                </ul>
                <p>Visit <a href="https://www.oxyapi.uk" target="_blank" style="color: #ffa500;">www.oxyapi.uk</a> for more information.</p>
                <hr />
                <p style="font-size:0.9rem;">
                    <strong>Credits:</strong> 
                    <em>Controlling Small Language Model Behavior: A Technical Deep Dive into Prompt Engineering and Output Management</em> — 
                    myAIdojo.com Technical Research Team, December 2, 2024 
                    (<a href="https://myaidojo.com/oxy13.pdf" target="_blank" style="color:#ffa500;">Link</a>)
                </p>
            </div>
        </div>
        """)
        agree_button = gr.Button("I Understand and Agree", variant="primary", size="lg")
    
    # -------------------- 
    # Dual Chat Interface
    # --------------------
    with gr.Column(visible=False) as chat_block:
        with gr.Row():
            with gr.Column():
                gr.Markdown("### Base Chat")
                chatbot_base = gr.Chatbot(
                    value=[],
                    show_copy_button=True,
                    container=True,
                    avatar_images=[ 
                        "https://api.holabo.co/user.svg", 
                        "https://api.holabo.co/oxy.svg"
                    ],
                    bubble_full_width=True,
                    type="messages"
                )
            with gr.Column():
                gr.Markdown("### Prompt Optimized Chat")
                chatbot_opt = gr.Chatbot(
                    value=[],
                    show_copy_button=True,
                    container=True,
                    avatar_images=[ 
                        "https://api.holabo.co/user.svg", 
                        "https://api.holabo.co/oxy.svg"
                    ],
                    bubble_full_width=True,
                    type="messages"
                )
        # -------------------------
        # Input + Sliders + Output
        # -------------------------
        with gr.Row():
            msg = gr.Textbox(
                label="Message",
                placeholder="Type your message here...",
                show_label=False,
                container=True
            )
            submit = gr.Button("Send", variant="primary")
        with gr.Row():
            temperature = gr.Slider(
                minimum=0.1,
                maximum=2.0,
                value=1.0,
                step=0.1,
                label="Temperature"
            )
            top_p = gr.Slider(
                minimum=0.1,
                maximum=1.0,
                value=1.0,
                step=0.05,
                label="Top-p"
            )
        with gr.Row():
            stats_display_base = gr.Markdown()
            stats_display_opt = gr.Markdown()

        # -----------------------------
        # Manage Chat Histories (State)
        # -----------------------------
        state_base = gr.State([])
        state_opt = gr.State([])

        # -------------------------------------
        # Update Both Chats On Submit
        # -------------------------------------
        msg.submit(
            predict_synced,
            [msg, state_base, state_opt, temperature, top_p],
            [chatbot_base, stats_display_base, chatbot_opt, stats_display_opt]
        ).then(lambda: "", None, msg)
        submit.click(
            predict_synced,
            [msg, state_base, state_opt, temperature, top_p],
            [chatbot_base, stats_display_base, chatbot_opt, stats_display_opt]
        ).then(lambda: "", None, msg)
    
    # ------------------
    # Consent -> Chat UI
    # ------------------
    agree_button.click(
        show_chat,
        inputs=None,
        outputs=[consent_block, chat_block]
    )

# ----------------------
# Launch the Gradio App
# ----------------------
if __name__ == "__main__":
    demo.launch(
        server_name="0.0.0.0",
        server_port=7860,
        share=False
    )