app.py

import os
import json
import hashlib
from cryptography.fernet import Fernet
from cryptography.hazmat.primitives import hashes
from cryptography.hazmat.primitives.kdf.pbkdf2 import PBKDF2HMAC
import base64
import openai
import gradio as gr
from epub2txt import epub2txt

# --- Configuration and Constants ---
MODEL_NAME = os.getenv("POE_MODEL", "GPT-5-mini")
PROMPT = os.getenv("prompt", "Summarize the following text:")
KEYS_FILE = "user_keys.json"

# --- Helper Functions for Security and User Management ---
# By moving logic out of a class, we avoid complex state issues.

def get_user_id(username):
    """Generates a unique, stable ID from a username."""
    return hashlib.sha256(username.encode()).hexdigest()

def generate_fernet_key(user_id):
    """Generates a deterministic encryption key from a user ID."""
    salt = user_id.encode()[:16].ljust(16, b'\0') # Salt must be 16 bytes
    kdf = PBKDF2HMAC(
        algorithm=hashes.SHA256(),
        length=32,
        salt=salt,
        iterations=100000,
    )
    key = base64.urlsafe_b64encode(kdf.derive(user_id.encode()))
    return Fernet(key)

def load_user_keys():
    """Loads the entire user key database."""
    if os.path.exists(KEYS_FILE):
        with open(KEYS_FILE, 'r') as f:
            return json.load(f)
    return {}

def save_user_keys(keys_data):
    """Saves the entire user key database."""
    with open(KEYS_FILE, 'w') as f:
        json.dump(keys_data, f)

def get_decrypted_api_key(user_id):
    """Retrieves and decrypts a single user's API key."""
    keys_data = load_user_keys()
    encrypted_key_b64 = keys_data.get(user_id)
    if encrypted_key_b64:
        try:
            cipher = generate_fernet_key(user_id)
            encrypted_key = base64.urlsafe_b64decode(encrypted_key_b64.encode())
            return cipher.decrypt(encrypted_key).decode()
        except Exception as e:
            print(f"Decryption failed for user {user_id}: {e}")
    return None

def set_encrypted_api_key(user_id, api_key):
    """Encrypts and saves a single user's API key."""
    keys_data = load_user_keys()
    cipher = generate_fernet_key(user_id)
    encrypted_key = cipher.encrypt(api_key.encode())
    keys_data[user_id] = base64.urlsafe_b64encode(encrypted_key).decode()
    save_user_keys(keys_data)

def delete_api_key(user_id):
    """Deletes a user's API key."""
    keys_data = load_user_keys()
    if user_id in keys_data:
        del keys_data[user_id]
        save_user_keys(keys_data)

# --- Gradio Event Handlers ---

def initialize_client(api_key):
    """Initializes and returns an OpenAI client or None on failure."""
    if not api_key:
        return None
    try:
        client = openai.OpenAI(api_key=api_key, base_url="https://api.poe.com/v1")
        # Test connection
        client.models.list()
        return client
    except Exception as e:
        print(f"Failed to initialize client: {e}")
        return None

def update_ui_on_load(request: gr.Request):
    """
    This function runs after the user logs in and the page loads.
    It configures the UI based on whether the user has a saved API key.
    """
    user_id = get_user_id(request.username)
    api_key = get_decrypted_api_key(user_id)
    client = initialize_client(api_key)

    if client:
        # Key is valid and loaded
        welcome_msg = f"Welcome back, **{request.username}**! Your saved API key is loaded."
        return {
            welcome_md: gr.update(value=welcome_msg),
            api_key_section: gr.update(visible=False),
            inp: gr.update(visible=True),
            clear_key_btn: gr.update(visible=True)
        }
    else:
        # No key or invalid key found
        welcome_msg = f"Welcome, **{request.username}**! Please provide your Poe API key to begin."
        return {
            welcome_md: gr.update(value=welcome_msg),
            api_key_section: gr.update(visible=True),
            inp: gr.update(visible=False),
            clear_key_btn: gr.update(visible=False)
        }

def set_api_key(api_key, remember_key, request: gr.Request):
    """Handles the 'Set API Key' button click."""
    user_id = get_user_id(request.username)
    client = initialize_client(api_key)

    if not client:
        return {out: gr.update(value="❌ Invalid or incorrect API key. Please check it and try again.")}

    if remember_key:
        set_encrypted_api_key(user_id, api_key)
        msg = "✅ API key validated and saved! You can now upload an ePub."
    else:
        delete_api_key(user_id) # Ensure no old key is stored if user unchecks
        msg = "✅ API key validated for this session. You can now upload an ePub."

    return {
        out: gr.update(value=msg),
        inp: gr.update(visible=True),
        api_key_section: gr.update(visible=False),
        clear_key_btn: gr.update(visible=remember_key)
    }

def clear_saved_key(request: gr.Request):
    """Handles the 'Clear Saved Key' button click."""
    user_id = get_user_id(request.username)
    delete_api_key(user_id)
    return {
        out: gr.update(value="✅ Saved API key cleared. Please enter an API key to continue."),
        api_key_section: gr.update(visible=True),
        inp: gr.update(visible=False),
        clear_key_btn: gr.update(visible=False),
        api_key_input: gr.update(value="")
    }

def process_epub(file, request: gr.Request):
    """Processes the uploaded ePub file. This is a generator function."""
    user_id = get_user_id(request.username)
    api_key = get_decrypted_api_key(user_id)
    
    # Re-initialize client to ensure API key is available for this long-running task
    client = initialize_client(api_key)
    if not client:
        yield "⚠️ Your API key is missing or invalid. Please clear and set your key again."
        return
    
    # ... (The rest of your ePub processing logic is identical) ...
    try:
        ch_list = epub2txt(file.name, outputlist=True)
        title = epub2txt.title
        yield f"# {title}\n\nProcessing ePub file..."
        sm_list = []
        for text in ch_list[2:]:
            if text.strip():
                response = client.chat.completions.create(
                    model=MODEL_NAME,
                    messages=[{"role": "user", "content": PROMPT + "\n\n" + text[:4000]}], # Simple chunking
                ).choices[0].message.content
                sm_list.append(response)
        
        yield f"# {title}\n\n" + "\n\n---\n\n".join(sm_list)

    except Exception as e:
        yield f"An error occurred: {e}"


# --- Build the Gradio Interface ---

with gr.Blocks(title="ePub Summarizer") as demo:
    welcome_md = gr.Markdown("Welcome! Please log in to continue.")

    # API Key Section (initially hidden, made visible by update_ui_on_load)
    with gr.Column(visible=False) as api_key_section:
        api_key_input = gr.Textbox(label="Poe API Key", type="password")
        remember_key = gr.Checkbox(label="Remember my API key", value=True)
        api_key_btn = gr.Button("Set API Key", variant="primary")

    # Main App Section
    out = gr.Markdown()
    inp = gr.File(file_types=['.epub'], visible=False, label="Upload ePub File")
    clear_key_btn = gr.Button("Clear Saved Key", variant="secondary", visible=False)

    # --- Event Wiring ---

    # 1. After login, the 'load' event fires, calling update_ui_on_load
    demo.load(
        fn=update_ui_on_load,
        outputs=[welcome_md, api_key_section, inp, clear_key_btn]
    )

    # 2. User interactions trigger other handlers
    api_key_btn.click(
        fn=set_api_key,
        inputs=[api_key_input, remember_key],
        outputs=[out, inp, api_key_section, clear_key_btn]
    )
    clear_key_btn.click(
        fn=clear_saved_key,
        outputs=[out, api_key_section, inp, clear_key_btn, api_key_input]
    )
    inp.change(fn=process_epub, inputs=[inp], outputs=[out])

# --- Main Execution ---

if __name__ == "__main__":
    # This dummy function is used for local testing.
    # On Hugging Face Spaces, their login system is used automatically.
    def auth(username, password):
        return True

    # The 'auth' parameter is passed to launch(), not as a separate method call.
    demo.queue().launch(auth=auth, show_error=True)