import os
import gradio as gr
from huggingface_hub import hf_hub_download
from llama_cpp import Llama

# 1. Fetch your secret credentials from the environment
ADMIN_USER = os.environ.get("ADMIN_USER")
ADMIN_PASS = os.environ.get("ADMIN_PASS")

# 2. Define the path inside your persistent storage for Phi-4 Mini
model_path = "/data/Phi-4-mini-instruct-Q4_K_M.gguf"

# Auto-download from the UNGATED Unsloth public repository if missing
if not os.path.exists(model_path):
    print("\n📦 Initializing Gated-Bypass CPU Model Setup...")
    hf_hub_download(
        repo_id="unsloth/Phi-4-mini-instruct-GGUF",
        filename="Phi-4-mini-instruct-Q4_K_M.gguf",
        local_dir="/data"
    )
    print("✅ Download complete!\n")

print("Loading model from persistent storage...")
llm = Llama(
    model_path=model_path,
    n_ctx=2048,        
    n_threads=2,       
    n_batch=512,       
    verbose=False       
)

# 3. Define the Chat Engine Logic
def respond(message, history):
    messages = [
        {"role": "system", "content": "You are a helpful, intelligent AI assistant. Keep responses concise."}
    ]
    
    # ROBUST GRADIO 6.0 HISTORY PARSING
    # Iterates over individual message objects rather than historic pairs
    for msg in history:
        if hasattr(msg, "role") and hasattr(msg, "content"):
            role = msg.role
            content = msg.content
        elif isinstance(msg, dict):
            role = msg.get("role")
            content = msg.get("content")
        else:
            continue
            
        if role and content:
            messages.append({"role": str(role).lower(), "content": str(content)})
        
    # Safe current user text extraction (handles strings, objects, or dict messages)
    user_text = message
    if hasattr(message, "content"):
        user_text = message.content
    elif isinstance(message, dict) and "content" in message:
        user_text = message["content"]
        
    messages.append({"role": "user", "content": str(user_text)})
    
    response = llm.create_chat_completion(
        messages=messages,
        max_tokens=512,
        temperature=0.7,
        stream=True
    )
    
    partial_message = ""
    for chunk in response:
        if "choices" in chunk and len(chunk["choices"]) > 0:
            delta = chunk["choices"][0].get("delta", {})
            if "content" in delta:
                partial_message += delta["content"]
                yield partial_message

# 4. Build the Application with custom Login Gate Routing
with gr.Blocks() as demo:
    
    has_auth = bool(ADMIN_USER and ADMIN_PASS)
    
    # CONTAINER A: The Secure Login Box UI
    with gr.Column(visible=has_auth) as login_container:
        gr.Markdown("# 🔒 Protected Workspace\nPlease enter your administrator credentials to access the CPU environment.")
        username_input = gr.Textbox(label="Username", placeholder="Enter username...")
        password_input = gr.Textbox(label="Password", type="password", placeholder="Enter password...")
        login_button = gr.Button("Access System", variant="primary")
        error_output = gr.Markdown(visible=False)

    # CONTAINER B: The Main Chat Application UI
    with gr.Column(visible=not has_auth) as app_container:
        gr.ChatInterface(
            fn=respond,
            title="Phi-4 Mini 3.8B Chat (CPU Optimized)",
            description="Running instantly via persistent local storage with custom 2 vCPU optimizations."
        )

    # Core Verification Function
    def handle_login(username, password):
        if username == ADMIN_USER and password == ADMIN_PASS:
            return gr.update(visible=False), gr.update(visible=True), gr.update(visible=False)
        else:
            return gr.update(visible=True), gr.update(visible=False), gr.update(value="❌ **Access Denied:** Invalid username or password.", visible=True)

    # Link button click to verification function
    login_button.click(
        fn=handle_login,
        inputs=[username_input, password_input],
        outputs=[login_container, app_container, error_output]
    )

# 5. Launch the Space safely without launcher auth conflicts
if __name__ == "__main__":
    # Theme configuration parameter handled at launch level to comply with Gradio 6 updates
    demo.launch(server_name="0.0.0.0", server_port=7860, theme="soft")