import gradio as gr
import time
import uuid
from util import (
    create_task_v3, 
    get_task_result,
)


IP_Dict = {}

def generate_trump_voice_with_realtime_updates(text, word_num, request: gr.Request):
    """
    Trump AI voice generation function with real-time status updates
    """
    client_ip = request.client.host
    x_forwarded_for = dict(request.headers).get('x-forwarded-for')
    if x_forwarded_for:
        client_ip = x_forwarded_for   
    if client_ip not in IP_Dict:
        IP_Dict[client_ip] = 0
    IP_Dict[client_ip] += 1
    print(f"client_ip: {client_ip}, count: {IP_Dict[client_ip]}")
    if IP_Dict[client_ip] > 2:
        msg = "You have reached the maximum number of requests, please go to https://donaldtrumpaivoice.com for more"
        yield msg, None, msg
        return msg, None, msg

    if not text or len(text.strip()) < 3:
        return "Text too short, please enter at least 3 characters", None, "No task information"
    
    try:
        task_type = "voice"
        
        # Create task
        task_result = create_task_v3(task_type, text.strip(), word_num, is_rewrite=False)
        if not task_result:
            return "Failed to create task", None, "Task creation failed"
        else:
            yield "Task created successfully", None, "Task creation successful"

        max_polls = 300
        poll_interval = 1
        for i in range(max_polls):
            time.sleep(poll_interval)
            task = get_task_result(task_result['uuid'])
            # print(task, i, "get_task_result")
            if task.get('data', {}):
                status = task.get('data').get('status', '')
                text_final = task.get('data').get('text_final', '')
                if status in ['completed',]:
                    voice_url = task.get('data').get('voice_url', '')
                    task_url = f"https://donaldtrumpaivoice.com/task/{task_result['uuid']}"
                    yield f"✅ success，go to task page {task_url} to check status", voice_url, text_final
                    return "✅ Generation successful!", voice_url, "success"
                elif status in ['failed', 'voice_error', 'no_credits']:
                    yield "❌ Generation failed!", None, None
                    return "❌ Generation failed!", None, None
                else:
                    task_url = f"https://donaldtrumpaivoice.com/task/{task_result['uuid']}"
                    yield f"query {i} times, on processing, go to task page {task_url} to check status", None, text_final
        return "❌ Generation failed!", None, None
    except Exception as e:
        error_msg = f"Generation failed: {str(e)}"
        yield error_msg, None, f"❌ Error message: {error_msg}"
        return error_msg, None, f"❌ Error message: {error_msg}"

# Create Gradio Interface
with gr.Blocks(title="Donald Trump AI Voice", theme=gr.themes.Soft()) as demo:
    gr.Markdown("""
    <div style='text-align: center;'>
        <h1>🎤 Trump AI Voice </h1>
        <p><strong>🔥 Try the most advanced Trump AI Voice and Video generator for FREE at <a href='https://donaldtrumpaivoice.com/' target='_blank'>donaldtrumpaivoice.com</a>! Generate authentic Trump-style audio and video from any text input, instantly and for free!</strong></p>
        <p style='font-size: 14px; color: #666;'>
            You can now try generating Trump's audio for free, come and try it!<br>
            <strong>For the best experience and more features, visit <a href='https://donaldtrumpaivoice.com/' target='_blank'>donaldtrumpaivoice.com</a> now!</strong>
        </p>
    </div>
    """)
    
    with gr.Row():
        with gr.Column(scale=2):
            text_input = gr.Textbox(
                label="📝 Input Text",
                lines=4,
                placeholder="Enter what you want Trump to say...",
                value="Hello everyone, this is a demonstration of the Trump AI Voice system with real-time status monitoring."
            )
        
        with gr.Column(scale=1):
            word_num_slider = gr.Slider(
                10, 50, value=25, step=5,
                label="⏱️ Duration (Word Count)"
            )

            submit_btn = gr.Button(
                "🚀 Generate Trump AI Voice",
                variant="primary",
                size="lg"
            )
    
    with gr.Row():
        status_output = gr.Textbox(
            label="📊 Status",
            interactive=False,
            placeholder="Waiting for generation..."
        )
    
    with gr.Row():
        audio_output = gr.Audio(
            label="🎵 Trump AI Voice",
            interactive=False
        )
    
    with gr.Row():
        task_info = gr.Textbox(
            label="📋 AI Rewritten Text with Latest News",
            interactive=False,
            lines=12,
            placeholder="AI rewritten text with the latest news will be shown here..."
        )

    # Bind event
    submit_btn.click(
        generate_trump_voice_with_realtime_updates,
        inputs=[text_input, word_num_slider],
        outputs=[status_output, audio_output, task_info]
    )

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