import gradio as gr
import boto3
import os
import shutil
import uuid

custom_css = """
footer {display: none !important;}
.gradio-container {min-height: 0px !important;}
#download-header {text-align: center; margin-bottom: 20px;}
"""

# --- CONFIGURATION ---
STORAGE_PATH = "/data"
APP_TOKEN = os.environ.get("APP_TOKEN")

# Setup Amazon SES
ses = boto3.client(
    'ses',
    aws_access_key_id=os.environ.get("AWS_ACCESS_KEY_ID"),
    aws_secret_access_key=os.environ.get("AWS_SECRET_ACCESS_KEY"),
    region_name=os.environ.get("AWS_REGION")
)

def process_submission(name, company, message, files, url_token):
    if url_token != APP_TOKEN:
        return gr.update(visible=True, value="### ❌ Access Denied: Invalid Security Token")

    if not files:
        return gr.update(visible=True, value="### ⚠️ Please select at least one file to upload.")

    job_id = str(uuid.uuid4())[:8]
    job_folder = os.path.join(STORAGE_PATH, job_id)
    os.makedirs(job_folder, exist_ok=True)
    
    space_id = os.environ.get('SPACE_ID', 'vocalmetric/storage')
    direct_url = f"https://{space_id.replace('/', '-')}.hf.space"
    
    file_links = []
    for file in files:
        file_name = os.path.basename(file.name)
        dest_path = os.path.join(job_folder, file_name)
        shutil.copy(file.name, dest_path)
        
        proxy_link = f"{direct_url}?job={job_id}&file={file_name}"
        file_links.append(f"<li><a href='{proxy_link}' style='color: #3b82f6;'>{file_name}</a></li>")

    email_body = f"""
    <div style="font-family: sans-serif; color: #333;">
        <h2>Vocalmetric: New Evidence Submission</h2>
        <p><strong>Name:</strong> {name}</p>
        <p><strong>Company:</strong> {company}</p>
        <p><strong>Instructions:</strong><br>{message}</p>
        <hr>
        <p><strong>Click to Download Files:</strong></p>
        <ul>{"".join(file_links)}</ul>
    </div>
    """

    try:
        ses.send_email(
            Source='noreply@songevolution.com',
            Destination={'ToAddresses': ['ettiennelane@gmail.com', 'riaanrsteyn@gmail.com']},
            Message={
                'Subject': {'Data': f'Vocalmetric: New Submission from {name}'},
                'Body': {'Html': {'Data': email_body}}
            }
        )
    except Exception as e:
        return gr.update(visible=True, value=f"### ❌ Email Error: {str(e)}")

    return gr.update(visible=True, value="### ✅ Success! Files securely uploaded.")

# --- UI HELPER FUNCTIONS ---
def disable_ui():
    # Instantly disable the button and show a loading message
    return gr.update(interactive=False, value="Sending..."), gr.update(visible=True, value="### ⏳ Encrypting and sending. Please wait...")

def enable_ui():
    # Re-enable the button after the process finishes
    return gr.update(interactive=True, value="Send Files")

# --- USER INTERFACE ---
theme = gr.themes.Soft(primary_hue="blue", neutral_hue="slate").set(
    body_background_fill="#0f172a",
    block_background_fill="#1e293b",
    block_label_text_color="#94a3b8",
    button_primary_background_fill="#3b82f6"
)

with gr.Blocks() as demo:
    url_token = gr.State()
    
    # Download Screen (Hidden by default)
    with gr.Column(visible=False) as download_ui:
        gr.Markdown("## 📥 Download Forensic Evidence", elem_id="download-header")
        gr.Markdown("This file was securely retrieved from the Vocalmetric private bucket.")
        file_display = gr.File(label="Ready for Download")
        back_btn = gr.Button("Return to Upload Page", variant="secondary", size="sm")

    # Upload Screen (Shown by default)
    with gr.Column(visible=True) as upload_ui:
        gr.Markdown("# Evidence Upload")
        with gr.Row():
            name = gr.Textbox(label="Your Name")
            company = gr.Textbox(label="Company Name")
        message = gr.TextArea(label="Forensic Analysis Instructions", lines=4)
        file_input = gr.File(label="Select Audio Files", file_count="multiple")
        send_btn = gr.Button("Send Files", variant="primary")
        status = gr.Markdown(visible=False)

    def on_load(request: gr.Request):
        token = request.query_params.get("token", "")
        job = request.query_params.get("job")
        filename = request.query_params.get("file")
        
        if job and filename:
            file_path = os.path.join(STORAGE_PATH, job, filename)
            if os.path.exists(file_path):
                return gr.update(visible=False), gr.update(visible=True), gr.update(value=file_path), token, gr.update(visible=False)
            else:
                return gr.update(visible=True), gr.update(visible=False), gr.update(value=None), token, gr.update(visible=True, value=f"### ❌ Error: File '{filename}' not found. It may have been deleted.")
            
        return gr.update(visible=True), gr.update(visible=False), gr.update(value=None), token, gr.update(visible=False)

    demo.load(on_load, outputs=[upload_ui, download_ui, file_display, url_token, status])
    
    # --- CHAINED CLICK EVENTS ---
    send_btn.click(
        disable_ui, None, [send_btn, status]
    ).then(
        process_submission, [name, company, message, file_input, url_token], [status]
    ).then(
        enable_ui, None, [send_btn]
    )
    
    back_btn.click(lambda: (gr.update(visible=True), gr.update(visible=False), gr.update(visible=False)), outputs=[upload_ui, download_ui, status])

demo.launch(theme=theme, css=custom_css, allowed_paths=[STORAGE_PATH])