import gradio as gr
import os
from groq import Groq

# Password authentication setup
PASSWORD = os.environ.get("APP_PASS")  # Get from Space secrets
authenticated = False  # Track authentication state

# Initialize Groq client
client = Groq(api_key=os.environ.get("GROQ_API_KEY"))  # Get API key from Space secrets

def respond(
    message,
    history: list[tuple[str, str]],
    system_message,
    model_name,
    max_tokens,
    temperature,
    top_p,
):
    messages = [{"role": "system", "content": system_message}]

    for val in history:
        if val[0]:
            messages.append({"role": "user", "content": val[0]})
        if val[1]:
            messages.append({"role": "assistant", "content": val[1]})

    messages.append({"role": "user", "content": message})

    response = ""
    
    # Use Groq's streaming API
    for chunk in client.chat.completions.create(
        model=model_name,
        messages=messages,
        max_tokens=max_tokens,
        stream=True,
        temperature=temperature,
        top_p=top_p,
    ):
        if chunk.choices[0].delta.content is not None:
            token = chunk.choices[0].delta.content
            response += token
            yield response


def check_password(password_attempt):
    if password_attempt == PASSWORD:
        return True, "Login successful! Redirecting to chat..."
    else:
        return False, "Incorrect password. Please try again."


# Authentication interface
with gr.Blocks() as login_interface:
    gr.Markdown("## Authentication Required")
    gr.Markdown("Please enter the password to access this application.")
    password_input = gr.Textbox(type="password", label="Password")
    login_button = gr.Button("Login")
    result = gr.Markdown()
    
    # Chat interface (initially hidden)
    with gr.Accordion("Chat", visible=False) as chat_accordion:
        # Add model selection dropdown
        model_dropdown = gr.Dropdown(
            choices=[
                "llama3-8b-8192",
                "llama3-70b-8192", 
                "mixtral-8x7b-32768",
                "gemma-7b-it"
            ],
            value="mixtral-8x7b-32768",
            label="Model"
        )
        
        chat_interface = gr.ChatInterface(
            respond,
            additional_inputs=[
                gr.Textbox(value="You are a friendly Chatbot.", label="System message"),
                model_dropdown,
                gr.Slider(minimum=1, maximum=4096, value=512, step=1, label="Max new tokens"),
                gr.Slider(minimum=0.0, maximum=2.0, value=0.7, step=0.1, label="Temperature"),
                gr.Slider(
                    minimum=0.1,
                    maximum=1.0,
                    value=0.95,
                    step=0.05,
                    label="Top-p (nucleus sampling)",
                ),
            ],
        )
    
    # Login logic
    def handle_login(password):
        success, message = check_password(password)
        if success:
            return message, gr.update(visible=True), gr.update(visible=False)
        else:
            return message, gr.update(visible=False), gr.update(visible=True)
    
    login_button.click(
        handle_login,
        inputs=[password_input],
        outputs=[result, chat_accordion, password_input],
    )


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