import gradio as gr

# Hardcoded credentials
VALID_CREDENTIALS = {
    "admin": "officechat888",
    "vip": "vip888",
    "user": "user888"
}

# Models loading code remains the same
model1 = gr.load("models/Jonny001/NSFW_master")
model2 = gr.load("models/Jonny001/Alita-v1")
model3 = gr.load("models/Jonny001/WBMH-v1.1")
model5 = gr.load("models/prashanth970/flux-lora-uncensored")

def check_login(username, password):
    if username in VALID_CREDENTIALS and VALID_CREDENTIALS[username] == password:
        return True
    return False

def generate_images(text, selected_model):
    if selected_model == "模型1 (NSFW Master)":
        model = model1
    elif selected_model == "模型2 (Alita)":
        model = model2
    elif selected_model == "模型3 (WBMH-v1.1)":
        model = model3
    elif selected_model == "模型5 (Lora Uncensored)":
        model = model5
    else:
        return "无效的模型选择。"
    
    results = []
    for i in range(3):
        modified_text = f"{text} 变体 {i+1}"
        result = model(modified_text)
        results.append(result)
    
    return results

custom_css = """
/* Login page styles */
.login-container {
    max-width: 400px !important;
    margin: 100px auto !important;
    padding: 30px !important;
    background: #ffffff !important;
    border-radius: 15px !important;
    box-shadow: 0 4px 20px rgba(0,0,0,0.1) !important;
}

.login-title {
    text-align: center !important;
    font-size: 24px !important;
    color: #333 !important;
    margin-bottom: 30px !important;
}

.login-button {
    background: #4a90e2 !important;
    width: 100% !important;
    margin-top: 20px !important;
}

/* Main interface styles */
.gradio-container {
    max-width: 1200px !important;
    margin: 0 auto !important;
    padding: 20px !important;
}

.control-panel {
    background: #ffffff !important;
    padding: 30px;
    border-radius: 15px;
    box-shadow: 0 2px 10px rgba(0,0,0,0.05);
    margin-bottom: 30px;
    border: 1px solid #e1e4e8 !important;
}

.title {
    font-size: 18px !important;
    margin-bottom: 25px !important;
    text-align: center !important;
    font-weight: bold !important;
    color: #24292e !important;
    background: #ffffff !important;
    padding: 15px !important;
    border-radius: 10px !important;
    border: 1px solid #e1e4e8 !important;
}

.input-area {
    display: flex;
    gap: 20px;
    align-items: flex-start;
}

.gradio-button {
    background: #4a90e2 !important;
    color: white !important;
    padding: 12px 30px !important;
    border-radius: 8px !important;
    border: none !important;
    cursor: pointer !important;
    transition: all 0.3s !important;
    font-size: 16px !important;
}

.results-panel {
    display: flex;
    flex-direction: column;
    gap: 30px;
}

.gradio-image {
    background: #ffffff !important;
    padding: 15px !important;
    border-radius: 15px !important;
    box-shadow: 0 3px 15px rgba(0,0,0,0.1) !important;
    min-height: 400px !important;
    width: 100% !important;
    border: 1px solid #e1e4e8 !important;
}

.gradio-image img {
    width: 100% !important;
    height: auto !important;
    min-height: 400px !important;
    object-fit: contain !important;
}

/* Dark mode styles */
@media (prefers-color-scheme: dark) {
    body {
        background: #0d1117 !important;
    }
    
    .login-container, .control-panel, .gradio-image {
        background: #161b22 !important;
        border: 1px solid #30363d !important;
    }
    
    .login-title, .title {
        color: #c9d1d9 !important;
        background: #161b22 !important;
        border: 1px solid #30363d !important;
    }
    
    label, input, textarea {
        color: #c9d1d9 !important;
    }
    
    input, textarea {
        background: #0d1117 !important;
        border: 1px solid #30363d !important;
    }
}
"""

with gr.Blocks(css=custom_css) as interface:
    # Login state
    logged_in = gr.State(False)
    
    # Login page
    with gr.Column(visible=True) as login_page:
        with gr.Column(elem_classes=["login-container"]):
            gr.Markdown("### Office Chat VIP登录", elem_classes=["login-title"])
            username = gr.Textbox(label="用户名", placeholder="请输入用户名...")
            password = gr.Textbox(label="密码", placeholder="请输入密码...", type="password")
            login_button = gr.Button("登录", elem_classes=["login-button"])
            login_status = gr.Markdown("")

    # Main interface
    with gr.Column(visible=False) as main_interface:
        with gr.Column(elem_classes=["control-panel"]):
            gr.Markdown(
                "### 本程序专为Office chat 超级VIP使用用户免费使用。因为免费，可能会比较慢，感谢您的理解。感谢您选择Office chat。", 
                elem_classes=["title"]
            )
            with gr.Row(elem_classes=["input-area"]):
                text_input = gr.Textbox(
                    label="在这里输入你的想象(英文):",
                    placeholder="输入你的提示...",
                    scale=1
                )
                model_selection = gr.Radio(
                    ["模型1 (NSFW Master)", "模型2 (Alita)", "模型3 (WBMH-v1.1)", 
                     "模型5 (Lora Uncensored)"],
                    label="选择模型（尝试所有模型并获得不同结果）",
                    value="模型1 (NSFW Master)",
                    scale=2
                )
            generate_button = gr.Button("生成图片")

        with gr.Column(elem_classes=["results-panel"]):
            image1 = gr.Image(label="生成的图片1")
            image2 = gr.Image(label="生成的图片2")
            image3 = gr.Image(label="生成的图片3")

    def login(username, password):
        if check_login(username, password):
            return {
                login_page: gr.update(visible=False),
                main_interface: gr.update(visible=True),
                login_status: ""
            }
        else:
            return {
                login_page: gr.update(visible=True),
                main_interface: gr.update(visible=False),
                login_status: "用户名或密码错误！"
            }

    login_button.click(
        login,
        inputs=[username, password],
        outputs=[login_page, main_interface, login_status]
    )

    generate_button.click(
        fn=generate_images,
        inputs=[text_input, model_selection],
        outputs=[image1, image2, image3]
    )

interface.launch()